How to Configure FP32 Gradient Accumulation in Nanotron

Set accumulate_grad_in_fp32: true in your optimizer configuration with zero_stage >= 1 to enable full-precision gradient accumulation across data-parallel ranks.

When training large language models with mixed precision (bfloat16 or float16) in the Hugging Face Nanotron framework, numerical stability can degrade during gradient accumulation. Nanotron solves this by allowing gradients to accumulate in full FP32 precision before synchronization across data-parallel (DP) ranks. This feature integrates with ZeRO optimization to maintain high numerical precision while keeping model weights in memory-efficient half-precision formats.

Prerequisites for FP32 Gradient Accumulation

FP32 gradient accumulation in Nanotron has specific requirements that must be met for the feature to activate:

• zero_stage must be greater than 0. The FP32 accumulation hook is built on top of the ZeRO-distributed optimizer infrastructure. Without ZeRO, Nanotron falls back to the standard DDP all-reduce path.
• accumulate_grad_in_fp32 must be set to True in your optimizer arguments.

If you enable the flag but set zero_stage: 0, Nanotron will silently ignore the FP32 accumulation setting and use default half-precision accumulation instead.

Configuration Methods

You can enable FP32 gradient accumulation either through YAML configuration files or programmatically via Python.

YAML Configuration

Add the accumulate_grad_in_fp32 flag to your optimizer configuration block:

optimizer:
  optimizer_factory:
    name: adamW
    adam_eps: 1.0e-8
    adam_beta1: 0.9
    adam_beta2: 0.95
    torch_adam_is_fused: true
  zero_stage: 1               # Required: Enable ZeRO stage 1

  weight_decay: 0.1
  clip_grad: null
  accumulate_grad_in_fp32: true   # Enable FP32 accumulation

  learning_rate_scheduler:
    name: cosine
    learning_rate: 5e-4
    warmup_steps: 1000

This flag is parsed in src/nanotron/config/config.py within the OptimizerArgs data class.

Programmatic Configuration

When building your training configuration in Python, set the flag in your OptimizerArgs instantiation:

from nanotron.config import Config, OptimizerArgs, AdamWOptimizerArgs, LRSchedulerArgs

cfg = Config(
    general=...,
    parallelism=...,
    model=...,
    optimizer=OptimizerArgs(
        optimizer_factory=AdamWOptimizerArgs(
            adam_eps=1e-8,
            adam_beta1=0.9,
            adam_beta2=0.95,
            torch_adam_is_fused=True,
            name="adamW",
        ),
        zero_stage=1,                     # Required: ZeRO stage 1+

        weight_decay=0.1,
        clip_grad=None,
        accumulate_grad_in_fp32=True,    # Enable FP32 accumulation

        learning_rate_scheduler=LRSchedulerArgs(
            name="cosine",
            learning_rate=5e-4,
            warmup_steps=1000,
        ),
    ),
    # …

)

When DistributedTrainer initializes with this configuration, the helper chain in nanotron.helpers.build_optimizer automatically constructs the required gradient accumulator wrapper.

Internal Implementation

Understanding the internal flow helps debug configuration issues and customize advanced training loops.

Configuration and Building

In src/nanotron/config/config.py, the OptimizerArgs class defines the accumulate_grad_in_fp32: bool attribute. When the trainer initializes, src/nanotron/helpers.py (lines 390-405) checks this flag and builds an OptimizerFromGradientAccumulator that wraps your base optimizer (AdamW, SGD, etc.) together with an FP32GradientAccumulator.

Gradient Accumulation Engine

The FP32GradientAccumulator class in src/nanotron/optim/gradient_accumulator.py maintains full-precision copies of each parameter and matching FP32 gradient buffers. During the backward pass, half-precision gradients are added into these FP32 buffers, preventing the numerical drift that occurs when accumulating low-precision gradients directly.

Communication Hook Registration

In src/nanotron/helpers.py (lines 442-462), Nanotron registers a custom DDP communication hook (get_fp32_accum_hook) on the DistributedDataParallel model. This hook ensures gradients are accumulated in FP32 during the backward pass before any cross-rank synchronization occurs.

DDP Wrapping Guard

