PEFT support for Online DPO

[qgallouedec]

What does this PR do?

Fixes # (issue)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a GitHub issue? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[qgallouedec]

accelerate launch examples/scripts/dpo_online.py
    --model_name_or_path trl-lib/pythia-1b-deduped-tldr-sft \
    --reward_model_path trl-lib/pythia-1b-deduped-tldr-rm \
    --dataset_name trl-lib/tldr \
    --output_dir pythia-1b-tldr-online-dpo \
    --learning_rate 5.0e-7 \
    --logging_first_step \
    --logging_steps 10 \
    --per_device_train_batch_size 8 \
    --gradient_accumulation_steps 8 \
    --warmup_ratio 0.1 \
    --missing_eos_penalty 1.0 \
    --push_to_hub \
    --dataset_num_proc 32 \
    --use_peft

PEFT stabilises training a lot!
Note that orange training suffer from overoptimization here

[kashif]

can you kindly add the above peft/lora usage command in the script?

[qgallouedec]

Particular attention for reviewing this one. I didn't use the code from DPO mostly because I don't understand it. So I'm afraid to have missed something.

1. Why do you need to merge and unload?
   

   trl/trl/trainer/dpo_trainer.py

   Lines 533 to 535 in e51a5ac

   	# if model is a peft model and we have a peft_config, we merge and unload it first
   	if isinstance(model, PeftModel):
   	model = model.merge_and_unload()

2. I don't support k-bit training yet
   

   trl/trl/trainer/dpo_trainer.py

   Lines 544 to 556 in e51a5ac

   	if getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False):
   	_support_gc_kwargs = hasattr(
   	args, "gradient_checkpointing_kwargs"
   	) and "gradient_checkpointing_kwargs" in list(
   	inspect.signature(prepare_model_for_kbit_training).parameters
   	)
   	prepare_model_kwargs = {"use_gradient_checkpointing": args.gradient_checkpointing}
   	if _support_gc_kwargs:
   	prepare_model_kwargs["gradient_checkpointing_kwargs"] = args.gradient_checkpointing_kwargs
   	model = prepare_model_for_kbit_training(model, **prepare_model_kwargs)

3. Can we drop it for new trainer?
   

   trl/trl/trainer/dpo_trainer.py

   Lines 556 to 565 in e51a5ac

   	model = prepare_model_for_kbit_training(model, **prepare_model_kwargs)
   	elif getattr(args, "gradient_checkpointing", False):
   	# For backward compatibility with older versions of transformers
   	if hasattr(model, "enable_input_require_grads"):
   	model.enable_input_require_grads()
   	else:
   	def make_inputs_require_grad(module, input, output):
   	output.requires_grad_(True)

4. Why do we need to cast? And why only when loaded in 4-bits. Btw why always using the getattr? We know that this arg exists in TrainingArguments
   

   trl/trl/trainer/dpo_trainer.py

   Lines 570 to 574 in e51a5ac

   	if args.bf16 and getattr(model, "is_loaded_in_4bit", False):
   	peft_module_casting_to_bf16(model)
   	# If args.bf16 we need to explicitly call `generate` with torch amp autocast context manager
   	self._peft_has_been_casted_to_bf16 = True

Overall I've choose to go for a code that I understand, even if it implies introducing bug that have been fixed in the past for other trainers.

[lewtun]

Thanks for adding PEFT support! Overall LGTM - have you run an experiment on e.g. TLDR to see if it looks OK?

Edit: sorry, just saw your comment!

[lewtun]

1. Why do you need to merge and unload?

This is needed if we have e.g. an SFT LoRA that needs to be merged into the base model, before finally inserting an adapter for DPO. See here for some considerations.

2. I don't support k-bit training yet

Sure, let's leave this as follow-up PR

3. Can we drop it for new trainer?

Fine with me. We can add it later if people open an issue.

4. Why do we need to cast? And why only when loaded in 4-bits. Btw why always using the getattr? We know that this arg exists in TrainingArguments

Good question, this was added by Younes but maybe @BenjaminBossan can help clarify here why PEFT models need upcasting in 4-bit?

[BenjaminBossan]

7. Why do we need to cast? And why only when loaded in 4-bits. Btw why always using the getattr? We know that this arg exists in TrainingArguments

Good question, this was added by Younes but maybe @BenjaminBossan can help clarify here why PEFT models need upcasting in 4-bit?

The original addition of this function was from #1110 (which in turn references this repo). I don't know the exact context, but I think the reason is to avoid indiscriminately casting all layers to bf16 -- specifically, layer norm stays in float32. This is probably based on some empirical finding that this is better for training, but after skimming the QLoRA paper, I could not find any mention of that, so I'm unsure.

[qgallouedec]

[qgallouedec]

Thanks a lot @BenjaminBossan and @lewtun. I'll further investigate in a dedicated branch.