> 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)                         | -                                                                                  |
| [`_combine_norms`](#nemo_automodel-components-training-utils-_combine_norms)                                     | -                                                                                  |
| [`clip_grad_norm`](#nemo_automodel-components-training-utils-clip_grad_norm)                                     | Common gradient clipping helper.                                                   |
| [`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 in a single pass, then clip.                             |

### Data

[`_TE_EXPERT_PARAM_PATTERN`](#nemo_automodel-components-training-utils-_TE_EXPERT_PARAM_PATTERN)

### 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
) -> torch.Tensor
```

```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.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
) -> torch.Tensor | float
```

Common gradient clipping helper.

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 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.

---

**`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.

---

**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
) -> torch.Tensor | float
```

Scale gradients for PP/EP in a single pass, then clip.

* 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`.
* Finally, perform grad clipping with PP/EP-aware reductions.

**Returns:** `torch.Tensor | float`

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

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