Trainer floor-divides valid-token counts under model parallelism, silently mis-scaling gradients and optimizer updates #50212

System Info

• Transformers: 5.23.0.dev0
• Upstream main inspected: ca8c871
• PyTorch: 2.11.0+cu128
• Python: 3.13
• Platform: Linux
• GPUs: 2 × NVIDIA Tesla T4
• Model: Randomly initialized LlamaForCausalLM (1 decoder layer, hidden size 32)
• Precision: FP32
• Distributed configuration: FSDP2 + context parallelism, cp_size=2
• Sequence length: 64
• Optimizer: SGD, learning rate 1e-3

Who can help?

@SunMarc @3outeille

Reproduction

Summary

Trainer._get_num_items_in_batch() performs floor division when normalizing valid-token counts by the tensor/context/sequence parallelism factor.

When the number of valid targets is not divisible by that factor, the denominator is truncated, silently changing the training loss and gradients.

I reproduced this using a real two-GPU FSDP2 + context-parallel training run, including actual optimizer updates.

Root cause

In src/transformers/trainer.py:

num_items_in_batch = num_items_in_batch // (
    self.get_tp_size() * self.get_cp_size() * self.get_sp_size()
)

For three valid targets and cp_size=2, this produces 1 instead of 1.5.

The truncated denominator reaches the causal LM cross-entropy loss, causing the loss and gradients to be over-scaled.

The expression is present in upstream main at commit ca8c871e.

Reproduction

import torch
from types import SimpleNamespace
from transformers import Trainer

trainer = Trainer.__new__(Trainer)
trainer.args = SimpleNamespace(
    average_tokens_across_devices=False,
    n_gpu=1,
)
trainer.model_accepts_loss_kwargs = True
trainer.compute_loss_func = None
trainer._loss_shifts_labels = True

trainer.get_tp_size = lambda: 1
trainer.get_cp_size = lambda: 2
trainer.get_sp_size = lambda: 1

def check(valid_tokens):
    labels = torch.full((1, 8), -100, dtype=torch.long)
    labels[:, -valid_tokens:] = 5

    result = trainer._get_num_items_in_batch(
        [{"labels": labels}],
        torch.device("cpu"),
    )

    expected = valid_tokens / 2

    print(
        f"valid={valid_tokens} "
        f"actual={float(result)} "
        f"expected={expected}"
    )

check(4)
check(3)
check(1)

Output from the current implementation:

valid=4 actual=2.0 expected=2.0
valid=3 actual=1.0 expected=1.5
valid=1 actual=0.0 expected=0.5

Dual-GPU FSDP2 + context-parallel validation

The distributed experiment used:

• A randomly initialized tiny LlamaForCausalLM
• Two NVIDIA Tesla T4 GPUs
• ParallelismConfig(cp_size=2) with FSDP2
• Sequence length 64
• FP32 training
• One SGD optimizer step, learning rate 1e-3
• average_tokens_across_devices=False
• No gradient clipping
• Context-parallel load balancing disabled for the diagnostic

Both runs used identical model initialization, inputs, and training settings.

The reference differed only in replacing the floor-divided valid-token denominator with exact floating-point division.

Divisible-token control: 4 valid targets

Non-divisible-token case: 3 valid targets

Verification

• Both GPU processes completed the FSDP2 + context-parallel training step.
• The divisible-token control matched exactly.
• The non-divisible-token case produced exactly 1.5× larger loss, gradient norm, and SGD parameter-update norm.
• No Transformers source files were modified in the original run.
• The reference changed only the denominator calculation.
• A separate Trainer backward test reproduced incorrect scaling with simulated tensor, context, and sequence parallelism factors.
• A zero denominator produced non-finite loss and gradients in the single-process audit.

The real distributed validation establishes the bug for FSDP2 context parallelism. Tensor parallelism, sequence parallelism, and DataParallel should be tested separately before claiming equivalent distributed impact.

Expected behavior

Trainer should preserve fractional valid-token denominators when normalizing losses under model parallelism.

For three valid targets and cp_size=2, the effective denominator should be 1.5, not 1.

The loss, gradient magnitude, and optimizer update should not acquire an additional scaling factor solely because the valid-token count is not divisible by the parallelism factor.

A nonzero valid-token count should also not become a zero denominator due to integer truncation.

Suggested fix

Replace the relevant floor division with true division, preserving fractional denominators:

num_items_in_batch = num_items_in_batch / (
    self.get_tp_size() * self.get_cp_size() * self.get_sp_size()
)

Add regression coverage for divisible and non-divisible valid-token counts, including distributed gradient and optimizer-update comparisons.

This reproduction uses average_tokens_across_devices=False, establishing a distinct failure mechanism.

Expected behavior

Trainer should preserve fractional valid-token denominators when normalizing losses under model parallelism.

Specifically:

• Three valid targets with cp_size=2 should produce an effective normalization denominator of 1.5, not 1.
• Losses, gradients, and optimizer updates should remain correctly normalized regardless of whether the valid-token count is divisible by the parallelism factor.
• A nonzero valid-token count must not become a zero denominator solely because of parallelism normalization.
• The fix should preserve existing behavior for divisible-token cases and ordinary single-device training.

Replacing the relevant floor division with true division is a possible fix, subject to validation across supported parallelism modes.