LLM Gradient Accumulation in Production: Token-Normalized Loss, DDP no_sync, and Scheduler Alignment to Avoid Equivalent Batch Distortion

Background: Memory Saved, but the Equivalent Batch May Have Changed

Gradient Accumulation is often understood as a simple memory-saving technique: split a large batch into multiple micro-batches, perform several forward and backward passes consecutively, and call optimizer.step() only at the end of the accumulation window.

In simple tasks with fixed-length sequences, single GPU, and full precision, dividing each micro-batch’s loss by the number of accumulation steps usually yields results close to training with a large batch. However, LLM fine-tuning frequently involves all of the following conditions:

• Each sample has a different number of effective tokens;
• Prompts, padding, and masked labels do not contribute to the loss;
• The number of effective tokens varies across DDP ranks;
• BF16, FP16, or GradScaler is used;
• Learning rate, warmup, logging, and checkpointing are driven by steps;
• The end of an epoch may have an incomplete accumulation window.

When these factors combine, “same number of micro-batches” does not equal “same data weight seen by the optimizer”. The most insidious problem is not a program crash, but that training runs normally, loss keeps decreasing, yet it is no longer mathematically equivalent to the intended large-batch baseline.

Core Principle: LLM Effective Batch Should Be Calculated by Trainable Tokens

Sample Count is Just a Shell; the True Denominator of Loss is the Number of Effective Labels

The common formula for equivalent batch size is:

samples_per_update = micro_batch_size × accumulation_steps × data_parallel_world_size

This is suitable for describing throughput and DataLoader behavior, but it is insufficient for autoregressive language models. Causal LM cross-entropy typically only computes over positions where labels != -100. Therefore, the true statistical quality of a single update is closer to:

valid_tokens_per_update = total tokens with labels != -100 across all ranks and all micro-batches

Suppose two micro-batches have 100 and 900 effective tokens respectively. If you first compute the Mean Loss for each micro-batch, then sum the two means and divide by 2, each micro-batch gets 50% weight. The correct full-batch loss should give them 10% and 90% weight respectively.

Therefore, for token-level tasks, the preferred approach is:

loss = sum of loss over all effective tokens / total number of all effective tokens

Instead of:

loss = average of mean_loss over each micro-batch

Hugging Face, when fixing gradient accumulation issues, explicitly stated that token-level tasks like Causal LM need to sum the loss over the entire accumulation window and divide by the total number of non-padding tokens, rather than averaging the already normalized losses of individual micro-batches.

The Denominator Must Be Aggregated Across Ranks

DDP by default averages gradients across ranks. If the number of effective tokens varies significantly per rank, normalizing only by local token counts distorts the gradient weights across ranks.

A clear implementation approach is:

1. Each rank counts the local effective tokens in its accumulation window;
2. Use all_reduce(SUM) to get the global effective token count;
3. For each micro-batch, use reduction="sum" to compute the local Loss Sum;
4. Since DDP will eventually divide by World Size, multiply the local loss by world_size / global_valid_tokens;
5. The last backward pass of the accumulation window triggers gradient synchronization.

This ensures the final gradient corresponds to the average loss over all global effective tokens.

Engineering Implementation: Treat One Optimizer Update as the Smallest Transaction

The code below demonstrates the core structure for BF16 + DDP scenarios, emphasizing boundary and normalization logic:

from contextlib import nullcontext
import torch
import torch.distributed as dist
import torch.nn.functional as F

IGNORE_INDEX = -100

