Skip to content

Adds the CPO Alignment Loss Function#382

Merged
ByronHsu merged 3 commits into
mainfrom
pramodith/chunked_cpo_loss
Nov 19, 2024
Merged

Adds the CPO Alignment Loss Function#382
ByronHsu merged 3 commits into
mainfrom
pramodith/chunked_cpo_loss

Conversation

@pramodith

@pramodith pramodith commented Nov 14, 2024

Copy link
Copy Markdown
Collaborator

Summary

CPO is almost the same as DPO with the major difference being that the Reference Model in CPO is assumed to be a Uniform distribution. This assumption leads to the cancellation of all terms related to the reference model.

$$CPOLoss = -\log(\sigma(\beta\log(\pi_\theta(y_c|x)) - \beta\log(\pi_\theta(y_r|x))))$$

This corresponds to equation 3 in the paper. Additionally CPO also assumes a scaling factor alpha for the NLL loss on the preferred response. In TRL this corresponds to the CPOTrainer using a loss_type="sigmoid"

We also refactor the test cases for chunked loss functions to include a generic HFAlignmentLoss base class that takes care some of the plumbing work to correctly generate batches of input, calculate the NLLoss etc. All future test cases can inherit from this class and just implement the alignment_loss function to compare implementation in the TRL lib versus the custom impl.

Testing Done

A100-80G-SXM

Benchmark Results:

Screenshot 2024-11-14 at 5 17 42 PM
Fused Linear CPO Loss Speed

  • Hardware Type:
  • run make test to ensure correctness
  • run make checkstyle to ensure code style
  • run make test-convergence to ensure convergence

@shivam15s

This comment was marked as outdated.

@shivam15s
shivam15s force-pushed the pramodith/chunked_cpo_loss branch from 64a742c to 2d4f911 Compare November 19, 2024 02:40

@shivam15s shivam15s left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lgtm

beta=0.1,
compute_nll_loss=True,
compiled=True,
compiled=False,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd rather have compiled True by default

@ByronHsu
ByronHsu merged commit 16d06ed into main Nov 19, 2024
@ByronHsu
ByronHsu deleted the pramodith/chunked_cpo_loss branch November 19, 2024 05:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants