nemo_automodel.components.distributed.tp_replicas

Synchronization helpers for parameters replicated across tensor parallel ranks.

Module Contents

Functions

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

Data

_MAX_FLAT_BUFFER_BYTES

_MODEL_OWNED_GRAD_DIVISOR_ATTR

_TP_REPLICA_GRAD_REDUCTION_ATTR

API

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.

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.

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

Return an FP32 communication dtype for low-precision gradients.

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

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.

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.

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.

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

Return rank-local storage for a Tensor or DTensor.

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

strDefaults to 'tp'

Name of the tensor-parallel mesh dimension.

Returns: int

Number of rank-local tensors broadcast on this rank.

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.

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.

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

strDefaults to '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.

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

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

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