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

Synchronization helpers for parameters replicated across tensor parallel ranks.

## Module Contents

### Functions

| Name                                                                                                                          | Description                                                                  |
| ----------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------- |
| [`_get_tp_mesh`](#nemo_automodel-components-distributed-tp_replicas-_get_tp_mesh)                                             | Return a non-trivial TP submesh when one is present.                         |
| [`_gradient_chunks`](#nemo_automodel-components-distributed-tp_replicas-_gradient_chunks)                                     | Split gradients into bounded communication buffers.                          |
| [`_gradient_reduce_dtype`](#nemo_automodel-components-distributed-tp_replicas-_gradient_reduce_dtype)                         | Return an FP32 communication dtype for low-precision gradients.              |
| [`_is_tp_replicated`](#nemo_automodel-components-distributed-tp_replicas-_is_tp_replicated)                                   | Return whether rank-local tensor storage is replicated across TP peers.      |
| [`_iter_unique_buffers`](#nemo_automodel-components-distributed-tp_replicas-_iter_unique_buffers)                             | Yield model buffers once in rank-stable fully qualified name order.          |
| [`_iter_unique_parameters`](#nemo_automodel-components-distributed-tp_replicas-_iter_unique_parameters)                       | Yield unmarked parameters once with their module reduction semantic.         |
| [`_iter_unique_parameters_by_name`](#nemo_automodel-components-distributed-tp_replicas-_iter_unique_parameters_by_name)       | Yield parameters once in rank-stable fully qualified name order.             |
| [`_local_tensor`](#nemo_automodel-components-distributed-tp_replicas-_local_tensor)                                           | Return rank-local storage for a Tensor or DTensor.                           |
| [`broadcast_tp_replicas`](#nemo_automodel-components-distributed-tp_replicas-broadcast_tp_replicas)                           | Broadcast replicated parameters and buffers from the first TP rank.          |
| [`exclude_from_tp_replica_sync`](#nemo_automodel-components-distributed-tp_replicas-exclude_from_tp_replica_sync)             | Exclude an owner-sharded module subtree from TP replica synchronization.     |
| [`mark_tp_replica_gradient_reduction`](#nemo_automodel-components-distributed-tp_replicas-mark_tp_replica_gradient_reduction) | Mark direct parameters of a module with their TP-replica reduction semantic. |
| [`synchronize_tp_replica_gradients`](#nemo_automodel-components-distributed-tp_replicas-synchronize_tp_replica_gradients)     | Reduce TP-replicated gradients once at the optimizer boundary.               |

### Data

[`_MAX_FLAT_BUFFER_BYTES`](#nemo_automodel-components-distributed-tp_replicas-_MAX_FLAT_BUFFER_BYTES)

[`_MODEL_OWNED_GRAD_DIVISOR_ATTR`](#nemo_automodel-components-distributed-tp_replicas-_MODEL_OWNED_GRAD_DIVISOR_ATTR)

[`_TP_REPLICA_GRAD_REDUCTION_ATTR`](#nemo_automodel-components-distributed-tp_replicas-_TP_REPLICA_GRAD_REDUCTION_ATTR)

### API

```python
nemo_automodel.components.distributed.tp_replicas._get_tp_mesh(
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None,
    tp_axis_name: str
) -> torch.distributed.device_mesh.DeviceMesh | None
```

Return a non-trivial TP submesh when one is present.

```python
nemo_automodel.components.distributed.tp_replicas._gradient_chunks(
    gradients: list[torch.Tensor]
) -> collections.abc.Iterator[list[torch.Tensor]]
```

Split gradients into bounded communication buffers.

**Parameters:**

**`gradients`** `list[torch.Tensor]`

Dense rank-local gradients of arbitrary shape on one device
with one dtype. Low-precision elements are budgeted at their FP32
communication size.

---

```python
nemo_automodel.components.distributed.tp_replicas._gradient_reduce_dtype(
    dtype: torch.dtype
) -> torch.dtype
```

Return an FP32 communication dtype for low-precision gradients.

```python
nemo_automodel.components.distributed.tp_replicas._is_tp_replicated(
    tensor: torch.Tensor,
    tp_group_ranks: tuple[int, ...],
    current_rank: int,
    tp_axis_name: str
) -> bool
```

Return whether rank-local tensor storage is replicated across TP peers.

A DTensor may live on a DP-only submesh even though physical copies exist for
every TP coordinate. Conversely, an MoE mesh may fold the TP ranks into axes
named `ep` or `ep_shard`. Comparing mesh coordinates and placements keeps
both cases explicit instead of assuming every unnamed TP dimension is a copy.

**Parameters:**

**`tensor`** `torch.Tensor`

Tensor or DTensor of arbitrary global and rank-local shape.

---

**`tp_group_ranks`** `tuple[int, ...]`

Ordered global ranks in the current TP group.

---

**`current_rank`** `int`

Current global rank.

---

**`tp_axis_name`** `str`

Name of the tensor-parallel mesh dimension.

---

**Returns:** `bool`

Whether every TP peer represented in the tensor mesh differs only along

```python
nemo_automodel.components.distributed.tp_replicas._iter_unique_buffers(
    model_parts: list[torch.nn.Module]
) -> collections.abc.Iterator[torch.Tensor]
```

Yield model buffers once in rank-stable fully qualified name order.

```python
nemo_automodel.components.distributed.tp_replicas._iter_unique_parameters(
    model_parts: list[torch.nn.Module]
) -> collections.abc.Iterator[tuple[torch.nn.Parameter, typing.Literal['mean', 'sum']]]
```

Yield unmarked parameters once with their module reduction semantic.

```python
nemo_automodel.components.distributed.tp_replicas._iter_unique_parameters_by_name(
    model_parts: list[torch.nn.Module]
) -> collections.abc.Iterator[tuple[torch.nn.Parameter, typing.Literal['mean', 'sum']]]
```

Yield parameters once in rank-stable fully qualified name order.

```python
nemo_automodel.components.distributed.tp_replicas._local_tensor(
    tensor: torch.Tensor
) -> torch.Tensor
```

Return rank-local storage for a Tensor or DTensor.

```python
nemo_automodel.components.distributed.tp_replicas.broadcast_tp_replicas(
    model_parts: list[torch.nn.Module],
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None,
    tp_axis_name: str = 'tp'
) -> int
```

Broadcast replicated parameters and buffers from the first TP rank.

Intended TP shards are excluded by their DTensor placement, and model-owned
shards are excluded by their explicit parameter or module marker. The
collective operates on each tensor's rank-local storage, whose shape is
identical within a TP group even when another mesh axis shards the global
tensor.

**Parameters:**

**`model_parts`** `list[torch.nn.Module]`

Local pipeline-stage modules containing tensors of arbitrary
shape. Shared tensors are visited once.

---

**`device_mesh`** `DeviceMesh | None`

Root mesh containing the tensor-parallel axis.

---

**`tp_axis_name`** `str` — default: 'tp'

Name of the tensor-parallel mesh dimension.

---

**Returns:** `int`

Number of rank-local tensors broadcast on this rank.

```python
nemo_automodel.components.distributed.tp_replicas.exclude_from_tp_replica_sync(
    module: torch.nn.Module
) -> None
```

Exclude an owner-sharded module subtree from TP replica synchronization.

Some parameter owners use a mesh that folds physical TP ranks into another
logical axis. Their local tensors are different shards, even though they are
not DTensors carrying an explicit `tp` placement. Marking the complete
subtree keeps both parameter and buffer synchronization from treating those
owner-local tensors as TP replicas.

**Parameters:**

**`module`** `torch.nn.Module`

Owner-sharded module whose current subtree must remain local.

---

```python
nemo_automodel.components.distributed.tp_replicas.mark_tp_replica_gradient_reduction(
    module: torch.nn.Module,
    reduction: typing.Literal['mean', 'sum']
) -> None
```

Mark direct parameters of a module with their TP-replica reduction semantic.

**Parameters:**

**`module`** `torch.nn.Module`

Module whose direct parameters are replicated across TP ranks.

---

**`reduction`** `Literal['mean', 'sum']`

`"mean"` for redundant full computation or `"sum"` for
disjoint partial contributions.

---

```python
nemo_automodel.components.distributed.tp_replicas.synchronize_tp_replica_gradients(
    model_parts: list[torch.nn.Module],
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None,
    tp_axis_name: str = 'tp'
) -> int
```

Reduce TP-replicated gradients once at the optimizer boundary.

Call this exactly once per optimizer update, after gradient accumulation and
before gradient scaling, norm calculation, clipping, or `optimizer.step()`.
Full-computation replicas are mean-reduced so each logical parameter has one
optimizer update. Modules explicitly marked as producing disjoint partial
contributions are sum-reduced, which makes repeated synchronization within
one update non-idempotent. Gradients are flattened into bounded buffers by
reduction, device, and dtype before communication. DTensor gradients are
reduced through rank-local storage while retaining their global shape and
placements. Model-owned shards and parameters or gradients sharded or partial
on the TP axis are untouched. Low-precision gradients are reduced in FP32 and
cast back to their storage dtype after synchronization.

**Parameters:**

**`model_parts`** `list[torch.nn.Module]`

Local pipeline-stage modules whose gradients have arbitrary
shapes and have completed accumulation for this optimizer update.

---

**`device_mesh`** `DeviceMesh | None`

Root mesh containing the tensor-parallel axis.

---

**`tp_axis_name`** `str` — default: 'tp'

Name of the tensor-parallel mesh dimension.

---

**Returns:** `int`

Number of rank-local gradients synchronized on this rank.

**Raises:**

* `RuntimeError`: If gradient presence differs across TP replicas, or if a
  replicated parameter has a sparse gradient that cannot be flattened
  without changing its representation.

```python
nemo_automodel.components.distributed.tp_replicas._MAX_FLAT_BUFFER_BYTES = 256 * 1024 * 1024
```

```python
nemo_automodel.components.distributed.tp_replicas._MODEL_OWNED_GRAD_DIVISOR_ATTR = '_nemo_model_owned_grad_divisor'
```

```python
nemo_automodel.components.distributed.tp_replicas._TP_REPLICA_GRAD_REDUCTION_ATTR = '_nemo_tp_replica_grad_reduction'
```