TPOTrainer.evaluate() returns NaN eval_loss while training loss is finite
🐛 bug
### Reproduction
# TPOTrainer.evaluate() returns NaN eval_loss while training loss is finite
## Description
`TPOTrainer.evaluate()` returns `eval_loss: nan` even though training runs normally and `train_loss` is finite.
I tested the experimental `TPOTrainer` from `trl.experimental.tpo` using the public dataset `tpo-alignment/triple-preference-ultrafeedback-40K` and the model `Qwen/Qwen3-0.6B`.
The dataset has the expected TPO columns:
```python
['prompt', 'reference', 'chosen', 'rejected']
```
I converted the dataset to conversational/chat format before passing it to `TPOTrainer`:
```python
{
"prompt": [
{"role": "user", "content": prompt}
],
"chosen": [
{"role": "assistant", "content": chosen}
],
"rejected": [
{"role": "assistant", "content": rejected}
],
"reference": [
{"role": "assistant", "content": reference}
],
}
```
Training works and produces finite losses, but `trainer.evaluate()` returns `nan` both before and after training.
## Reproduction
```python
import os
os.environ["TRL_EXPERIMENTAL_SILENCE"] = "1"
os.environ["TOKENIZERS_PARALLELISM"] = "false"
import torch
from datasets import load_dataset
from peft import LoraConfig
from transformers import AutoTokenizer
from trl.experimental.tpo import TPOConfig, TPOTrainer
MODEL_NAME = "Qwen/Qwen3-0.6B"
DATASET_NAME = "tpo-alignment/triple-preference-ultrafeedback-40K"
TRAIN_SAMPLES = 1000
EVAL_SAMPLES = 100
MAX_LENGTH = 384
MAX_STEPS = 100
raw_train = load_dataset(DATASET_NAME, split="train")
raw_eval = load_dataset(DATASET_NAME, split="test")
raw_train = raw_train.select(range(TRAIN_SAMPLES))
raw_eval = raw_eval.select(range(EVAL_SAMPLES))
print("Raw train columns:", raw_train.column_names)
print("Raw eval columns:", raw_eval.column_names)
print("Raw train size:", len(raw_train))
print("Raw eval size:", len(raw_eval))
def clean_str(x):
if x is None:
return ""
return str(x).strip()
def to_chat_format(example):
return {
"prompt": [
{"role": "user", "content": clean_str(example["prompt"])}
],
"chosen": [
{"role": "assistant", "content": clean_str(example["chosen"])}
],
"rejected": [
{"role": "assistant", "content": clean_str(example["rejected"])}
],
"reference": [
{"role": "assistant", "content": clean_str(example["reference"])}
],
}
train_ds = raw_train.map(
to_chat_format,
remove_columns=raw_train.column_names,
)
eval_ds = raw_eval.map(
to_chat_format,
remove_columns=raw_eval.column_names,
)
print("Processed columns:", train_ds.column_names)
print("Example prompt:", train_ds[0]["prompt"])
tokenizer = AutoTokenizer.from_pretrained(
MODEL_NAME,
trust_remote_code=True,
use_fast=True,
)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"
print("pad_token:", tokenizer.pad_token)
print("eos_token:", tokenizer.eos_token)
peft_config = LoraConfig(
r=8,
lora_alpha=16,
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
target_modules="all-linear",
)
args = TPOConfig(
output_dir="Qwen3-0.6B-TPO-LoRA-fast-test",
per_device_train_batch_size=1,
per_device_eval_batch_size=1,
gradient_accumulation_steps=4,
gradient_checkpointing=True,
max_steps=MAX_STEPS,
learning_rate=1e-5,
warmup_steps=10,
lr_scheduler_type="cosine",
max_length=MAX_LENGTH,
loss_type="sigmoid",
beta=0.01,
tpo_alpha=1.0,
fp16=True,
bf16=False,
logging_steps=5,
eval_strategy="steps",
eval_steps=25,
save_strategy="no",
remove_unused_columns=False,
report_to="none",
dataloader_num_workers=2,
)
trainer = TPOTrainer(
model=MODEL_NAME,
args=args,
train_dataset=train_ds,
eval_dataset=eval_ds,
processing_class=tokenizer,
peft_config=peft_config,
)
print("Eval before training:")
print(trainer.evaluate())
train_result = trainer.train()
print(train_result.metrics)
print("Eval after training:")
print(trainer.evaluate())
```
outputs:
```text
================================================================================
Torch: 2.10.0+cu128
CUDA: True
GPU: Tesla T4
================================================================================
Raw train columns: ['prompt', 'reference', 'chosen', 'rejected']
Raw eval columns: ['prompt', 'reference', 'chosen', 'rejected']
Raw train size: 40000
Raw eval size: 1910
Fast train size: 1000
Fast eval size: 100
Converting train dataset: 100% 1000/1000
Converting eval dataset: 100% 100/100
Processed columns: ['prompt', 'reference', 'chosen', 'rejected']
Example prompt: [{'role': 'user', 'content': 'Write about the importance of self-reflection and self-improvement during workplace conflicts.'}]
pad_token: <|endoftext|>
eos_token: <|im_end|>
Loading weights: 100% 311/311
The tied weights mapping and config for this model specifies to tie model.embed_tokens.weight to lm_head.weight, but both are present in the checkpoints, so we will NOT tie them. You should update the config with `tie_word_embeddings=False` to silence this warning
Tokenizing train dataset: 100% 1000/1000
Tokenizing eval dataset: 100% 100/100
================================================================================
EVAL BEFORE TRAINING
================================================================================
The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'bos_token_id': None, 'pad_token_id': 151643}.
{'eval_loss': nan, 'eval_model_preparation_time': 0.0189, 'eval_runtime': 25.1482, 'eval_samples_per_second': 3.976, 'eval_steps_per_second': 3.976}
================================================================================
START FAST TRAINING
================================================================================
Step Training Loss Validation Loss Model Preparation Time
25 4.514127 nan 0.018900
50 3.317068 nan 0.018900
75 2.523718 nan 0.018900
100 3.386643 nan 0.018900
================================================================================
TRAIN RESULT
================================================================================
{'train_runtime': 412.4126, 'train_samples_per_second': 0.97, 'train_steps_per_second': 0.242, 'total_flos': 1452732883402752.0, 'train_loss': 3.5000829219818117}
================================================================================
EVAL AFTER TRAINING
================================================================================
{'eval_loss': nan, 'eval_model_preparation_time': 0.0189, 'eval_runtime': 24.9824, 'eval_samples_per_second': 4.003, 'eval_steps_per_second': 4.003}
```
## Actual behavior
`trainer.evaluate()` returns:
```python
{'eval_loss': nan, ...}
```
This happens both before and after training.
During training, validation loss is also reported as `nan`:
```text
Step 25 Training Loss 4.514127 Validation Loss nan
Step 50 Training Loss 3.317068 Validation Loss nan
Step 75 Training Loss 2.523718 Validation Loss nan
Step 100 Training Loss 3.386643 Validation Loss nan
```
However, training itself appears to work because the training losses are finite and the final training metrics are finite:
```python
{
'train_runtime': 412.4126,
'train_samples_per_second': 0.97,
'train_steps_per_second': 0.242,
'total_flos': 1452732883402752.0,
'train_loss': 3.5000829219818117,
}
```
## Expected behavior
`trainer.evaluate()` should return a finite `eval_loss`, or at least expose a clear reason why eval loss cannot be computed for the current batch/data configuration.
Since training loss is finite and the same dataset format is used for training and evaluation, I expected evaluation loss to also be finite.
## Additional observations
The issue does not seem to be caused by missing columns. The dataset has the expected TPO fields:
```python
['prompt', 'reference', 'chosen', 'rejected']
```
I initially saw tokenizer prefix mismatch warnings when using plain string format. Those warnings disappeared after converting the examples to chat format. After the conversion, tokenization completes successfully and training starts normally.
The following warning appears, but training continues:
```text
The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config.
The model config and generation config were aligned accordingly.
Updated tokens: {'bos_token_id': None, 'pad_token_id': 151643}.
```
I also manually checked that training itself is not failing: the model trains and produces coherent generations after 100 steps. The NaN seems specific to `TPOTrainer.evaluate()` / validation loss reporting rather than a full training failure.
No Python traceback is produced because the script does not crash. The issue is that `TPOTrainer.evaluate()` returns `{'eval_loss': nan, ...}` while training continues with finite losses. The full relevant output is included in the `outputs:` block above.
## Environment
```text
TRL: 1.3.0
Torch: 2.10.0+cu128
CUDA available: True
GPU: Tesla T4
Model: Qwen/Qwen3-0.6B
Dataset: tpo-alignment/triple-preference-ultrafeedback-40K
Precision: fp16=True, bf16=False
PEFT: LoRA, target_modules="all-linear"
```
## Question
Is this a bug in `TPOTrainer.evaluate()` / evaluation loss aggregation, or is there an additional configuration required for evaluation with `TPOTrainer`?
If this is expected behavior for the current experimental implementation, it would be helpful to document the limitation or provide recommended evaluation metrics for TPO, such as chosen-vs-rejected preference accuracy or reward margin.
### System Info
Platform: Linux-6.6.113+-x86_64-with-glibc2.35
Python: 3.12.13 (main, Mar 4 2026, 09:23:07) [GCC 11.4.0]
PyTorch: 2.10.0+cu128
CUDA available: True
GPU: Tesla T4
CUDA version: 12.8
Transformers: 5.0.0
TRL: 1.3.0
Accelerate: 1.13.0
Datasets: 4.8.4
PEFT: 0.18.1
### 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
2 条评论