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`	-
`_ep_local_expert_param_ids`	ids of expert params whose grads live on the EP axis without carrying it.
`_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

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,
    ep_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    ep_param_ids: set[int] | 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.

ep_mesh

DeviceMesh | NoneDefaults to None

Optional expert-parallel mesh. Expert parameters live on the EP axis without carrying it in their own mesh (TE grouped experts are plain tensors, or DTensors on the ep_shard sub-mesh only), so their norm contribution must additionally be reduced over this mesh or every EP group computes a different “global” norm and clips shards of the same logical parameter inconsistently.

ep_param_ids

set[int] | NoneDefaults to None

ids of the parameters whose gradients differ across the EP axis and need the ep_mesh reduction.

grad_norm_backend

_GradNormBackendDefaults to 'triton'

Local L2 reducer. "triton" uses the 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._ep_local_expert_param_ids(
    model_parts: list[torch.nn.Module],
    ep_axis_name: str | None
) -> set[int]

ids of expert params whose grads live on the EP axis without carrying it.

Modules stamp _nemo_ep_local_expert_params at construction (TE grouped experts, see GroupedExpertsTE.__init__) because their parameters are either plain tensors or DTensors sharded only over ep_shard — so no mesh dimension of theirs reflects that gradients differ per EP rank. EP grad scaling and the grad-clip EP norm reduction both use this set, and it depends only on model structure, so every rank derives the same answer (collective participation must be rank-uniform).

Params whose own mesh already carries the EP axis are excluded: the per-mesh-dim reductions handle those, and reducing them again over EP would double count.

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(
    representative: torch.Tensor,
    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 kernel itself is CUDA-only (_reduce routes anything else through its reference reduction, so a non-CUDA group would pay the wrapper for nothing).

The choice is made from the group’s representative parameter rather than from this rank’s gradients. The fused branch issues one collective where the fallback issues two, so a group that resolved differently on two ranks would mismatch their collectives and hang. Every rank holds the representative for every group, including the ranks that contributed no gradients this step, so deriving the choice from it keeps the decision identical everywhere.

Parameters:

representative

torch.Tensor

The group’s representative parameter, equal in device and layout across ranks by construction.

norm_type

float

Norm exponent, including inf.

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,
    moe_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    ep_axis_name: str | None = None,
    ep_param_ids: set[int] | 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.

moe_mesh

DeviceMesh | NoneDefaults to None

Optional MoE mesh. TE grouped expert parameters don’t carry the EP axis on their own mesh (they are plain tensors, or DTensors sharded only over ep_shard), so without this their norm contribution is never reduced across EP and every EP group clips with a different “global” norm.

ep_axis_name

str | NoneDefaults to None

Name of the expert-parallel axis in moe_mesh.

ep_param_ids

set[int] | NoneDefaults to None

Optional precomputed result of _ep_local_expert_param_ids (callers that already scanned the modules pass it to avoid a second walk); derived when omitted.

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']