nemo_automodel.components.distributed.context_parallel.utils

Module Contents

Functions

Name	Description
`_attention_backend`	Read the configured attention backend from a live model.
`_get_submesh`	The named submesh, or None when the mesh or dimension is absent.
`_is_multimodal_model`	Whether a live model owns a vision or audio tower.
`_magi_state_from_model`	Recreate the per-forward Magi handle from the model and device mesh.
`_make_cp_batch_and_ctx`	Resolve a ContextParallelSharder and shard the batch; a no-op when no CP prep applies.
`_mesh_dim_size`	Size of a named submesh, 0 when the mesh or dimension is absent.
`_prepare_cp_sharder`	Resolve and configure a CP sharder for its public constructor.
`_resolve_cp_sharder`	Resolve the ContextParallelSharder for this forward: model-owned > magi > TE > generic > none.
`_shard_grad_buffer_for_cp`	Shard a gradient-bearing buffer with CP’s head-tail load-balancing order.
`_shard_thd_chunk_for_te`	-
`_uses_te_attention`	Whether a live model uses native or injected TE attention.
`attach_context_parallel_hooks`	Attach forward pre-hooks to self_attn modules to fix attention masks for context parallelism.
`attach_cp_sdpa_hooks`	Inject CP-aware SDPA into self_attn modules for compile + CP>1 correctness.
`attach_te_context_parallel`	Configure Transformer Engine attention modules for context and tensor parallelism.
`cp_dispatcher_suspended`	Temporarily suspend torch’s `context_parallel` ring-SDPA monkeypatch.
`create_context_parallel_ctx`	Create a context parallel context.
`get_train_context`	Create a train context.
`make_cp_batch_for_te`	Build a CP batch for Transformer Engine using THD format.
`unshard_context_parallel_tensor`	Restore a tensor from PyTorch’s load-balanced context-parallel layout.

API

nemo_automodel.components.distributed.context_parallel.utils._attention_backend(
    model
) -> str | None

Read the configured attention backend from a live model.

nemo_automodel.components.distributed.context_parallel.utils._get_submesh(
    device_mesh,
    dim: str
)

The named submesh, or None when the mesh or dimension is absent.

nemo_automodel.components.distributed.context_parallel.utils._is_multimodal_model(
    model
) -> bool

Whether a live model owns a vision or audio tower.

nemo_automodel.components.distributed.context_parallel.utils._magi_state_from_model(
    model,
    device_mesh
)

Recreate the per-forward Magi handle from the model and device mesh.

nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx(
    device_mesh,
    batch,
    loss_mask = None,
    use_te: bool = False,
    padding_token_id: int = 0,
    num_chunks: int = 1,
    seq_lens_padding_value: int = -1000,
    magi = None,
    model = None,
    cp_sharder: typing.Optional[nemo_automodel.components.distributed.context_parallel.sharder.ContextParallelSharder] = None,
    extra_seq_buffers: typing.Optional[dict[str, int]] = None
)

Resolve a ContextParallelSharder and shard the batch; a no-op when no CP prep applies.

Every CP backend is a :class:ContextParallelSharder. A model that owns its CP attention returns one from its prepare_model_inputs_for_cp hook (threaded here as cp_sharder — an explicit parameter, never a batch key, so the batch stays pure tensors); the framework constructs one for magi, TE/THD, and the default load-balanced torch context_parallel path. Resolution order: model-owned > magi > TE > generic. magi and TE also run at cp_size <= 1 (packing conversion / mask-spec activation); model is passed through opaquely to magi for its per-step key/spec stamping.

Parameters:

device_mesh

DeviceMesh

The device mesh; its cp/tp submeshes are read.

batch

Dict[str, torch.Tensor]

The input batch containing (string, torch.Tensor)

Returns: (contextmanager, dict[str, torch.Tensor], ContextParallelSharder)

The forward

nemo_automodel.components.distributed.context_parallel.utils._mesh_dim_size(
    device_mesh,
    dim: str
) -> int

