DataCollatorForCompletionOnlyLM for ORPO Trainer

[abpani]

Is DataCollatorForCompletionOnlyLM supported for ORPOTrainer.

[Advait251206]

Hi @abpani,

DataCollatorForCompletionOnlyLM is primarily designed for supervised fine-tuning (SFT), where you want the model to calculate loss only on the completion portion of a prompt-response sequence.

For ORPO, the situation is slightly different because it is a preference optimization method that uses both chosen and rejected responses.

1. Is DataCollatorForCompletionOnlyLM supported by ORPOTrainer?

ORPOTrainer has its own data preprocessing and collation requirements. It expects preference data containing at least:

• prompt
• chosen
• rejected

The trainer tokenizes these fields and constructs the appropriate inputs and labels for its preference optimization objective.

In particular, ORPO computes:

• A supervised negative log-likelihood loss on the chosen response.
• An odds-ratio preference loss comparing chosen and rejected responses.

The trainer already handles masking the prompt tokens from the chosen and rejected completion labels during its preprocessing.

Therefore, you generally do not need DataCollatorForCompletionOnlyLM when using ORPOTrainer.

2. What if I want to train only on the completion tokens?

That behavior is already incorporated into ORPO's training logic.

For example, your dataset can look like this:

{
    "prompt": "Explain what a neural network is.",
    "chosen": "A neural network is a machine learning model...",
    "rejected": "A neural network is simply a database..."
}

ORPO processes the prompt and both responses, then applies the appropriate masking and preference loss internally.

You should avoid applying an SFT-specific collator that changes the expected batch structure unless you have verified that it is compatible with the trainer's implementation.

3. What if I want to customize the data collator?

ORPOTrainer does expose a data_collator argument, so custom collators are possible.

However, the custom collator must preserve the structure and tensors expected by ORPO's training and loss computation.

For most use cases, I would recommend using the trainer's default collator and allowing ORPOTrainer to handle tokenization, padding, and label preparation.

Conclusion

For standard ORPO training, use the preference dataset format (prompt, chosen, rejected) and let ORPOTrainer handle completion masking internally.

DataCollatorForCompletionOnlyLM is not generally required, and directly substituting it for ORPO's default collator may cause compatibility issues.

Hope this helps!

---

If this answer helped or pointed you in the right direction, I'd appreciate it if you could mark it as the accepted answer so it's easier for others with the same issue to find.

Also, if you found my contribution useful, I'd appreciate it if you could check out my GitHub profile, follow me, and star any repositories you find interesting.

GitHub: https://github.com/Advait251206