Mixed-Precision Training

NeMo AutoModel uses FSDP2’s MixedPrecisionPolicy to control compute precision during forward and backward, and the model’s storage dtype (model.dtype) to control the precision of the resident sharded parameter. Together, these decide what numeric precision the optimizer state ends up in, which is the part that determines whether long full-parameter training runs converge cleanly.

This page describes the recommended precision patterns and the pattern to avoid. For any long full-parameter training run (pretraining or extended fine-tuning), the key rule is: do not combine torch.optim.AdamW with bf16 resident parameters unless you have explicitly accepted bf16 master weights and Adam state.

For NeMoAutoModel loaders, use model.dtype to select storage precision. The deprecated model.torch_dtype spelling remains supported. If both are supplied, a non-null dtype takes precedence, including dtype: auto; dtype: null falls back to torch_dtype. String values such as float32 and torch.float32, and Python torch.dtype values, are accepted by both spellings. Choose the storage dtype explicitly for full-parameter training; the optimizer does not select it for you. Diffusion pipelines use their separate model.torch_dtype and model.compute_dtype settings.

Storage Dtype Compared to Compute Dtype

Precision settings have distinct effects:

Setting	Controls	Effect on Optimizer State
`model.dtype`	Base model storage dtype requested at loading. FSDP2 PEFT can lower eligible frozen parameters to the compute dtype before sharding.	For trainable parameters, `torch.optim.Adam`/`AdamW` EMA buffers (`exp_avg`, `exp_avg_sq`) use the resident parameter dtype. Frozen parameters have no Adam state. Transformer Engine (TE) FusedAdam configures master weights and EMA dtypes separately.
`peft.lora_dtype`	Storage dtype of trainable LoRA adapters.	With `torch.optim.Adam`/`AdamW`, adapter EMA buffers use the adapter dtype, independently of frozen base storage.
`mp_policy.param_dtype`	Compute dtype FSDP2 casts to during forward/backward.	None directly; this only affects matmul/activation precision.
`mp_policy.reduce_dtype`	Dtype used for gradient reduce-scatter/all-reduce across data-parallel (DP) ranks.	None directly; this only affects how gradients are summed.
`mp_policy.output_dtype`	Dtype FSDP2 casts module outputs to.	None directly; this affects activation tensors, including tensors that cross pipeline-parallel boundaries.

When model.dtype: bfloat16 is used with torch.optim.AdamW, the AdamW EMA buffers (exp_avg and exp_avg_sq) are also stored in bf16. This is fragile for long full-parameter training runs: bf16 has only a 7-bit mantissa, so small EMA updates can be rounded away even though the values themselves are still in range. Symptoms range from silent degradation (slower convergence, that is, higher final loss at the same step count, with no visible instability) to overt failure: unstable grad_norm, loss bumps, loss spikes, or divergence.

Freeze Exact Parameters Before Checkpoint Loading

Use freeze_config.freeze_parameters for exact parameter exclusions alongside existing module freezes. Only named parameters are frozen; unlisted parameters retain their existing trainability. Use the loader’s dtype argument for trainable storage precision and FSDP2’s mixed-precision policy for compute.

import torch
from torch.distributed.fsdp import MixedPrecisionPolicy
from nemo_automodel import NeMoAutoModelForCausalLM
from nemo_automodel.components.distributed.config import DistributedSetup, FSDP2Config
from nemo_automodel.components.utils.model_utils import FreezeConfig
model = NeMoAutoModelForCausalLM.from_pretrained(
    "/path/to/nemotron-checkpoint",
    dtype=torch.float32,
    freeze_config=FreezeConfig(
        freeze_parameters=("model.layers.0.norm.weight",),
    ),
    distributed_setup=DistributedSetup.build(
        FSDP2Config(mp_policy=MixedPrecisionPolicy(
            param_dtype=torch.bfloat16, reduce_dtype=torch.float32,
        )),
    ),
)