Size of a named submesh, 0 when the mesh or dimension is absent.

nemo_automodel.components.distributed.context_parallel.utils._prepare_cp_sharder(
    model: typing.Any,
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None,
    batch: dict[str, typing.Any],
    padding_token_id: int = 0,
    num_chunks: int = 1,
    loss_mask: torch.Tensor | None = None,
    invoke_pre_embed: bool = True,
    extra_seq_buffers: typing.Optional[dict[str, int]] = None
) -> nemo_automodel.components.distributed.context_parallel.sharder.ContextParallelSharder

Resolve and configure a CP sharder for its public constructor.

The model hook may return a ContextParallelSharder; otherwise this function resolves a framework-owned sharder from the live model’s attention backend and the batch’s token layout. When CP is active and the model exposes prepare_model_inputs_for_cp, that sharder-only hook is invoked directly as a plain method (it constructs a sharder and touches no weights; embed / vision splice / sequence shard run in the model’s own forward per microbatch).

Returns: The resolved and mesh-configured :class:ContextParallelSharder (the identity sharder when no CP prep applies).

Parameters:

model

Any

The (first) model part, or None (e.g. no-model contexts).

device_mesh

DeviceMesh | None

The full device mesh (cp/tp submeshes are read).

batch

dict[str, Any]

The full-sequence batch. Model hook updates are merged in place; :meth:ContextParallelSharder.shard performs the actual sharding.

padding_token_id

intDefaults to 0

Pad sentinel for input_ids.

num_chunks

intDefaults to 1

THD chunk count, forwarded to the hook and TE sharding.

loss_mask

torch.Tensor | NoneDefaults to None

Optional per-token mask forwarded to the batch sharding.

invoke_pre_embed

boolDefaults to True

Invoke the model hook when CP is active or the model owns native THD preparation. Recipes pass False for PP stages without embeddings and for KD paths that never wired model-owned CP.

extra_seq_buffers

Optional[dict[str, int]]Defaults to None

Additional batch keys mapped to their sequence axis (e.g. {"teacher_logits": 1}), padded and sharded alongside the batch on the generic torch path (rejected on the TE THD path; ignored by backends that own their transport).

nemo_automodel.components.distributed.context_parallel.utils._resolve_cp_sharder(
    cp_mesh,
    model_sharder: typing.Optional[nemo_automodel.components.distributed.context_parallel.sharder.ContextParallelSharder],
    magi,
    is_thd: bool,
    num_chunks: int,
    seq_lens_padding_value: int,
    model,
    extra_seq_buffers: typing.Optional[dict[str, int]] = None
) -> nemo_automodel.components.distributed.context_parallel.sharder.ContextParallelSharder

Resolve the ContextParallelSharder for this forward: model-owned > magi > TE > generic > none.

Always returns a sharder: when no CP prep applies, an identity sharder, so callers hold working token verbs at every cp_size and need no branches. The generic torch context_parallel path only shards at cp_size > 1; model-owned, magi, and TE sharders may also run at cp_size <= 1 for native THD packing conversion / mask-spec activation. The magi and THD token layouts depend on batch content (cu_seqlens partitioning / dispatch solver), not just (cp_mesh, seq_len), so their sharders construct without local_token_global_indices and install the index map computed during shard_batch (token verbs raise before the first shard).

nemo_automodel.components.distributed.context_parallel.utils._shard_grad_buffer_for_cp(
    buffer: torch.Tensor,
    seq_dim: int,
    cp_mesh: torch.distributed.device_mesh.DeviceMesh
) -> torch.Tensor

Shard a gradient-bearing buffer with CP’s head-tail load-balancing order.

nemo_automodel.components.distributed.context_parallel.utils._shard_thd_chunk_for_te(
    batch,
    cp_mesh,
    qkv_format,
    seq_lens_padding_value,
    padding_token_id
)

nemo_automodel.components.distributed.context_parallel.utils._uses_te_attention(
    model
) -> bool

