nemo_automodel.components.training.utils

Module Contents

Classes

Name	Description
`ScopedModuleOffloading`	Context manager that temporarily moves a module between CPU and CUDA.

Functions

Name	Description
`_all_reduce_scalar`	All-reduce a 0-dim norm accumulator over `mesh`, communicating on the mesh device.
`_clip_grad_norm_impl`	Compute and clip the norm of local and DTensor gradients.
`_combine_norms`	-
`_local_te_l2_norm`	Reduce local gradients with Transformer Engine where supported.
`_use_fused_grad_norm`	Whether the fused multi-tensor reduction applies to this group.
`clip_grad_norm`	Apply sharding-aware gradient clipping.
`count_tail_padding`	Counts the total number of padding token in the tail of labels
`get_expert_tp_replication_factor`	Return the TP token-replication factor for custom-MoE expert gradients.
`move_to_device`	Move a model and its buffers to a device and release stale CUDA cache.
`prepare_after_first_microbatch`	Disable first-microbatch flag after the first forward-backward pass.
`prepare_for_final_backward`	Prepare model parts before the final backward pass.
`prepare_for_grad_accumulation`	Prepare model parts before starting gradient accumulation.
`scale_grads_and_clip_grad_norm`	Scale gradients for PP/EP and model-owned shards, then clip.

Data

_GradNormBackend

_TE_EXPERT_PARAM_PATTERN

API

class nemo_automodel.components.training.utils.ScopedModuleOffloading(
    model,
    enabled = False
)

Context manager that temporarily moves a module between CPU and CUDA.

nemo_automodel.components.training.utils.ScopedModuleOffloading.__enter__()

nemo_automodel.components.training.utils.ScopedModuleOffloading.__exit__(
    exc_type,
    exc_val,
    exc_tb
)

nemo_automodel.components.training.utils._all_reduce_scalar(
    scalar: torch.Tensor,
    op: torch.distributed.ReduceOp,
    mesh: torch.distributed.device_mesh.DeviceMesh,
    mesh_dim: int | None = None
) -> torch.Tensor

All-reduce a 0-dim norm accumulator over mesh, communicating on the mesh device.

The norm math stays on the gradients’ own device, which under FSDP2 CPUOffloadPolicy is CPU while the mesh’s process group is NCCL and has no CPU backend. Only the scalar hops to mesh.device_type for the collective and comes straight back, so a genuinely-CPU (gloo) mesh is never forced onto an accelerator.

Parameters:

scalar

torch.Tensor

0-dim tensor to reduce, on the gradients’ device.

torch.distributed.ReduceOp

Reduction operation.

mesh

DeviceMesh

Device mesh whose process group performs the collective.

mesh_dim

int | NoneDefaults to None

Mesh dimension to reduce over, or None for the whole mesh.

Returns: torch.Tensor

The reduced scalar on scalar’s original device. Callers must use the return

nemo_automodel.components.training.utils._clip_grad_norm_impl(
    parameters: torch.Tensor | typing.Iterable[torch.Tensor],
    max_norm: float,
    norm_type: float = 2.0,
    error_if_nonfinite: bool = False,
    foreach: bool | None = None,
    pp_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    grad_norm_backend: nemo_automodel.components.training.utils._GradNormBackend = 'triton'
) -> torch.Tensor

Compute and clip the norm of local and DTensor gradients.

Parameters:

parameters

torch.Tensor | Iterable[torch.Tensor]

One parameter tensor or an iterable of parameter tensors with arbitrary shapes. DTensors retain their declared mesh and placements.

max_norm

float

Maximum allowed global gradient norm.

norm_type

floatDefaults to 2.0

Norm exponent, including inf.

error_if_nonfinite

boolDefaults to False

Whether to raise for a non-finite global norm.

foreach

bool | NoneDefaults to None

Optional foreach implementation preference for clipping.

pp_mesh

DeviceMesh | NoneDefaults to None

Optional pipeline mesh over which the scalar norm is reduced.

grad_norm_backend

_GradNormBackendDefaults to 'triton'

Local L2 reducer. "triton" uses this PR’s FP64 multi-tensor kernel; "te" uses Transformer Engine where eligible.

Returns: torch.Tensor

Scalar tensor containing the pre-clipping global gradient norm.

