[CPO/SimPO] Truncation in CPOTrainer can silently produce empty completions -> NaN loss with loss_type="simpo"
- Dominant language
- Python
- Stars
- 19.3k
- Forks
- 3k
- Avg merge
- 1d 20h
- Merged PRs (30d)
- 194
Description
### Reproduction
## Description
In `trl/experimental/cpo/cpo_trainer.py`, the response truncation logic in
`tokenize_row` truncates **both** responses using the length of the **longer**
response of the pair:
```python
longer_response_length = max(len(chosen_tokens["input_ids"]), len(rejected_tokens["input_ids"]))
# if combined sequence is too long, truncate the response
for answer_tokens in [chosen_tokens, rejected_tokens]:
if len(answer_tokens["prompt_input_ids"]) + longer_response_length > self.max_length:
for k in ["input_ids", "attention_mask"]:
answer_tokens[k] = answer_tokens[k][: self.max_length - longer_response_length]
```
(https://github.com/huggingface/trl/blob/4708879046dde897bdb84667a97e781bfc1965d2/trl/experimental/cpo/cpo_trainer.py#L554-L560)
When `longer_response_length > self.max_length`, the slice bound
`self.max_length - longer_response_length` is **negative**, so `lst[:negative]`
cuts from the end. If the *shorter* response of the pair has fewer tokens than
that negative offset, it is truncated to an **empty sequence** (0 completion
tokens, all labels are -100).
Consequences:
- With `loss_type="simpo"`, log-probs are length-normalized, so an empty
completion produces a `0/0` division -> `NaN` in `logps`, which propagates to
the loss and destroys all model weights after the first optimizer step.
Training metrics show `loss = 0`, `grad_norm = nan`, all rewards/logps `nan`.
- With other loss types the empty completion contributes a silent `logp = 0`,
which is incorrect but does not crash — arguably worse, since it corrupts
training without any visible symptom.
This is easy to hit in practice with self-generated preference data
(e.g., math reasoning), where response lengths are heavy-tailed: a single pair
with one very long response (> max_length) and one normal-length response
triggers it.
## Reproduction
Pure-Python demonstration of the slicing logic (3 lines):
```python
max_length = 2048
chosen = list(range(5000)) # long response, > max_length
rejected = list(range(600)) # normal short response
longer = max(len(chosen), len(rejected)) # 5000
print(len(chosen[: max_length - longer])) # 2048 (truncated, survives)
print(len(rejected[: max_length - longer])) # 0 (silently EMPTY)
```
End-to-end reproduction through the trainer (any small public model):
```python
import torch
from datasets import Dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl.experimental.cpo import CPOTrainer, CPOConfig
model_id = "Qwen/Qwen2.5-0.5B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32)
pair = {
"prompt": "Question: what is 2+2?\nAnswer:",
"chosen": " step" * 5000, # ~5000 tokens, exceeds max_length
"rejected": " The answer is 4.", # ~7 tokens, healthy short response
}
def count_completion_tokens(max_length):
args = CPOConfig(
output_dir="/tmp/cpo_repro",
loss_type="simpo",
cpo_alpha=0.0,
beta=0.1,
max_length=max_length,
per_device_train_batch_size=1,
report_to="none",
)
trainer = CPOTrainer(
model=model,
args=args,
train_dataset=Dataset.from_list([pair]),
processing_class=tokenizer,
)
row = trainer.train_dataset[0]
n_chosen = sum(t != -100 for t in row["chosen_labels"])
n_rejected = sum(t != -100 for t in row["rejected_labels"])
return n_chosen, n_rejected
print("max_length=2048:", count_completion_tokens(2048))
# -> (2048, 0) rejected is EMPTY -> NaN loss with loss_type="simpo"
print("max_length=8192:", count_completion_tokens(8192))
# -> (5001, 7) control: everything fits, both responses intact
```
Same data, only `max_length` differs: the short response is destroyed exactly
when the *other* response exceeds `max_length`.
## Expected behavior
Each response should be truncated based on **its own** length, e.g.:
```python
for answer_tokens in [chosen_tokens, rejected_tokens]:
if len(answer_tokens["prompt_input_ids"]) + len(answer_tokens["input_ids"]) > self.max_length:
for k in ["input_ids", "attention_mask"]:
answer_tokens[k] = answer_tokens[k][
: max(0, self.max_length - len(answer_tokens["prompt_input_ids"]))
]
```
so that a short response is never emptied because its pair partner is long.
Additionally, a guard (warning or error) when a completion ends up with zero
tokens after truncation would make failures like the SimPO NaN immediately
diagnosable instead of silent.
### System Info
- Platform: Linux-6.8.0-124-generic-x86_64-with-glibc2.39
- Python version: 3.12.3
- TRL version: 1.6.0
- PyTorch version: 2.11.0+cu130
- accelerator(s): NVIDIA L40S, NVIDIA L40S
- Transformers version: 5.12.1
- Accelerate version: 1.14.0
- Accelerate config:
- compute_environment: LOCAL_MACHINE
- distributed_type: FSDP
- mixed_precision: bf16
- use_cpu: False
- debug: False
- num_processes: 2
- machine_rank: 0
- num_machines: 1
- rdzv_backend: static
- same_network: True
- main_training_function: main
- enable_cpu_affinity: False
- fsdp_config: {'fsdp_activation_checkpointing': False, 'fsdp_auto_wrap_policy': 'TRANSFORMER_BASED_WRAP', 'fsdp_backward_prefetch': 'BACKWARD_PRE', 'fsdp_cpu_ram_efficient_loading': True, 'fsdp_forward_prefetch': False, 'fsdp_offload_params': False, 'fsdp_sharding_strategy': 'FULL_SHARD', 'fsdp_state_dict_type': 'SHARDED_STATE_DICT', 'fsdp_sync_module_states': True, 'fsdp_use_orig_params': False, 'fsdp_version': 1}
- downcast_bf16: yes
- tpu_use_cluster: False
- tpu_use_sudo: False
- tpu_env: []
- Datasets version: 5.0.0
- HF Hub version: 1.23.0
- bitsandbytes version: 0.49.2
- DeepSpeed version: not installed
- Liger-Kernel version: not installed
- PEFT version: 0.19.1
- vLLM version: 0.23.0
### Checklist
- [x] I have checked that my issue isn't already filed (see [open issues](https://github.com/huggingface/trl/issues?q=is%3Aissue))
- [x] I have included my system information
- [x] Any code provided is minimal, complete, and reproducible ([more on MREs](https://docs.github.com/en/get-started/writing-on-github/working-with-advanced-formatting/creating-and-highlighting-code-blocks))
- [x] Any code provided is properly formatted in code blocks, (no screenshot, [more on code blocks](https://docs.github.com/en/get-started/writing-on-github/working-with-advanced-formatting/creating-and-highlighting-code-blocks))
- [x] Any traceback provided is complete
Contributor guide
Research direction
Read trl/experimental/cpo/cpo_trainer.py, especially tokenize_row and the response-truncation logic. Run the pure-Python slicing reproduction or the supplied CPOTrainer example, then verify that truncation uses each response's own length and that the short completion is not empty or associated with a NaN SimPO loss.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 78/100