Whether a live model uses native or injected TE attention.

nemo_automodel.components.distributed.context_parallel.utils.attach_context_parallel_hooks(
    model: torch.nn.Module
)

Attach forward pre-hooks to self_attn modules to fix attention masks for context parallelism.

Context parallelism shards Q/K/V on the sequence dimension as DTensors, so explicit 4D attention masks would have mismatched shapes. This function registers a hook on every self_attn sub-module that strips the attention_mask kwarg and sets is_causal=True instead, letting SDPA handle causal masking internally.

Based on accelerate.big_modeling._attach_context_parallel_hooks.

nemo_automodel.components.distributed.context_parallel.utils.attach_cp_sdpa_hooks(
    model: torch.nn.Module,
    cp_mesh
) -> None

Inject CP-aware SDPA into self_attn modules for compile + CP>1 correctness.

Problem: when per-layer torch.compile is active, Dynamo traces through the decoder layer including Q/K/V projections. At the F.scaled_dot_product_attention call site, Q/K/V are already local tensors (DTensor metadata was never propagated through the compiled graph). The DTensor SDPA dispatch — which triggers the CP allgather — never fires, so each rank silently attends only to its local sequence shard.

Fix: swap F.scaled_dot_product_attention with a @torch._dynamo.disable wrapper for the duration of each self_attn forward. Dynamo sees the disabled function and creates a graph break there, so:

Everything before (Q/K/V proj + RoPE) is compiled and fused.
The disabled wrapper runs eagerly: re-wraps local Q/K/V as DTensors with Shard(2) on the CP mesh so the DTensor SDPA dispatch fires the allgather.
Everything after (O proj + residual + MLP) is compiled and fused.

Seq dim at the SDPA call is 2: tensors are [B, nH, S/cp_size, D] after HF reshape.

nemo_automodel.components.distributed.context_parallel.utils.attach_te_context_parallel(
    model: torch.nn.Module,
    cp_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    tp_mesh: torch.distributed.device_mesh.DeviceMesh | None = None
) -> int

Configure Transformer Engine attention modules for context and tensor parallelism.

Parameters:

model

torch.nn.Module

Model or pipeline stage containing self_attn modules.

cp_mesh

DeviceMesh | NoneDefaults to None

Optional one-dimensional context-parallel device mesh. When its size is greater than one, every attention module communicates over this mesh; no tensor is mutated.

tp_mesh

DeviceMesh | NoneDefaults to None

Optional one-dimensional tensor-parallel device mesh. When its size is greater than one, Q/K/V use per-rank head shards and every attention module is configured with the corresponding process group.

Returns: int

Number of Transformer Engine attention modules configured.

nemo_automodel.components.distributed.context_parallel.utils.cp_dispatcher_suspended(
    cp_mesh
)

Temporarily suspend torch’s context_parallel ring-SDPA monkeypatch.

A model that embeds and sequence-shards its own primary stream in forward (Megatron-style per-microbatch CP) runs any auxiliary non-CP attention — a VLM vision tower over image patches — inside the ring-SDPA context_parallel context while CP is active. That attention is bidirectional and unsharded, so torch’s load-balanced ring SDPA all-gathers its Q/K/V and rejects it with “Load balancing requires is_causal=True”. The legacy buffer-API context_parallel installs the ring by monkeypatching F.scaled_dot_product_attention; this restores the original SDPA for the wrapped vision forward, then re-installs the ring for the sharded text decoder.

Parameters:

cp_mesh

The context-parallel submesh. A no-op when None or size <= 1, so callers pass the same mesh they shard with; re-enabling needs it.

nemo_automodel.components.distributed.context_parallel.utils.create_context_parallel_ctx(
    cp_mesh: torch.distributed.device_mesh.DeviceMesh,
    cp_buffers: typing.List[torch.Tensor],
    cp_seq_dims: typing.List[int],
    cp_no_restore_buffers: typing.Set[torch.Tensor],
    cp_rotate_method: typing.Optional[str] = None
)