The trainer in src/nanotron/trainer.py (lines 93-95) contains a guard that disables standard DDP wrapping when FP32 accumulation is active and ZeRO stage is greater than 0. This prevents duplicate gradient synchronization, as the custom hook already handles gradient reduction across data-parallel ranks.

Verifying FP32 Accumulation is Active

After trainer initialization, inspect the optimizer and gradient accumulator objects to confirm configuration:

from nanotron.optim.gradient_accumulator import FP32GradientAccumulator

# After trainer creation

trainer = DistributedTrainer("path/to/config.yaml")
opt, grad_acc = trainer.optimizer, trainer.gradient_accumulator

print(type(opt))               

# <class 'nanotron.optim.optimizer_from_gradient_accumulator.OptimizerFromGradientAccumulator'>

print(isinstance(grad_acc, FP32GradientAccumulator))

# True → gradients will be accumulated in FP32

If isinstance returns False, verify that zero_stage is set correctly in your configuration.

Manual Hook Registration (Advanced)

For custom training loops or model wrappers outside the standard helper flow, manually register the FP32 accumulation hook:

from nanotron.optim.gradient_accumulator import FP32GradientAccumulator, FP32GradBucketManager, get_fp32_accum_hook
import torch.distributed as dist

grad_acc = FP32GradientAccumulator(named_parameters=model.named_parameters())
ddp_model = torch.nn.parallel.DistributedDataParallel(
    model,
    process_group=parallel_context.dp_cp_pg,
    broadcast_buffers=False,
)

ddp_model.register_comm_hook(
    state=FP32GradBucketManager(
        dp_cp_pg=parallel_context.dp_cp_pg,
        accumulator=grad_acc,
        param_id_to_name={id(p): n for n, p in model.named_parameters()},
    ),
    hook=get_fp32_accum_hook(
        reduce_scatter=False,   # Set True if using ZeRO-2/3 reduce-scatter

        reduce_op=dist.ReduceOp.AVG,
    ),
)

This manual approach mirrors the automatic setup performed by nanotron.helpers.build_optimizer.

Summary

• Enable FP32 accumulation by setting accumulate_grad_in_fp32: true in your optimizer configuration.
• Require ZeRO: You must set zero_stage to 1 or higher; otherwise, the setting is ignored.
• Core files: Configuration is defined in src/nanotron/config/config.py, built in src/nanotron/helpers.py, and executed by src/nanotron/optim/gradient_accumulator.py.
• Verification: Check that your trainer's optimizer is an instance of OptimizerFromGradientAccumulator and the gradient accumulator is FP32GradientAccumulator.
• No DDP conflicts: The trainer automatically disables standard DDP wrapping when FP32 accumulation is active to avoid duplicate gradient synchronization.

Frequently Asked Questions

Why does FP32 gradient accumulation require ZeRO in Nanotron?

The FP32 accumulation hook is implemented on top of the ZeRO-distributed optimizer infrastructure, specifically using the OptimizerFromGradientAccumulator wrapper. Without ZeRO (zero_stage = 0), Nanotron defaults to standard DDP all-reduce operations that do not support the custom gradient buffer management required for full-precision accumulation.

What happens if I enable accumulate_grad_in_fp32 but forget to set zero_stage?

Nanotron will silently ignore the FP32 accumulation setting and proceed with default half-precision gradient accumulation. The trainer checks this condition during initialization and falls back to standard behavior to prevent runtime errors, but you will lose the numerical stability benefits of FP32 accumulation.

How do I know if gradients are actually being accumulated in FP32?

After initializing the DistributedTrainer, inspect trainer.gradient_accumulator. If it is an instance of FP32GradientAccumulator (from src/nanotron/optim/gradient_accumulator.py), the feature is active. Additionally, the optimizer type should be OptimizerFromGradientAccumulator, indicating the wrapper is managing the FP32 buffers.

Can I use FP32 gradient accumulation with custom optimizers or learning rate schedulers?

Yes. The OptimizerFromGradientAccumulator wrapper is optimizer-agnostic and works with any base optimizer (AdamW, SGD, Adam, etc.) specified in your configuration. The FP32 accumulation logic operates on the gradient buffers before the optimizer step, so learning rate schedules and optimizer-specific hyperparameters remain unaffected.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}