nemo_automodel.components.training.utils._combine_norms(
    norms: list[torch.Tensor],
    norm_type: float,
    target_device: torch.device
) -> torch.Tensor

nemo_automodel.components.training.utils._local_te_l2_norm(
    gradients: list[torch.Tensor],
    target_device: torch.device
) -> torch.Tensor

Reduce local gradients with Transformer Engine where supported.

Parameters:

gradients

list[torch.Tensor]

Plain local tensors of arbitrary shape, without DTensor placements. Contiguous CUDA FP16/BF16/FP32 tensors use TE, grouped by device and dtype. Other layouts, devices, and dtypes use PyTorch’s FP64 vector norm. Inputs are read-only and may alias parameter gradients.

target_device

torch.device

Device for the returned scalar; only scalars move between devices.

Returns: torch.Tensor

Independent scalar FP64 L2 norm on target_device. TE squares and accumulates

Raises:

RuntimeError: If eligible CUDA gradients require TE but TE is unavailable.

nemo_automodel.components.training.utils._use_fused_grad_norm(
    params,
    norm_type: float
) -> bool

Whether the fused multi-tensor reduction applies to this group.

Only the 2-norm and inf-norm are implemented by the kernel, and the whole group has to be CUDA — a mixed CPU/CUDA group would silently take two different reduction paths.

nemo_automodel.components.training.utils.clip_grad_norm(
    max_grad_norm: float | None,
    model_parts: list[torch.nn.Module],
    norm_type: float = 2.0,
    pp_enabled: bool = False,
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    pp_axis_name: str | None = None,
    foreach: bool = True,
    use_torch_clip_grad_norm: bool = False,
    grad_norm_backend: nemo_automodel.components.training.utils._GradNormBackend = 'triton'
) -> torch.Tensor | float

Apply sharding-aware gradient clipping.

Handles all parallelism strategies (TP, PP, EP/MoE) with automatic sharding-aware grouping. Returns the gradient norm as a scalar tensor on the gradients’ device, or 0.0 if clipping is skipped. This function does not synchronize TP-replicated gradients; optimizer loops must do that exactly once before calling this function.

This function automatically:

Groups parameters by sharding pattern (device mesh + placements)
Computes norms correctly across different sharding strategies
Handles MoE with separate DP/EP meshes
Reduces norms across pipeline parallel stages when enabled

Parameters:

max_grad_norm

float | None

Maximum gradient norm. If None, skips clipping.

model_parts

list[torch.nn.Module]

List of model modules to clip.

norm_type

floatDefaults to 2.0

Type of norm to use (default: 2.0 for L2).

pp_enabled

boolDefaults to False

Whether pipeline parallelism is enabled.

device_mesh

DeviceMesh | NoneDefaults to None

Device mesh for parallelism.

pp_axis_name

str | NoneDefaults to None

Pipeline parallel axis name.

foreach

boolDefaults to True

Whether to use foreach implementation for clipping.

use_torch_clip_grad_norm

boolDefaults to False

Use PyTorch’s optimized regular-tensor clipping path when possible.

grad_norm_backend

_GradNormBackendDefaults to 'triton'

Local L2 reducer, either "triton" or "te".

Returns: torch.Tensor | float

Scalar tensor containing the total gradient norm without synchronizing it to the host,

nemo_automodel.components.training.utils.count_tail_padding(
    labels,
    ignore_label = -100
)

Counts the total number of padding token in the tail of labels

