Improper buggy SFT training of Llama-3.2-3B-Instruct.

### Reproduction

@qgallouedec @kashif When I perform SFT on Llama-3.2-3B-Instruct using the recent trl library (By recent, I mean the ones which don't use DataCollatorForCompletionOnly for calculating completion only loss anymore), I see that the model is not training properly. Whereas training is good when I use trl==0.13.0 with the said collator for calculating completion-only loss.

I use the following code to run the training.

training data: ```rrvaswin/llama_star```

accelerate_config.yaml

```bash
compute_environment: LOCAL_MACHINE
debug: false
deepspeed_config:
deepspeed_multinode_launcher: standard
offload_optimizer_device: none
offload_param_device: none
zero3_init_flag: true
zero3_save_16bit_model: true
zero_stage: 3
distributed_type: DEEPSPEED
downcast_bf16: 'no'
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 4
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
```

run.sh
```bash

source .venv/bin/activate

export CUDA_VISIBLE_DEVICES=6,7
NUM_GPUS=$(echo $CUDA_VISIBLE_DEVICES | awk -F',' '{print NF}')

echo "Running with $NUM_GPUS GPUs"

MODEL_NAME="meta-llama/Llama-3.2-3B-Instruct"
DATASET_NAME="rrvaswin/llama_star"
SPLIT="train"
PROMPT_FIELD="question"
RESPONSE_FIELD="response"
OUTPUT_DIR="models/llama_star_sft"
MAX_LENGTH=5192
LEARNING_RATE=5e-6
NUM_EPOCHS=1
SAVE_STRATEGY="epoch"
EFF_BATCH_SIZE=16
PER_DEVICE_BATCH_SIZE=4
OPTIM="adamw_torch"
LR_SCHEDULER="cosine"
SAVE_TOTAL_LIMIT=3
REPORT_TO="wandb"
LOGGING_STEPS=10
BF16=True
GRADIENT_CHECKPOINTING=True
WARMUP_RATIO=0.1
PACKING=False
SEED=42

RUN_NAME="llama_star_sft_pc_mode"
DATASET_MODE="pc_mode"

ACCELERATE_LOG_LEVEL=info accelerate launch \
--main_process_port 0 \
--config_file config/accelerate_config.yaml \
--num_processes=$NUM_GPUS \
sft.py \
--model_name $MODEL_NAME \
--dataset_name $DATASET_NAME \
--split $SPLIT \
--prompt_field $PROMPT_FIELD \
--response_field $RESPONSE_FIELD \
--shuffle \
--output_dir $OUTPUT_DIR \
--max_length $MAX_LENGTH \
--learning_rate $LEARNING_RATE \
--num_epochs $NUM_EPOCHS \
--save_strategy $SAVE_STRATEGY \
--eff_batch_size $EFF_BATCH_SIZE \
--per_device_batch_size $PER_DEVICE_BATCH_SIZE \
--optim $OPTIM \
--lr_scheduler $LR_SCHEDULER \
--save_total_limit $SAVE_TOTAL_LIMIT \
--report_to $REPORT_TO \
--logging_steps $LOGGING_STEPS \
--bf16 \
--gradient_checkpointing \
--warmup_ratio $WARMUP_RATIO \
--seed $SEED \
--run_name $RUN_NAME \
--dataset_mode $DATASET_MODE \
--completion_only_loss

```
sft.py
```python

import os
import torch
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import SFTConfig, SFTTrainer
import argparse

def get_config(args, gradient_accumulation_steps):
training_args = SFTConfig(
output_dir=args.output_dir,
per_device_train_batch_size=args.per_device_batch_size,
gradient_accumulation_steps=gradient_accumulation_steps,
learning_rate=args.learning_rate,
num_train_epochs=args.num_epochs,
logging_steps=args.logging_steps,
save_strategy=args.save_strategy,
save_total_limit=args.save_total_limit,
bf16=args.bf16,
gradient_checkpointing=args.gradient_checkpointing,
optim=args.optim,
lr_scheduler_type=args.lr_scheduler,
warmup_ratio=args.warmup_ratio,
packing=args.packing,
seed=args.seed,
report_to=args.report_to,
run_name=args.run_name,
max_length=args.max_length,
completion_only_loss=args.completion_only_loss,
)
return training_args

def format_dataset_messages(example, prompt_field, response_field):
"""Convert the dataset to the TRL conversational format."""
if prompt_field in example and response_field in example:
messages = [{"role": "user", "content": example[prompt_field]}, {"role": "assistant", "content": example[response_field]}]
return {"messages": messages}
else:
raise ValueError(f"Dataset must have {prompt_field} and {response_field} fields or give the correct fields!:)")

def format_dataset_pc(example, prompt_field, response_field):
if prompt_field in example and response_field in example:
prompt = [{"role": "user", "content": example[prompt_field]}]
completion = [{"role": "assistant", "content": example[response_field]}]
return {"prompt": prompt, "completion": completion}
else:
raise ValueError(f"Dataset must have {prompt_field} and {response_field} fields or give the correct fields!:)")

def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--model_name", type=str, default="meta-llama/Llama-3.2-3B-Instruct")
parser.add_argument("--dataset_name", type=str, default="rrvaswin/llama_star")
parser.add_argument("--split", type=str, default="train")
parser.add_argument("--prompt_field", type=str, default="question")
parser.add_argument("--response_field", type=str, default="response")
parser.add_argument("--shuffle", action="store_true")
parser.add_argument("--output_dir", type=str, default="models/llama_star_sft")
parser.add_argument("--learning_rate", type=float, default=5e-6)
parser.add_argument("--eff_batch_size", type=int, default=16)
parser.add_argument("--per_device_batch_size", type=int, default=4)
parser.add_argument("--num_epochs", type=int, default=3)
parser.add_argument("--optim", type=str, default="adamw_torch")
parser.add_argument("--lr_scheduler", type=str, default="cosine")
parser.add_argument("--save_strategy", type=str, default="epoch")
parser.add_argument("--save_total_limit", type=int, default=3)
parser.add_argument("--report_to", type=str, default="wandb")
parser.add_argument("--logging_steps", type=int, default=10)
parser.add_argument("--bf16", action="store_true")
parser.add_argument("--gradient_checkpointing", action="store_true")
parser.add_argument("--warmup_ratio", type=float, default=0.1)
parser.add_argument("--packing", action="store_true")
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--run_name", type=str, default="llama_star_sft")
parser.add_argument("--max_length", type=int, default=2048)
parser.add_argument("--completion_only_loss", action="store_true")
parser.add_argument("--dataset_mode", type=str, default="pc_mode", choices=["pc_mode", "messages_mode"])
return parser.parse_args()

def main():
args = parse_args()

if args.dataset_mode == "messages_mode" and args.completion_only_loss:
raise ValueError(f"Completion only loss is not supported in messages mode. By default it will calculate the over the entire conversation. So please don't use completion only loss in messages mode. But you can use \"assistant_only_loss\" to calculate the loss only over the assistant responses, provided the presence of {{generation}} tags in chat_template.")

print(f"\033[95mLoading Dataset from {args.dataset_name}...\033[0m") # Magenta colored text
dataset = load_dataset(args.dataset_name, split=args.split)
print(f"\033[95mDataset loaded with {len(dataset)} examples...\033[0m") # Magenta colored text

print(f"\033[95mLoading Model and Tokenizer from {args.model_name}...\033[0m") # Magenta colored text
model = AutoModelForCausalLM.from_pretrained(args.model_name,dtype=torch.bfloat16,attn_implementation="flash_attention_2")
tokenizer = AutoTokenizer.from_pretrained(args.model_name)
print(f"\033[95mModel loaded with {model.num_parameters()} parameters...\033[0m") # Magenta colored text
print(f"\033[95mTokenizer loaded with {tokenizer.vocab_size} tokens...\033[0m") # Magenta colored text

if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
model.config.use_cache = False

print("\033[95mFormatting Dataset...\033[0m") # Magenta colored text

if args.dataset_mode == "messages_mode":
dataset = dataset.map(format_dataset_messages, fn_kwargs={"prompt_field": args.prompt_field, "response_field": args.response_field}, remove_columns=[col for col in dataset.column_names if col != "messages"])
elif args.dataset_mode == "pc_mode":
dataset = dataset.map(
format_dataset_pc,
fn_kwargs={"prompt_field": args.prompt_field, "response_field": args.response_field},
remove_columns=[col for col in dataset.column_names if col not in ["prompt", "completion"]]
)
else:
raise ValueError(f"Invalid dataset mode: {args.dataset_mode}. Please choose from 'messages_mode' or 'pc_mode'.")

print(f"\033[95mDataset formatted with {len(dataset)} examples...\033 with the columns only being {dataset.column_names}...\033[0m") # Magenta colored text

per_device_batch_size = args.per_device_batch_size
num_gpus = int(os.environ.get("WORLD_SIZE", 1))
gradient_accumulation_steps = max(1, args.eff_batch_size // (per_device_batch_size * num_gpus))
print(f"\033[95mTraining config: {num_gpus} GPUs, batch_size={per_device_batch_size}, grad_accum={gradient_accumulation_steps}...\033[0m") # Magenta colored text

sft_config = get_config(args, gradient_accumulation_steps)

trainer = SFTTrainer(
model=model,
processing_class=tokenizer,
train_dataset=dataset,
args=sft_config,
)

print("\033[95mStarting training...\033[0m") # Magenta colored text
trainer.train()

print("\033[95mSaving model as a final checkpoint...\033[0m") # Magenta colored text
trainer.save_model()
tokenizer.save_pretrained(args.output_dir)

print(f"\033[95mTraining complete! Model saved to {args.output_dir}...\033[0m") # Magenta colored text

if __name__ == "__main__":
main()
```

Looking at the tokenize_fn in SFTTrainer:

```python

def tokenize_fn(example, processing_class, dataset_text_field, assistant_only_loss):
if "prompt" in example: # prompt-completion case
output = {}
if is_conversational(example):
if self._is_vlm:
prompt = prepare_multimodal_messages(example["prompt"], images=[])
completion = prepare_multimodal_messages(example["completion"], images=[])
else:
prompt = example["prompt"]
completion = example["completion"]
prompt_ids = processing_class.apply_chat_template(
prompt,
tools=example.get("tools"),
add_generation_prompt=True,
tokenize=True,
return_dict=False,
**example.get("chat_template_kwargs", {}),
)
# Fix transformers inconsistency: for VLMs, apply_chat_template returns lists of lists
# even for single examples, while for LLMs it returns lists of ints.
prompt_ids = prompt_ids[0] if isinstance(prompt_ids[0], list) else prompt_ids
prompt_completion_processed = processing_class.apply_chat_template(
prompt + completion,
tools=example.get("tools"),
tokenize=True,
return_dict=True,
return_assistant_tokens_mask=assistant_only_loss,
**example.get("chat_template_kwargs", {}),
)
# Fix transformers inconsistency: for VLMs, apply_chat_template returns lists of lists
# even for single examples, while for LLMs it returns lists of ints.
prompt_completion_processed = {
k: v[0] if isinstance(v[0], list) else v
for k, v in prompt_completion_processed.items()
}
prompt_completion_ids = prompt_completion_processed["input_ids"]
if "assistant_masks" in prompt_completion_processed:
output["assistant_masks"] = prompt_completion_processed["assistant_masks"]
else:
prompt_ids = processing_class(text=example["prompt"])["input_ids"]
prompt_completion_ids = processing_class(text=example["prompt"] + example["completion"])[
"input_ids"
]
# Fix transformers inconsistency: for VLMs, processing_class returns lists of lists
# even for single examples, while for LLMs it returns lists of ints.
prompt_ids = prompt_ids[0] if isinstance(prompt_ids[0], list) else prompt_ids
prompt_completion_ids = (
prompt_completion_ids[0]
if isinstance(prompt_completion_ids[0], list)
else prompt_completion_ids
)

# Check if the tokenized prompt starts with the tokenized prompt+completion
if not prompt_completion_ids[: len(prompt_ids)] == prompt_ids:
logger.warning(
"Mismatch between tokenized prompt and the start of tokenized prompt+completion. "
"This may be due to unexpected tokenizer behavior, whitespace issues, or special "
"token handling. Verify that the tokenizer is processing text consistently."
)

# Create completion mask
completion_mask = [0] * len(prompt_ids) + [1] * (len(prompt_completion_ids) - len(prompt_ids))
output["input_ids"] = prompt_completion_ids
output["completion_mask"] = completion_mask

else: # language modeling case
if is_conversational(example):
if self._is_vlm:
messages = prepare_multimodal_messages(example["messages"], images=[])
else:
messages = example["messages"]
processed = processing_class.apply_chat_template(
messages,
tools=example.get("tools"),
tokenize=True,
return_dict=True,
return_assistant_tokens_mask=assistant_only_loss,
**example.get("chat_template_kwargs", {}),
)
# Fix transformers inconsistency: for VLMs, apply_chat_template returns lists of lists
# even for single examples, while for LLMs it returns lists of ints.
processed = {k: v[0] if isinstance(v[0], list) else v for k, v in processed.items()}
output = {k: processed[k] for k in ("input_ids", "assistant_masks") if k in processed}
else:
output = {"input_ids": processing_class(text=example[dataset_text_field])["input_ids"]}

if "assistant_masks" in output and 1 not in output["assistant_masks"]:
raise RuntimeError(
"You're using `assistant_only_loss=True`, but at least one example has no assistant "
"tokens. This usually means the tokenizer's chat template doesn't generate assistant "
"masks — it may be missing the `{% generation %}` keyword. Please check the template and "
"ensure it's correctly configured to support assistant masking."
)
return output
```
If we have the input in messages format (messages mode) , then the model calculates the loss over both prompt+completion (pc_mode). If we train in prompt-completion case, then the model calculates loss over only the completion part.

However, in both the cases, the model's performance drops!

Performance (GSM8k Test Set)

| Setting | Accuracy |
|---------------------------------|---------|
| trl==0.13.0(Using DataCollatorForCompletionOnly) | 72.67 |
| messages_mode (Loss over prompt and completion) | 46.08 |
| pc_mode (Loss over completion_only) | 45.90 |

### System Info

- Platform: Linux-5.15.0-151-generic-x86_64-with-glibc2.35
- Python version: 3.12.9
- TRL version: 0.28.0
- PyTorch version: 2.10.0
- accelerator(s): NVIDIA H100 NVL, NVIDIA H100 NVL, NVIDIA H100 NVL, NVIDIA H100 NVL, NVIDIA H100 NVL, NVIDIA H100 NVL, NVIDIA H100 NVL, NVIDIA H100 NVL
- Transformers version: 5.2.0
- Accelerate version: 1.12.0
- Accelerate config:
- compute_environment: LOCAL_MACHINE
- distributed_type: DEEPSPEED
- mixed_precision: bf16
- use_cpu: False
- debug: True
- num_processes: 2
- machine_rank: 0
- num_machines: 1
- rdzv_backend: static
- same_network: True
- main_training_function: main
- enable_cpu_affinity: False
- deepspeed_config: {'gradient_accumulation_steps': 16, 'gradient_clipping': 1.0, 'offload_optimizer_device': 'cpu', 'offload_param_device': 'cpu', 'zero3_init_flag': True, 'zero3_save_16bit_model': False, 'zero_stage': 3}
- downcast_bf16: no
- tpu_use_cluster: False
- tpu_use_sudo: False
- tpu_env: []
- dynamo_config: {'dynamo_backend': 'INDUCTOR'}
- Datasets version: 4.5.0
- HF Hub version: 1.4.1
- bitsandbytes version: not installed
- DeepSpeed version: 0.18.6
- Liger-Kernel version: not installed
- LLM-Blender version: not installed
- OpenAI version: not installed
- PEFT version: not installed
- vLLM version: not installed

### 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