Create a context parallel context.

Parameters:

cp_mesh

DeviceMesh

The device mesh for context parallel.

cp_buffers

List[torch.Tensor]

The buffers for context parallel.

cp_seq_dims

List[int]

The sequence dimensions for context parallel.

cp_no_restore_buffers

Set[torch.Tensor]

The no restore buffers for context parallel.

cp_rotate_method

strDefaults to None

The rotation method for context parallel, such as “allgather” or “addtoall”.

nemo_automodel.components.distributed.context_parallel.utils.get_train_context(
    enable_loss_parallel: bool,
    enable_compiled_autograd: bool,
    cp_context = None
)

Create a train context.

Parameters:

enable_loss_parallel

bool

Whether to enable loss parallelism.

enable_compiled_autograd

bool

Whether to enable compiled autograd.

nemo_automodel.components.distributed.context_parallel.utils.make_cp_batch_for_te(
    cp_mesh,
    batch,
    qkv_format = 'thd',
    padding_token_id: int = 0,
    num_chunks: int = 1,
    seq_lens_padding_value: int = -1000,
    return_local_indices: bool = False
)

Build a CP batch for Transformer Engine using THD format.

This function converts BSHD format batches to THD format and shards them across context parallel ranks for use with Transformer Engine. It processes the batch in chunks if num_chunks > 1, allowing for better memory efficiency with large sequences.

The function performs three main steps:

Converts BSHD format to THD format using split_batch_into_thd_chunks
Optionally splits the batch into multiple chunks for memory efficiency
Shards each chunk across CP ranks using Transformer Engine’s partitioning

Parameters:

cp_mesh

DeviceMesh or None

The device mesh for context parallel. If None or size <= 1, returns the batch in THD format without sharding.

batch

Dict[str, torch.Tensor]

The input batch in BSHD format containing:

input_ids: Input token IDs [batch_size, seq_len] or [batch_size, seq_len, hidden_dim]
labels: Label token IDs [batch_size, seq_len]
position_ids (optional): Position IDs [batch_size, seq_len]
seq_lens: Actual sequence lengths [batch_size, num_packs]
seq_lens_padded: Padded sequence lengths [batch_size, num_packs]

qkv_format

strDefaults to 'thd'

Format for QKV tensors. Currently only “thd” is supported.

padding_token_id

intDefaults to 0

Token ID used for padding in input_ids (default: 0)

num_chunks

intDefaults to 1

Number of chunks to split the batch into. If > 1, the batch dimension is split and each chunk is processed separately (default: 1)

seq_lens_padding_value

intDefaults to -1000

Sentinel value used to indicate padding in seq_lens/seq_lens_padded tensors (default: -1000)

return_local_indices

boolDefaults to False

Also return this rank’s local-token global index map (the thd_get_partitioned_indices partition; an identity arange when CP is inactive; None in chunked mode, where each chunk is its own token space). Used by the THD ContextParallelSharder’s token verbs.

Returns:

Processed batch in THD format (or (dict, LongTensor | None)

Raises:

ValueError: If qkv_format is not “thd”
KeyError: If required fields (seq_lens, seq_lens_padded) are missing from batch

nemo_automodel.components.distributed.context_parallel.utils.unshard_context_parallel_tensor(
    cp_mesh: torch.distributed.device_mesh.DeviceMesh,
    tensor: torch.Tensor,
    seq_dim: int
) -> torch.Tensor

Restore a tensor from PyTorch’s load-balanced context-parallel layout.

Parameters:

cp_mesh

DeviceMesh

One-dimensional context-parallel mesh of size C.

tensor

torch.Tensor

Tensor of shape [..., local_sequence, ...] whose sequence axis is selected by seq_dim and uses PyTorch’s load-balanced CP layout.

seq_dim

int

Axis containing the local sequence extent.

Returns: torch.Tensor

Replicated tensor of shape [..., sequence, ...] with the same axis