基于DPO优化LLM时评估出现NaN值的问题排查

问题背景 #

用DPO（Direct Preference Optimization）优化Pythia-14m大模型，训练完成后评估时所有指标都出现NaN值，以下是训练代码及评估结果，需排查问题原因和代码错误。

训练代码 #

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


model_name = "EleutherAI/pythia-14m"
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

def preprocess_data(item):
    return {
        'prompt': 'Instruct: ' + item['prompt'] + '\n',
        'chosen': 'Output: ' + item['chosen'],
        'rejected': 'Output: ' + item['rejected']
    }

dataset = load_dataset('jondurbin/truthy-dpo-v0.1', split="train")
dataset = dataset.map(preprocess_data)
split_dataset = dataset.train_test_split(test_size=0.1)
train_dataset = split_dataset['train']
val_dataset = split_dataset['test']

print(f"Length of train data: {len(train_dataset)}")
print(f"Length of validation data: {len(val_dataset)}")


# Tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.unk_token

# Model to fine-tune
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    low_cpu_mem_usage=True,
    torch_dtype=torch.float16
).to(device)

model_ref = AutoModelForCausalLM.from_pretrained(
    model_name,
    low_cpu_mem_usage=True,
    torch_dtype=torch.float16
).to(device)

# Config
training_args = DPOConfig(
    output_dir="./output",
    beta=0.1,
    max_length=512,
    max_prompt_length=128,
    remove_unused_columns=False,
)

# Load trainer
dpo_trainer = DPOTrainer(
    model,
    model_ref,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset,  
    tokenizer=tokenizer,
)

# Train
dpo_trainer.train()

# Evaluate
evaluation_results = dpo_trainer.evaluate()
print("Evaluation Results:", evaluation_results)

评估结果 #

Evaluation Results: {'eval_loss': nan, 'eval_runtime': 0.5616, 'eval_samples_per_second': 181.61, 'eval_steps_per_second': 12.463, 'eval_rewards/chosen': nan, 'eval_rewards/rejected': nan, 'eval_rewards/accuracies': 0.0, 'eval_rewards/margins': nan, 'eval_logps/rejected': nan, 'eval_logps/chosen': nan, 'eval_logits/rejected': nan, 'eval_logits/chosen': nan, 'epoch': 3.0}

问题原因分析 #

1. Pad Token设置错误：代码中将tokenizer.pad_token设为tokenizer.unk_token，UNK token是模型用于表示未知词汇的token，在计算损失和log概率时不会被忽略，会干扰DPO的损失计算逻辑，进而产生NaN。
2. 数值精度问题：Pythia-14m是极小模型，使用torch.float16训练时容易出现数值不稳定（如梯度溢出、log计算下溢），导致中间计算结果变为NaN。
3. 训练配置缺失关键参数：DPOConfig未设置batch size、学习率等核心参数，默认值可能导致训练过程不稳定，引发NaN。
4. 未过滤异常样本：数据集可能存在过长的prompt或回复，超过max_length和max_prompt_length限制，截断后可能导致计算异常。

修复方案 #

1. 正确设置Pad Token #

将pad token改为eos token（序列结束token），模型在计算损失时会自动忽略pad部分：

tokenizer.pad_token = tokenizer.eos_token

2. 调整模型数据类型 #

改用torch.float32或torch.bfloat16（GPU支持的话）避免数值不稳定：

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    low_cpu_mem_usage=True,
    torch_dtype=torch.float32  # 或torch.bfloat16
).to(device)

model_ref = AutoModelForCausalLM.from_pretrained(
    model_name,
    low_cpu_mem_usage=True,
    torch_dtype=torch.float32
).to(device)

3. 完善训练配置 #

添加必要的训练参数，稳定训练过程：

training_args = DPOConfig(
    output_dir="./output",
    beta=0.1,
    max_length=512,
    max_prompt_length=128,
    remove_unused_columns=False,
    per_device_train_batch_size=4,  # 根据GPU内存调整
    per_device_eval_batch_size=4,
    learning_rate=5e-5,
    num_train_epochs=3,
    gradient_accumulation_steps=2,
    logging_steps=10,
    evaluation_strategy="epoch",  # 每轮epoch评估，便于监控
)

4. 过滤异常样本 #

预处理后过滤掉长度超标的样本，避免截断导致的计算异常：

def filter_long_samples(item):
    prompt_len = len(tokenizer(item['prompt'])['input_ids'])
    chosen_len = len(tokenizer(item['chosen'])['input_ids'])
    rejected_len = len(tokenizer(item['rejected'])['input_ids'])
    # 确保prompt不超过max_prompt_length，回复不超过max_length - max_prompt_length
    return prompt_len <= 128 and chosen_len <= 384 and rejected_len <= 384

dataset = dataset.filter(filter_long_samples)

总结 #

优先修复Pad Token设置和数据类型问题，这两个是导致NaN最常见的原因。调整后重新训练评估，应该能解决指标为NaN的问题。

内容的提问来源于stack exchange，提问作者Refinath