# How to Configure FP32 Gradient Accumulation in Nanotron

> Unlock superior model training by enabling FP32 gradient accumulation in Nanotron. Simply set accumulate_grad_in_fp32: true in your optimizer config with zero_stage 1 or higher. Train with higher precision.

- Repository: [Hugging Face/nanotron](https://github.com/huggingface/nanotron)
- Tags: how-to-guide
- Published: 2026-03-03

---

**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:

```yaml
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`](https://github.com/huggingface/nanotron/blob/main/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:

```python
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`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/config/config.py), the `OptimizerArgs` class defines the `accumulate_grad_in_fp32: bool` attribute. When the trainer initializes, [`src/nanotron/helpers.py`](https://github.com/huggingface/nanotron/blob/main/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`](https://github.com/huggingface/nanotron/blob/main/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`](https://github.com/huggingface/nanotron/blob/main/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`](https://github.com/huggingface/nanotron/blob/main/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:

```python
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:

```python
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`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/config/config.py), built in [`src/nanotron/helpers.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/helpers.py), and executed by [`src/nanotron/optim/gradient_accumulator.py`](https://github.com/huggingface/nanotron/blob/main/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`](https://github.com/huggingface/nanotron/blob/main/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.