e.g. labels = torch.tensor([ [-100, 1, 1, -100, -100], # 2 tail -100s [-100, -100, 2, 3, 4], # 0 tail -100s [5, 6, -100, -100, -100], # 3 tail -100s ]) count_tail_padding will return 5. Please do note there’s more than 5 ignore labels. Args: labels (torch.Tensor): the labels ignore_label (int, optional): ignore label index. Defaults to -100.

Returns:

total number of ignored tokens in the labels input.

nemo_automodel.components.training.utils.get_expert_tp_replication_factor(
    model_parts: list[torch.nn.Module],
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None
) -> int

Return the TP token-replication factor for custom-MoE expert gradients.

The custom-MoE tensor-parallel path keeps the token path (attention, router) replicated across TP ranks, so every TP rank feeds the same tokens into the expert-parallel all-gather and each expert gradient is accumulated tp_size times. scale_grads_and_clip_grad_norm divides expert gradients by this factor to restore the correct scale.

nemo_automodel.components.training.utils.move_to_device(
    model,
    device
)

Move a model and its buffers to a device and release stale CUDA cache.

nemo_automodel.components.training.utils.prepare_after_first_microbatch()

Disable first-microbatch flag after the first forward-backward pass.

Called after the first microbatch in gradient accumulation so that subsequent microbatches reuse cached FP8 weights instead of re-quantizing.

nemo_automodel.components.training.utils.prepare_for_final_backward(
    model_parts: list[torch.nn.Module],
    pp_enabled: bool = False
)

Prepare model parts before the final backward pass.

This is typically called before the final gradient accumulation step to prepare FSDP states for gradient synchronization and resharding.

Parameters:

model_parts

list[torch.nn.Module]

List of model parts (modules) to prepare.

pp_enabled

boolDefaults to False

Whether pipeline parallelism is enabled.

nemo_automodel.components.training.utils.prepare_for_grad_accumulation(
    model_parts: list[torch.nn.Module],
    pp_enabled: bool = False
)

Prepare model parts before starting gradient accumulation.

This is typically called once at the start of gradient accumulation to prepare FSDP states for the upcoming forward and backward passes.

Parameters:

model_parts

list[torch.nn.Module]

List of model parts (modules) to prepare.

pp_enabled

boolDefaults to False

Whether pipeline parallelism is enabled.

nemo_automodel.components.training.utils.scale_grads_and_clip_grad_norm(
    max_grad_norm: float | None,
    model_parts: list[torch.nn.Module],
    norm_type: float = 2.0,
    pp_enabled: bool = False,
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    moe_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    ep_axis_name: str | None = None,
    pp_axis_name: str | None = None,
    foreach: bool = True,
    num_label_tokens: int | None = None,
    dp_group_size: int | None = None,
    expert_tp_replication_factor: int = 1,
    use_torch_clip_grad_norm: bool = False,
    grad_norm_backend: nemo_automodel.components.training.utils._GradNormBackend = 'triton'
) -> torch.Tensor | float

Scale gradients for PP/EP and model-owned shards, then clip.

The caller must synchronize TP-replicated gradients once after accumulation and before calling this function. This helper does not synchronize replicas.

PP scaling: divide all local grads by (num_label_tokens / dp_group_size).
EP scaling: for parameters on the expert axis, divide grads by (dp_group_size / ep_shard_size) * expert_tp_replication_factor.
Owner-sharded scaling: divide each marked gradient by the explicit factor declared by its model-owned sharding contract.
Finally, perform grad clipping with PP/EP-aware reductions.

Parameters:

max_grad_norm

float | None

Maximum global gradient norm, or None to skip clipping.

model_parts

list[torch.nn.Module]

Model modules whose parameters have gradients of arbitrary shape. Gradients retain their original local or DTensor layout and are scaled in place.

norm_type

floatDefaults to 2.0

Norm order.

pp_enabled

boolDefaults to False

Whether pipeline-parallel normalization is required.

device_mesh

DeviceMesh | NoneDefaults to None

Training mesh used for gradient norm reductions.

moe_mesh

DeviceMesh | NoneDefaults to None

Expert-parallel mesh used to normalize expert gradients.

ep_axis_name

str | NoneDefaults to None

Expert axis in the parameter mesh.

pp_axis_name

str | NoneDefaults to None

Pipeline axis in the training mesh.

foreach

boolDefaults to True

Whether to use foreach for in-place clipping.

num_label_tokens

int | NoneDefaults to None

Global supervised-token count for PP normalization.

dp_group_size

int | NoneDefaults to None

Data-parallel group size, including CP when configured.

expert_tp_replication_factor

intDefaults to 1

Number of identical TP copies of expert tokens.

use_torch_clip_grad_norm

boolDefaults to False

Prefer PyTorch’s regular-tensor clipping fast path.

grad_norm_backend

_GradNormBackendDefaults to 'triton'

Local L2 reducer, either "triton" or "te".

Returns: torch.Tensor | float

Scalar tensor containing the total gradient norm without synchronizing it to the host,

nemo_automodel.components.training.utils._GradNormBackend = Literal['triton', 'te']

nemo_automodel.components.training.utils._TE_EXPERT_PARAM_PATTERN = re.compile('(^|\\.)mlp\\.experts\\.(gate_up_linear|down_linear)\\.(weight|bias)\...