def causal_lm_loss_sum(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
    shift_logits = logits[..., :-1, :].contiguous().float()
    shift_labels = labels[..., 1:].contiguous()
    return F.cross_entropy(
        shift_logits.view(-1, shift_logits.size(-1)),
        shift_labels.view(-1),
        ignore_index=IGNORE_INDEX,
        reduction="sum",
    )

def train_update(ddp_model, microbatches, optimizer, scheduler, max_grad_norm: float):
    """microbatches collectively form one optimizer update."""
    optimizer.zero_grad(set_to_none=True)
    device = next(ddp_model.parameters()).device
    local_tokens = torch.zeros((), device=device, dtype=torch.long)

    for batch in microbatches:
        local_tokens += batch["labels"].ne(IGNORE_INDEX).sum()

    global_tokens = local_tokens.clone()
    if dist.is_initialized():
        dist.all_reduce(global_tokens, op=dist.ReduceOp.SUM)
    world_size = dist.get_world_size() if dist.is_initialized() else 1

    if global_tokens.item() == 0:
        raise ValueError("Accumulation window contains no trainable tokens")

    denominator = global_tokens.to(torch.float32)

    for index, batch in enumerate(microbatches):
        is_last_microbatch = index == len(microbatches) - 1
        sync_context = (
            nullcontext() if is_last_microbatch else ddp_model.no_sync()
        )
        # DDP requires forward to also be within the no_sync context
        with sync_context:
            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                outputs = ddp_model(
                    input_ids=batch["input_ids"],
                    attention_mask=batch.get("attention_mask"),
                )
            loss_sum = causal_lm_loss_sum(outputs.logits, batch["labels"])
            # DDP averages gradients across ranks, so multiply by world_size to cancel that average
            loss = loss_sum * world_size / denominator
            loss.backward()

    grad_norm = torch.nn.utils.clip_grad_norm_(
        ddp_model.parameters(), max_norm=max_grad_norm,
    )
    optimizer.step()
    scheduler.step()

    return {
        "global_valid_tokens": int(global_tokens.item()),
        "microbatches": len(microbatches),
        "grad_norm": float(grad_norm),
    }

This code has three key constraints:

• Loss uses reduction="sum", normalized uniformly at the accumulation window level;
• Intermediate micro-batches use no_sync(), and only the last micro-batch triggers DDP synchronization;
• optimizer.step() and scheduler.step() are executed only once at the complete update boundary.

Note: If using Hugging Face Trainer, Accelerate, DeepSpeed, or FSDP, prioritize following the framework’s own loss, synchronization, and step contracts. Do not mechanically apply the World Size compensation formula on top of the framework’s already implemented normalization logic.

Correct Boundaries for no_sync

PyTorch DDP’s no_sync() pauses cross-process gradient synchronization and keeps gradients local until the first Forward/Backward that exits this context triggers synchronization.

A common mistake is to only place backward() inside no_sync() while leaving Forward outside. The PyTorch documentation explicitly requires that Forward must also be within the context; otherwise, gradients might still be synchronized.

Recommended rules:

Micro-batch Position	Forward	Backward	Gradient Sync
First N-1	Inside no_sync()	Inside no_sync()	Not triggered
Nth	Normal context	Normal context	Trigger cumulative sync

Each rank must enter and exit the synchronization boundary in a consistent order.

AMP, Gradient Clipping, and Step Skipping Must Be Within the Same Update Boundary

Under FP16 + GradScaler, gradients within the accumulation window should maintain the same scale. PyTorch’s AMP documentation requires:

• For intermediate micro-batches, only execute scaler.scale(loss).backward();
• After accumulation is complete, call scaler.unscale_(optimizer) once;
• Perform gradient clipping after unscaling;
• Finally, call scaler.step(optimizer) and scaler.update();
• The scale must not change within a valid batch, and scaled gradients must not be mixed with unscaled gradients.

If Inf/NaN is detected, GradScaler may skip the current optimizer step. In this case, the scheduler should not advance independently; otherwise, the learning rate timeline will lead the parameter update count. Production code should use “whether the optimizer actually updated” as the common signal for the scheduler, EMA, global update ID, and checkpoint naming.

Tail Windows: Do Not Continue Dividing by a Fixed Accumulation Steps

At the end of an epoch, there are often only 1 to accumulation_steps - 1 micro-batches left. If you still divide by the fixed full accumulation steps, the tail gradient will be artificially reduced.

Two safer strategies are:

1. Dynamic Window: Normalize by the actual total effective tokens in the tail and perform a normal update;
2. Drop Tail: Use drop_last or pad at the data layer to make all update windows fixed, but explicitly record dropped or repeated data.

For streaming datasets, pay extra attention: check whether the framework automatically triggers synchronization and step at the end of the DataLoader. Accelerate’s Gradient Accumulation Plugin defaults to synchronizing with the current DataLoader end state; for custom infinite streams or cross-epoch accumulation, configure this behavior explicitly.

Scheduler and Training Budget: Unify Using optimizer_update_id

There are at least three types of steps in a training system:

Step Type	Trigger Frequency	Use Case
Micro Step	Once per micro-batch	Debugging, fine-grained logging only
Optimizer Update	Once per accumulation window	Warmup, LR Scheduler, EMA, Checkpoint
Token Step	Once per cumulative effective token count	Throughput statistics, cross-experiment alignment

Warmup, LR Scheduler, weight decay, EMA, and checkpoint intervals should typically be based on Optimizer Update; throughput, data budget, and cross-experiment alignment are better served by also recording cumulative effective tokens.

It is recommended to fix the following fields in training logs:

micro_step_id
optimizer_update_id
epoch
global_valid_tokens_in_update
cumulative_train_tokens
learning_rate
optimizer_step_skipped
grad_norm_before_clip
grad_norm_after_clip

Do not let a single ambiguous global_step carry three different semantics.

Applicable Scenarios

The solutions in this article are particularly suitable for:

• Causal LM or SFT on variable-length dialogue data;
• Tasks where prompt parts use -100 masking and only assistant outputs are trained;
• Sequence packing where each pack has a different number of effective training tokens;
• Multi-GPU DDP/FSDP where length distributions across ranks may not be perfectly balanced;
• Reproducible experiments comparing different micro-batch sizes or accumulation steps.

For sample-level objectives like classification, ranking, or preference learning, the denominator may not need to be token count. The core principle is not “divide by tokens for all tasks,” but to choose the correct statistical unit based on the objective function definition and ensure consistency before and after accumulation.

Common Misconceptions

Misconception 1: Dividing Mean Loss by Accumulation Steps is Always Correct

This is only valid when the statistical unit and effective count are the same across micro-batches. For variable-length token tasks, unify normalization over the loss sum and total effective tokens.

Misconception 2: Synchronizing Every Micro-batch is Just a Performance Issue

It usually doesn’t change the mathematical result, but it introduces unnecessary All-Reduce operations, significantly undermining the value of gradient accumulation in reducing communication frequency.

Misconception 3: The Scheduler Can Advance by DataLoader Batch

This causes the learning rate to advance with micro-steps rather than parameter updates. After changing accumulation steps, training curves become directly incomparable.

Misconception 4: Gradient Clipping on Every Micro-batch

This alters the direction and magnitude of the summed vectors. Gradient clipping should be applied to the final gradient of the complete accumulation window.

Misconception 5: Using a Fixed Divisor for Tail Windows

Incomplete windows must use actual statistics or be explicitly discarded; do not pretend they still contain the full number of micro-batches.

Production Checklist

• Compare parameter update differences between a “single large batch step” and “multiple micro-batch accumulation step” using the same random seed;
• Construct samples with vastly different lengths to verify short samples are not amplified with equal weight;
• Check that global effective token counts are aggregated across ranks;
• Profile the number of All-Reduce operations per optimizer update to confirm intermediate micro-batches are not synchronized;
• Verify that no_sync() covers both Forward and Backward;
• Confirm tail window, empty label window, and DataLoader end behavior;
• Under FP16, confirm the order: Unscale → Clip → Step → Update;
• When overflow skips a step, ensure the scheduler, EMA, and update ID do not advance;
• Log micro-step, optimizer update, and cumulative effective tokens simultaneously;
• After changing gradient_accumulation_steps, re-verify warmup and total update count.

References

• Hugging Face: Fixing Gradient Accumulation
• Hugging Face Accelerate: Performing Gradient Accumulation
• PyTorch DistributedDataParallel: no_sync
• PyTorch Automatic Mixed Precision Examples
• Hugging Face Transformers Trainer