> This page is for version 0.5.0 · 26.06. > For other versions, use one of these documentation indexes: > - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt > - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt > - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt > - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt > For clean Markdown of any page, append .md to the page URL. > For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt. > For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server. # nemo_automodel.components.training.utils ## Module Contents ### Classes | Name | Description | | -------------------------------------------------------------------------------------------- | --------------------------------------------------------------------- | | [`ScopedModuleOffloading`](#nemo_automodel-components-training-utils-ScopedModuleOffloading) | Context manager that temporarily moves a module between CPU and CUDA. | ### Functions | Name | Description | | ---------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------- | | [`_all_reduce_scalar`](#nemo_automodel-components-training-utils-_all_reduce_scalar) | All-reduce a 0-dim norm accumulator over `mesh`, communicating on the mesh device. | | [`_clip_grad_norm_impl`](#nemo_automodel-components-training-utils-_clip_grad_norm_impl) | Compute and clip the norm of local and DTensor gradients. | | [`_combine_norms`](#nemo_automodel-components-training-utils-_combine_norms) | - | | [`_ep_local_expert_param_ids`](#nemo_automodel-components-training-utils-_ep_local_expert_param_ids) | ids of expert params whose grads live on the EP axis without carrying it. | | [`_local_te_l2_norm`](#nemo_automodel-components-training-utils-_local_te_l2_norm) | Reduce local gradients with Transformer Engine where supported. | | [`_use_fused_grad_norm`](#nemo_automodel-components-training-utils-_use_fused_grad_norm) | Whether the fused multi-tensor reduction applies to this group. | | [`clip_grad_norm`](#nemo_automodel-components-training-utils-clip_grad_norm) | Apply sharding-aware gradient clipping. | | [`count_tail_padding`](#nemo_automodel-components-training-utils-count_tail_padding) | Counts the total number of padding token in the tail of labels | | [`get_expert_tp_replication_factor`](#nemo_automodel-components-training-utils-get_expert_tp_replication_factor) | Return the TP token-replication factor for custom-MoE expert gradients. | | [`move_to_device`](#nemo_automodel-components-training-utils-move_to_device) | Move a model and its buffers to a device and release stale CUDA cache. | | [`prepare_after_first_microbatch`](#nemo_automodel-components-training-utils-prepare_after_first_microbatch) | Disable first-microbatch flag after the first forward-backward pass. | | [`prepare_for_final_backward`](#nemo_automodel-components-training-utils-prepare_for_final_backward) | Prepare model parts before the final backward pass. | | [`prepare_for_grad_accumulation`](#nemo_automodel-components-training-utils-prepare_for_grad_accumulation) | Prepare model parts before starting gradient accumulation. | | [`scale_grads_and_clip_grad_norm`](#nemo_automodel-components-training-utils-scale_grads_and_clip_grad_norm) | Scale gradients for PP/EP and model-owned shards, then clip. | ### Data [`_GradNormBackend`](#nemo_automodel-components-training-utils-_GradNormBackend) ### API ```python class nemo_automodel.components.training.utils.ScopedModuleOffloading( model, enabled = False ) ``` Context manager that temporarily moves a module between CPU and CUDA. ```python nemo_automodel.components.training.utils.ScopedModuleOffloading.__enter__() ``` ```python nemo_automodel.components.training.utils.ScopedModuleOffloading.__exit__( exc_type, exc_val, exc_tb ) ``` ```python 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. --- **`op`** `torch.distributed.ReduceOp` Reduction operation. --- **`mesh`** `DeviceMesh` Device mesh whose process group performs the collective. --- **`mesh_dim`** `int | None` — default: 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 ```python 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`** `float` — default: 2.0 Norm exponent, including `inf`. --- **`error_if_nonfinite`** `bool` — default: False Whether to raise for a non-finite global norm. --- **`foreach`** `bool | None` — default: None Optional foreach implementation preference for clipping. --- **`pp_mesh`** `DeviceMesh | None` — default: None Optional pipeline mesh over which the scalar norm is reduced. --- **`ep_mesh`** `DeviceMesh | None` — default: 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] | None` — default: None ids of the parameters whose gradients differ across the EP axis and need the `ep_mesh` reduction. --- **`grad_norm_backend`** `_GradNormBackend` — default: '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. ```python nemo_automodel.components.training.utils._combine_norms( norms: list[torch.Tensor], norm_type: float, target_device: torch.device ) -> torch.Tensor ``` ```python 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. ```python 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. ```python 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`. --- ```python 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`** `float` — default: 2.0 Type of norm to use (default: 2.0 for L2). --- **`pp_enabled`** `bool` — default: False Whether pipeline parallelism is enabled. --- **`device_mesh`** `DeviceMesh | None` — default: None Device mesh for parallelism. --- **`pp_axis_name`** `str | None` — default: None Pipeline parallel axis name. --- **`moe_mesh`** `DeviceMesh | None` — default: 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 | None` — default: None Name of the expert-parallel axis in `moe_mesh`. --- **`ep_param_ids`** `set[int] | None` — default: 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`** `bool` — default: True Whether to use foreach implementation for clipping. --- **`use_torch_clip_grad_norm`** `bool` — default: False Use PyTorch's optimized regular-tensor clipping path when possible. --- **`grad_norm_backend`** `_GradNormBackend` — default: '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, ```python 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. ```python 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. ```python nemo_automodel.components.training.utils.move_to_device( model, device ) ``` Move a model and its buffers to a device and release stale CUDA cache. ```python 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. ```python 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`** `bool` — default: False Whether pipeline parallelism is enabled. --- ```python 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`** `bool` — default: False Whether pipeline parallelism is enabled. --- ```python 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`** `float` — default: 2.0 Norm order. --- **`pp_enabled`** `bool` — default: False Whether pipeline-parallel normalization is required. --- **`device_mesh`** `DeviceMesh | None` — default: None Training mesh used for gradient norm reductions. --- **`moe_mesh`** `DeviceMesh | None` — default: None Expert-parallel mesh used to normalize expert gradients. --- **`ep_axis_name`** `str | None` — default: None Expert axis in the parameter mesh. --- **`pp_axis_name`** `str | None` — default: None Pipeline axis in the training mesh. --- **`foreach`** `bool` — default: True Whether to use foreach for in-place clipping. --- **`num_label_tokens`** `int | None` — default: None Global supervised-token count for PP normalization. --- **`dp_group_size`** `int | None` — default: None Data-parallel group size, including CP when configured. --- **`expert_tp_replication_factor`** `int` — default: 1 Number of identical TP copies of expert tokens. --- **`use_torch_clip_grad_norm`** `bool` — default: False Prefer PyTorch's regular-tensor clipping fast path. --- **`grad_norm_backend`** `_GradNormBackend` — default: '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, ```python nemo_automodel.components.training.utils._GradNormBackend = Literal['triton', 'te'] ```