Loss Scales Incorrectly with Gradient Accumulation Steps in Trainer (Gemma3 and other models)

[yevvonlim]

System Info

• transformers version: 4.53
• Python version: 3.12
• Platform: Linux
• Models affected: google/gemma-3-12b-it (confirmed), likely affects other models
• Training setup: torchrun with DeepSpeed, LoRA fine-tuning

Description

The HuggingFace Trainer exhibits incorrect loss scaling behavior when using gradient accumulation, where logged loss values increase proportionally with gradient_accumulation_steps instead of remaining consistent. This affects model training monitoring, convergence analysis, and hyperparameter tuning.

Who can help?

@zach-huggingface @SunMarc

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Minimal Example

from transformers import Trainer, TrainingArguments, AutoTokenizer, AutoModelForCausalLM
from datasets import Dataset

# Load model and tokenizer
model = AutoModelForCausalLM.from_pretrained("google/gemma-3-12b-it")
tokenizer = AutoTokenizer.from_pretrained("google/gemma-3-12b-it")
tokenizer.padding_side = "right" 

# Create dummy dataset
dummy_data = {
    "input_ids": [[1] * 100] * 1000,  # Simple repeated tokens
    "labels": [[1] * 100] * 1000
}
dataset = Dataset.from_dict(dummy_data)

# Test with different gradient accumulation steps
for ga_steps in [1, 2, 4, 8]:
    trainer = Trainer(
        model=model,
        train_dataset=dataset,
        args=TrainingArguments(
            output_dir=f"./test_ga_{ga_steps}",
            per_device_train_batch_size=1,
            gradient_accumulation_steps=ga_steps,
            max_steps=10,
            logging_steps=1,
            save_steps=1000,  
        )
    )
    trainer.train()

Expected behavior

Loss values should remain mathematically consistent regardless of gradient_accumulation_steps, as the effective batch size and optimization steps should be equivalent.

Actual Behavior

• gradient_accumulation_steps=8: Loss = X
• gradient_accumulation_steps=16: Loss = 2X (exactly double)
• gradient_accumulation_steps=32: Loss = 4X

Root Cause Analysis

The issue stems from the model_accepts_loss_kwargs parameter in the Trainer class:
this issue has been raised by here but it is not properly handled yet (transformers 4.56.0 has same issue)

[effortprogrammer]

I encountered a similar issue before. It would be helpful if you could include which DeepSpeed version you're using and which ZeRO stage you're currently on.

[yevvonlim]

I had this issue with DeepSpeed 0.16.2 when using both ZeRO-2 and ZeRO-3, but I believe the root cause is indeed related to the self.model_accepts_loss_kwargs variable set in here.

According to the official documentation:

If you are not using num_items_in_batch when computing your loss, make sure to overwrite self.model_accepts_loss_kwargs to False. Otherwise, the loss calculation might be slightly inaccurate when performing gradient accumulation.

However, this variable is automatically set during Trainer initialization, and in practice, it is almost always set to True, since the check is based only on whether the model’s forward method accepts **kwargs in its signature. This naive detection can cause problems, as the assumption may not reflect the intended usage in all cases.

[SunMarc]

Hey @yevvonlim, for gemma3, we've set accepts_loss_kwargs to False, so normally it shouldn't happen anymore. Can you double check with the latest transformers ? #37208

[yevvonlim]

(transformers) root@zoiworld88-0:~/ye/transformers# uv pip list | grep transformers
transformers             4.57.0.dev0 /home/jovyan/ye/transformers

Unfortunately I get the same issue with the brief snippet below with the current main branch.

from transformers import Trainer, TrainingArguments, AutoTokenizer, AutoModelForCausalLM
from datasets import Dataset

# Load model and tokenizer
model = AutoModelForCausalLM.from_pretrained("google/gemma-3-4b-it", device_map="auto")
tokenizer = AutoTokenizer.from_pretrained("google/gemma-3-4b-it")
tokenizer.padding_side = "right" 

# Create dummy dataset
dummy_data = {
    "input_ids": [[1] * 100] * 1000,  # Simple repeated tokens
    "labels": [[1] * 100] * 1000
}
dataset = Dataset.from_dict(dummy_data)

# Test with different gradient accumulation steps
for ga_steps in [1, 2, 4, 8]:
    trainer = Trainer(
        model=model,
        train_dataset=dataset,
        args=TrainingArguments(
            output_dir=f"./test_ga_{ga_steps}",
            per_device_train_batch_size=1,
            gradient_accumulation_steps=ga_steps,
            max_steps=10,
            logging_steps=1,
            save_steps=1000,  
            bf16=True,
        )
    )
    trainer.train()

I tried pdb and got:

(Pdb) model.accepts_loss_kwargs
*** AttributeError: 'Gemma3ForConditionalGeneration' object has no attribute 'accepts_loss_kwargs'

It turns out that Gemma3 Model has accepts_loss_kwargs variable set to False whereas Gemma3ForConditionalGeneration doesn't.

[yevvonlim]

I just made a simple patch for this ! please check this out :) @SunMarc

[gonzalobenegas]

I also ran into this issue in our custom model. Fixed it by setting accepts_loss_kwargs = False similar to the fix for gemma. I hope I don't need to do this for every new model.