This freezes the first block’s norm weight. Other parameters retain their existing trainability, including model-required freezes. Names match canonical parameter paths, ignoring activation-checkpoint and compile wrappers. Freezing any alias of a tied parameter freezes that shared parameter. Empty exclusions are a no-op; duplicate or unknown names raise an error. Exact parameter freezes are applied after module selectors, so they remain frozen even when their containing module is explicitly unfrozen.

The policy is applied before checkpoint loading and rebound after module replacement, before distributed wrappers capture parameters. Construct the optimizer after the loader returns. Exact exclusions of TE expert parameters with expert parallelism are rejected because repartitioning renumbers local expert parameters; freeze their owning module instead.

For sharded FSDP2 with exact parameter selectors, the existing frozen-storage preparation reduces ordinary frozen parameters to the safe compute dtype while retaining trainable storage and model-protected FP32 parameters. With dtype=torch.float32 and BF16 compute, trainable weights load directly into FP32 destinations without rounding through BF16. Frozen buffers follow existing compute-dtype rules and precision protections. Existing guards for differing root/MoE compute policies remain authoritative.

For an exact update surface, freeze all modules and then enable only the named parameters. This keeps unselected biases and other siblings frozen:

from nemo_automodel.components.utils.model_utils import ModuleSelector
freeze_config = FreezeConfig(
    freeze_modules=[ModuleSelector(glob="*")],
    unfreeze_parameters=("model.layers.0.norm.weight",),
)

Pass this configuration with dtype=torch.float32 and the BF16 FSDP2 compute policy shown above. Exact exclusions take precedence over exact unfreezes; framework-required freezes remain protected. Unknown and duplicate unfreeze names are rejected. Exact TE expert selections require the same restrictions as exact exclusions.

On one rank, FSDP2 skips its mixed-precision wrappers. Exact parameter freezes therefore retain uniform loader storage and compute in that storage dtype, without introducing a mixed plain-tensor dtype boundary.

Nemotron V3 also accepts compute_mtp=False during an off-pipeline-parallel forward call. This skips auxiliary MTP computation while retaining all MTP parameters for checkpoint loading and saving. Response-only RL callers can combine it with compute_logits=False, output_hidden_states=True to consume the final hidden states without materializing full-vocabulary logits. The existing empty-vocabulary logits tensor remains the return contract.

Recommended fp32 Master Weight and Optimizer-State Patterns

Use one of these patterns for long full-parameter training (pretraining or extended fine-tuning):

Pattern	Model Storage Dtype	Optimizer State	When to Use
TE FusedAdam with bf16 model storage	bf16	fp32 master weights and fp32 Adam EMA buffers	When TE has a validated memory/runtime path for the model and you want training checkpoints to stay bf16.
torch AdamW with `model.dtype: float32`	fp32	fp32 master weight (the resident param) and fp32 Adam EMA buffers	Robust starting point for new or precision-sensitive models. Parameters, gradients, and optimizer state are all fp32. Trade-off: writes fp32 training checkpoints.

Both patterns keep forward and backward compute in bf16 through FSDP2 mixed precision. They differ in where the fp32 master weight lives, how much peak memory they use for a specific model, and what dtype is written to the training checkpoint.

On a single-rank or single-element mesh, AutoModel skips FSDP2 sharding and applies the policy’s compute dtype through forward autocast instead. This also applies to direct model calls. Resident parameters and buffers retain their storage precision, and activation-checkpoint recomputation retains autocast. An explicit distributed.autocast_dtype overrides the policy’s compute dtype on this unsharded path. This path does not apply mp_policy.output_dtype. Native AutoModel implementations protect explicitly fp32 math, including routing, rotary angle construction, sparse indexer scores, and hyperconnections, inside the forward autocast region. Models loaded with force_hf=True or through the Hugging Face fallback use the upstream implementation’s autocast behavior; AutoModel does not add fp32 compatibility patches to those implementations.

Pattern A Using TE FusedAdam and bf16 Model Storage

