DPO Trainer

[kashif]

Initial DPOTrainer class for #405 by copying the PPOTrainer RewardTrainer and started to implement changes in it

Fixes #405

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[Forbu]

Thank you @kashif.
I think one interesting stuff to do would be to make the forward pass of the model_ref in no_grad mode no ?
You don't need to compute the backward pass for the ref model because you don't want to modify the ref model weights.

[kashif]

agree @Forbu and I think we will refactor the data-collator so that we only have a mask on the positive and negative part of the sequence...

[gaetanlop]

Hello @kashif, thanks for the DPO integration. I did a refactoring of the data collator to compute the mean logprobs only on the positive and negative part of the sequence. Can you share with me the branch please?

[kashif]

@gaetanlop i have added you to my fork

[gaetanlop]

Thanks @kashif, just pushed the required changes to make the trainer not compute the mean logprobs on masked input_ids. I used a similar approach as the HF DataCollatorForTokenClassification (https://github.com/huggingface/transformers/blob/main/src/transformers/data/data_collator.py)

[TevenLeScao]

Flagging that at some point we'll want the ScriptArguments to be consistent with the TrainingArguments (log_with -> report_to, batch_size -> per_device_train_batch_size, model_name -> model_name_or_path) the second and third ones are especially important since the semantic of the argument is actually different.

[TevenLeScao]

@kashif @gaetanlop a fix to enable distributed training required a change to transformers.Trainer (huggingface/transformers#24326), you'll need to pull transformers from master if you pull it!

[eric-mitchell]

@kashif just to follow up again, the re-run of the DPO stage of Pythia on HH just completed, and training is basically unchanged there as well after the bug fix on our end. So comparing with our original runs should be fine.

You can see the SFT/DPO runs before and after the bug fix here.

[eric-mitchell]

@kashif just wanted to check in on this- curious if you've had the chance to re-run the replication experiment :)

[kashif]

@eric-mitchell i did but not with a SFT'ed pythia... also with Peft worked... We had validation loss around 0.6 or so as per your wandb so that was nice... if you can share your SFT'ed pythia model with me I can also run it now with that?

[eric-mitchell]

Here are the weights to our pre-trained Pythia. You can load with model.load_state_dict(torch.load(PATH)['state']).

It's nice that peft worked! What type of peft did you use?

[kashif]

@eric-mitchell thanks! yes we tried QLora and Lora as well.. let me confirm. Thanks for the weights!

[lvwerra]

Two small nits, then we can merge :)