> 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.context_parallel.utils ## Module Contents ### Functions | Name | Description | | ---------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------- | | [`_attention_backend`](#nemo_automodel-components-distributed-context_parallel-utils-_attention_backend) | Read the configured attention backend from a live model. | | [`_get_submesh`](#nemo_automodel-components-distributed-context_parallel-utils-_get_submesh) | The named submesh, or None when the mesh or dimension is absent. | | [`_is_multimodal_model`](#nemo_automodel-components-distributed-context_parallel-utils-_is_multimodal_model) | Whether a live model owns a vision or audio tower. | | [`_magi_state_from_model`](#nemo_automodel-components-distributed-context_parallel-utils-_magi_state_from_model) | Recreate the per-forward Magi handle from the model and device mesh. | | [`_make_cp_batch_and_ctx`](#nemo_automodel-components-distributed-context_parallel-utils-_make_cp_batch_and_ctx) | Resolve a ContextParallelSharder and shard the batch; a no-op when no CP prep applies. | | [`_mesh_dim_size`](#nemo_automodel-components-distributed-context_parallel-utils-_mesh_dim_size) | Size of a named submesh, 0 when the mesh or dimension is absent. | | [`_prepare_cp_sharder`](#nemo_automodel-components-distributed-context_parallel-utils-_prepare_cp_sharder) | Resolve and configure a CP sharder for its public constructor. | | [`_resolve_cp_sharder`](#nemo_automodel-components-distributed-context_parallel-utils-_resolve_cp_sharder) | Resolve the ContextParallelSharder for this forward: model-owned > magi > TE > generic > none. | | [`_shard_grad_buffer_for_cp`](#nemo_automodel-components-distributed-context_parallel-utils-_shard_grad_buffer_for_cp) | Shard a gradient-bearing buffer with CP's head-tail load-balancing order. | | [`_shard_thd_chunk_for_te`](#nemo_automodel-components-distributed-context_parallel-utils-_shard_thd_chunk_for_te) | - | | [`_uses_te_attention`](#nemo_automodel-components-distributed-context_parallel-utils-_uses_te_attention) | Whether a live model uses native or injected TE attention. | | [`attach_context_parallel_hooks`](#nemo_automodel-components-distributed-context_parallel-utils-attach_context_parallel_hooks) | Attach forward pre-hooks to self\_attn modules to fix attention masks for context parallelism. | | [`attach_cp_sdpa_hooks`](#nemo_automodel-components-distributed-context_parallel-utils-attach_cp_sdpa_hooks) | Inject CP-aware SDPA into self\_attn modules for compile + CP>1 correctness. | | [`attach_te_context_parallel`](#nemo_automodel-components-distributed-context_parallel-utils-attach_te_context_parallel) | Configure Transformer Engine attention modules for context and tensor parallelism. | | [`cp_dispatcher_suspended`](#nemo_automodel-components-distributed-context_parallel-utils-cp_dispatcher_suspended) | Temporarily suspend torch's `context_parallel` ring-SDPA monkeypatch. | | [`create_context_parallel_ctx`](#nemo_automodel-components-distributed-context_parallel-utils-create_context_parallel_ctx) | Create a context parallel context. | | [`get_train_context`](#nemo_automodel-components-distributed-context_parallel-utils-get_train_context) | Create a train context. | | [`make_cp_batch_for_te`](#nemo_automodel-components-distributed-context_parallel-utils-make_cp_batch_for_te) | Build a CP batch for Transformer Engine using THD format. | | [`unshard_context_parallel_tensor`](#nemo_automodel-components-distributed-context_parallel-utils-unshard_context_parallel_tensor) | Restore a tensor from PyTorch's load-balanced context-parallel layout. | ### API ```python nemo_automodel.components.distributed.context_parallel.utils._attention_backend( model ) -> str | None ``` Read the configured attention backend from a live model. ```python 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. ```python nemo_automodel.components.distributed.context_parallel.utils._is_multimodal_model( model ) -> bool ``` Whether a live model owns a vision or audio tower. ```python 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. ```python 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:** The device mesh; its `cp`/`tp` submeshes are read. The input batch containing (string, torch.Tensor) **Returns:** `(contextmanager, dict[str, torch.Tensor], ContextParallelSharder)` The forward ```python 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. ```python 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:** The (first) model part, or None (e.g. no-model contexts). The full device mesh (`cp`/`tp` submeshes are read). The full-sequence batch. Model hook updates are merged in place; :meth:`ContextParallelSharder.shard` performs the actual sharding. Pad sentinel for `input_ids`. THD chunk count, forwarded to the hook and TE sharding. Optional per-token mask forwarded to the batch sharding. 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. 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). ```python 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). ```python 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. ```python nemo_automodel.components.distributed.context_parallel.utils._shard_thd_chunk_for_te( batch, cp_mesh, qkv_format, seq_lens_padding_value, padding_token_id ) ``` ```python nemo_automodel.components.distributed.context_parallel.utils._uses_te_attention( model ) -> bool ``` Whether a live model uses native or injected TE attention. ```python 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`. ```python 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. ```python 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 or pipeline stage containing `self_attn` modules. 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. 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. ```python 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:** 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. ```python 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:** The device mesh for context parallel. The buffers for context parallel. The sequence dimensions for context parallel. The no restore buffers for context parallel. The rotation method for context parallel, such as "allgather" or "addtoall". ```python 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:** Whether to enable loss parallelism. Whether to enable compiled autograd. ```python 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: 1. Converts BSHD format to THD format using split\_batch\_into\_thd\_chunks 2. Optionally splits the batch into multiple chunks for memory efficiency 3. Shards each chunk across CP ranks using Transformer Engine's partitioning **Parameters:** The device mesh for context parallel. If None or size \<= 1, returns the batch in THD format without sharding. 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] Format for QKV tensors. Currently only "thd" is supported. Token ID used for padding in input\_ids (default: 0) Number of chunks to split the batch into. If > 1, the batch dimension is split and each chunk is processed separately (default: 1) Sentinel value used to indicate padding in seq\_lens/seq\_lens\_padded tensors (default: -1000) 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 ```python 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:** One-dimensional context-parallel mesh of size `C`. Tensor of shape `[..., local_sequence, ...]` whose sequence axis is selected by `seq_dim` and uses PyTorch's load-balanced CP layout. Axis containing the local sequence extent. **Returns:** `torch.Tensor` Replicated tensor of shape `[..., sequence, ...]` with the same axis