Use Transformer Engine FusedAdam when it has been validated for the model. The resident model parameters remain bf16, so model checkpoints can stay bf16, while TE keeps the optimizer’s master weights and Adam EMA buffers in fp32.

model:
  dtype: bfloat16             # resident sharded parameter + model checkpoint in bf16
optimizer:
  _target_: transformer_engine.pytorch.optimizers.fused_adam.FusedAdam
  lr: 3.0e-4
  adam_w_mode: true
  bias_correction: true
  master_weights: true
  store_param_remainders: true
  exp_avg_dtype: torch.float32
  exp_avg_sq_dtype: torch.float32
distributed:
  strategy: fsdp2
  # Defaults already provide bf16 forward/backward + fp32 gradient reduction; this block is shown explicitly for clarity.
  mp_policy:
    _target_: torch.distributed.fsdp.MixedPrecisionPolicy
    param_dtype: bfloat16
    reduce_dtype: float32
    output_dtype: bfloat16
    cast_forward_inputs: true

TE FusedAdam is the cleanest way to request fp32 optimizer state without making the resident model parameter fp32.

It is a common misconception that TE costs extra memory for a second fp32 master. With store_param_remainders: true (as above), TE does not keep a full extra fp32 master: it stores the bf16 parameter plus a 16-bit remainder that together reconstruct the fp32 master, costing approximately 4 bytes/param, the same as torch AdamW’s fp32 resident parameter. The fp32 Adam EMA buffers (exp_avg, exp_avg_sq) are the same in both patterns, so the two optimizers’ steady-state footprints are essentially equal. The practical difference is then simplicity (torch AdamW: no TE dependency, fewer moving parts) compared to keeping model storage and checkpoints in bf16 (TE).

Where the two can differ is the gradient buffer: TE keeps the resident parameter in bf16, so its gradients are bf16, whereas torch AdamW with fp32 storage keeps parameters and gradients in fp32 (approximately 2 bytes/param more on the gradient buffer). In practice, what you measure is peak memory, which is usually dominated by activations rather than the optimizer step. Depending on the model (fraction of intrinsically fp32 parameters, fragmentation, and where the peak falls), TE can come out lower, equal, or higher. Validate memory per model before making it the default.

Lower Optimizer Memory Using bf16 Adam States

With bf16 resident parameters, TE FusedAdam can preserve fp32 master weights while storing both Adam EMA buffers in bf16. This halves the storage for those two buffers compared to fp32 moments:

optimizer:
  _target_: transformer_engine.pytorch.optimizers.fused_adam.FusedAdam
  master_weights: true
  store_param_remainders: true
  exp_avg_dtype: torch.bfloat16
  exp_avg_sq_dtype: torch.bfloat16

Lower-precision moments trade update precision for memory, so validate convergence for your model. torch.optim.AdamW does not expose independent EMA dtype settings: its moment buffers follow the parameter dtype.

Pattern B Using torch AdamW and fp32 Model Storage

model:
  dtype: float32              # sharded parameter + AdamW state in fp32
optimizer:
  _target_: torch.optim.AdamW
  lr: 3.0e-4
  betas: [0.9, 0.95]
  weight_decay: 0.1
distributed:
  strategy: fsdp2
  # Defaults already provide bf16 forward/backward + fp32 gradient reduction; this block is shown explicitly for clarity.
  mp_policy:
    _target_: torch.distributed.fsdp.MixedPrecisionPolicy
    param_dtype: bfloat16           # forward/backward compute in bf16 (fast)
    reduce_dtype: float32           # safe gradient reduction
    output_dtype: bfloat16
    cast_forward_inputs: true

This is the PyTorch AdamW version of the master-weights pattern. Forward and backward run in bf16, the all-reduce or reduce-scatter runs in fp32, and the optimizer applies updates against fp32 resident parameters. With torch.optim.AdamW, the resident fp32 sharded parameter is the master weight, so there is no separate fp32 master-weight copy.

