🕊️ DPO padding free

[qgallouedec]

What does this PR do?

demo; further experiments in next comment

from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import DPOConfig, DPOTrainer
import torch

model_id = "Qwen/Qwen2-0.5B-Instruct"
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2")
tokenizer = AutoTokenizer.from_pretrained(model_id)
dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train[:10%]")
training_args = DPOConfig(output_dir="Gemma2-2B-DPO-pf", max_prompt_length=128, max_completion_length=128, logging_steps=10, padding_free=True)
trainer = DPOTrainer(model=model, args=training_args, train_dataset=dataset, processing_class=tokenizer)
trainer.train()

With and without padding-free (not sure why they don't match exactly, the logits do precisely match though)

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]

Regression test:

from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import DPOConfig, DPOTrainer
import torch

model_id = "Qwen/Qwen2-0.5B-Instruct"
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2")
tokenizer = AutoTokenizer.from_pretrained(model_id)
dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train[:10%]")
# dataset = load_dataset("trl-internal-testing/zen", "standard_preference", split="train")
training_args = DPOConfig(output_dir="Qwen2-0.5B-DPO-no_pf", max_prompt_length=128, max_completion_length=128, logging_steps=10, padding_free=False)
trainer = DPOTrainer(model=model, args=training_args, train_dataset=dataset, processing_class=tokenizer)
trainer.train()

Is the new padding_free=False (no_pf in screenshot) equivalent to DPO on current main branch (main in screenshot)? -> yes

Does padding_free=True (pf in screenshot) results match padding_free=False (no_pf in screenshot) results? -> Yes

(note: on screenshots its written "Gemma" but it's actually a Qwen model trained)

[qgallouedec]

rolling works in both cases (flattened tensors and batched)

# Padding case
# input_ids = [[1, 2, 3, 4],
#              [5, 6, 7, 8]]
labels = input_ids[:, 1:].clone()
# labels =    [[2, 3, 4],
#              [6, 7, 8]]

# But

# Padding-free case
# input_ids = [[1, 2, 3, 4, 5, 6, 7, 8]]
labels = input_ids[:, 1:].clone()
# labels =    [[2, 3, 4, 5, 6, 7, 8]]

The first token of the first sequence (1) is removed but not the first token of the second sequence (5). To align the labels while keeping a consistent behaviour across sequences in the batch, we use roll instead. The only difference is that the first token, instead of being discarded, is appended to the end.

# Padding case
# input_ids = [[1, 2, 3, 4],
#              [5, 6, 7, 8]]
labels = torch.roll(input_ids, shifts=-1, dims=1)
# labels =    [[2, 3, 4, 1],
#              [6, 7, 8, 5]]

# And

# Padding-free case
# input_ids = [[1, 2, 3, 4, 5, 6, 7, 8]]
labels = torch.roll(input_ids, shifts=-1, dims=1)
# labels =    [[2, 3, 4, 5, 6, 7, 8, 1]]

[qgallouedec]

Read https://github.com/huggingface/trl/pull/2520/files#r1905992872 before

since we have an additional token in the end, we need to keep one additional token.
This update make things at this point nicer imo.

[August-murr]

LGTM!