The trade-off is the dtype of the training checkpoint: NeMo AutoModel’s Distributed Checkpoint (DCP) stores the resident model parameters, so this pattern writes them in fp32, which you keep for exact resume. This concerns only the training checkpoint; a consolidated Hugging Face (HF) checkpoint for inference or release is exported separately and follows the model’s intended dtype (for fine-tuning, this matches the original HF checkpoint, typically bf16). See Checkpointing for details.

Precision and Robustness

Both patterns keep the master weights and Adam EMA in fp32, so for most models they converge equivalently. The remaining difference is the gradient: under TE (Pattern A), the resident parameter is bf16, so the gradient feeding the Adam update carries bf16 rounding, while torch AdamW with fp32 storage (Pattern B) keeps the parameter, gradient, and optimizer state all in fp32.

This is usually negligible. When bringing up a new model, fp32 storage (Pattern B) is a robust starting point because every part of the update path is fp32; move to TE (Pattern A) after it is validated for that model and you want bf16 checkpoints.

Risky Pattern Using torch AdamW with bf16 Model Storage

This pattern is easy to enter accidentally. In several NeMo AutoModel paths, leaving model.dtype unset, or setting it to auto, resolves the resident model parameter dtype to bf16. If that is paired with torch.optim.AdamW, the AdamW EMA buffers are also bf16 because the optimizer initializes state from the parameter dtype.

model:
  # dtype omitted, or dtype: auto / bfloat16
optimizer:
  _target_: torch.optim.AdamW
distributed:
  strategy: fsdp2

This keeps the resident model parameters in bf16 instead of fp32, so it can reduce memory usage compared to the torch AdamW and model.dtype: float32 pattern. It is common in existing fine-tuning example configs, but convergence should be validated for the intended workload. It is not recommended for long full-parameter training (pretraining or extended fine-tuning): bf16 EMA quantization can quietly slow convergence (higher final loss at the same step count) and, in worse cases, produce unstable grad_norm, loss bumps, loss spikes, or divergence. For LoRA/PEFT, check adapter storage separately as described below.

LoRA Adapter Storage and Precision Warnings

Frozen base parameters do not receive optimizer updates or acquire Adam state. Keeping them in bf16 therefore does not reduce the precision of adapter updates. With torch.optim.Adam or AdamW, it is the trainable adapter storage dtype that determines the precision of adapter updates and EMA buffers.

For multi-rank, unquantized FSDP2 LoRA training, set the adapter dtype explicitly to keep adapter storage and Adam moments in fp32 while allowing frozen base storage and forward and backward compute to remain bf16:

peft:
  lora_dtype: float32

The default peft.lora_dtype: null inherits the base layer dtype when adapters are created; it does not guarantee fp32 adapters. Setting peft.lora_dtype: bfloat16 explicitly gives bf16 adapter storage and, with torch Adam/AdamW, bf16 moments. Validate convergence before choosing that setting. In this multi-rank FSDP2 path, changing model.dtype alone does not override an explicit adapter dtype.

On a single-rank FSDP2 setup, parameters remain unsharded and mixed LoRA submodules are cast to the compute dtype, so peft.lora_dtype: float32 does not guarantee fp32 adapter storage or optimizer moments there.

The automatic warning about torch Adam/AdamW with trainable bf16 parameters currently skips PEFT runs. Absence of this warning does not mean your adapters or their optimizer state are fp32. Check peft.lora_dtype and the actual trainable parameter dtypes; use peft.lora_dtype: float32 to request fp32 adapter storage with torch Adam/AdamW on multi-rank FSDP2 runs.

Example Configs

examples/llm_pretrain/llama3_70b_pretrain.yaml: TE FusedAdam example. This keeps model storage and training checkpoints in bf16 while using fp32 master weights and optimizer state.
examples/llm_pretrain/megatron_pretrain_moonlight_16b_te_slurm.yaml: torch AdamW example. This uses model.torch_dtype: float32 so AdamW state stays fp32 while compute remains bf16 through FSDP2 mixed precision.