nemo_automodel.components.distributed.context_parallel.sharder

Context-parallel batch-sharding contract.

Every CP backend is a :class:ContextParallelSharder. A model that owns its CP batch sharding and attention transport returns one from prepare_model_inputs_for_cp under the "cp_sharder" batch key; the framework constructs its own for the remaining backends (torch context_parallel round-robin, TE/THD, MagiAttention) so constructing :class:ContextParallelSharder resolves and configures the backend; callers then invoke sharder.shard(batch). This replaces the retired private batch keys (_cp_make_batch_fn, _cp_metadata_seq_dims, _cp_metadata_pad_values, _cp_full_logits_grad_touch).

The contract is a closed verb set with open implementations: the dataclass slots are plain callables filled with functions from the owning model’s directory or from the framework. local_token_global_indices — the global position of every local token — is the universal layout coordinate system: the default token-tensor shard/gather are synthesized from it, so a sharder only overrides them when it has a cheaper communication pattern. Layouts that are a pure function of (cp_mesh, padded_seq_len) (contiguous, round-robin) provide it at construction; data-dependent layouts (THD cu_seqlens partitioning, magi’s dispatch solver) start as None and report the partition they computed as :class:ShardLayout, which the dispatch stores on the sharder — sharders are built per resolution/hook call, so the layout never leaks across steps. Before the first shard_batch their token verbs raise. The sharder carries no backend tag: nothing may branch on which backend produced it.

This module also hosts the framework’s shard_batch implementations: the shared contiguous-shard batch prep used by models whose CP ranks own contiguous sequence slices (Gemma4, DeepSeek V4), and the torch context_parallel round-robin load-balanced prep with its index map. The TE/THD and magi preps live with their dependencies (context_parallel.utils, context_parallel.magi); the dispatcher wraps them into sharders at resolution time.

Module Contents

Classes

Name	Description
`ContextParallelSharder`	CP backend description: how a batch is sharded and where local tokens live.
`ShardLayout`	What a `shard_batch` learned about the layout it just applied.

Functions

Name	Description
`_cp_rank`	Resolve this rank’s index within the CP submesh (0 without distributed).
`_normalize_cp_primary_and_positions`	Resolve the primary stream and normalize `position_ids` for a CP shard.
`_pad_cp_buffers_to_divisor`	Right-pad every CP buffer on its sequence dim to `cp_divisor` in place.
`_pad_position_ids_seq_dim_`	-
`_pad_tensor_seq_dim_`	-
`_reorder_gathered_token_tensor`	Reassemble rank-order gathered shards into global sequence order.
`contiguous_local_indices`	Global positions owned by this rank under contiguous rank-ordered sharding.
`convert_attention_mask_to_padding_mask`	Pop `attention_mask` and derive `padding_mask` (True == pad) in place.
`gather_token_tensor_by_indices`	Differentiably gather a token-aligned local shard back to the full sequence.
`identity_local_indices`	Global positions owned by this rank when CP is inactive: all of them.
`round_robin_local_indices`	Global positions owned by this rank under torch `context_parallel` load balancing.
`shard_batch_aux_only`	Round-robin CP shard of the no-grad AUXILIARY streams only.
`shard_batch_contiguous`	Prepare and contiguously shard a batch for model-owned CP.
`shard_batch_identity`	`shard_batch` of the identity sharder: a no-op that returns a trivial shard layout.
`shard_batch_load_balanced`	Shard a batch with torch `context_parallel` round-robin load balancing.
`shard_sequence_for_cp_contiguous`	Pad and contiguously shard a full-length sequence tensor inside a model forward.
`shard_sequence_for_cp_round_robin`	Pad and round-robin shard a full-length sequence tensor inside a model forward.
`shard_token_tensor_by_indices`	Keep this rank’s slice of a full-length token-aligned tensor.

Data

_NO_SHARD_LAYOUT

_ROUND_ROBIN_PAD_FILL

API

class nemo_automodel.components.distributed.context_parallel.sharder.ContextParallelSharder(
    model: torch.nn.Module | None = None,
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    batch: dict[str, typing.Any] | None = None,
    shard_batch: collections.abc.Callable[..., tuple[collections.abc.Callable, dict[str, typing.Any], 'ShardLayout | None']] | None = None,
    local_token_global_indices: collections.abc.Callable[..., torch.Tensor] | None = None,
    shard_layout: 'ShardLayout | None' = None,
    padding_token_id: int = 0,
    num_chunks: int = 1,
    loss_mask: torch.Tensor | None = None,
    invoke_pre_embed: bool = True,
    extra_seq_buffers: dict[str, int] | None = None
)

CP backend description: how a batch is sharded and where local tokens live.

_cp_mesh

Any = resolved._cp_mesh

_loss_mask

Tensor | None = resolved._loss_mask

_padding_token_id

int = resolved._padding_token_id

_tp_mesh

Any = resolved._tp_mesh

local_token_global_indices

Callable[..., Tensor] | None = resolved.local_token_global_indices

shard_batch

Callable[..., tuple[Callable, dict[str, Any], 'ShardLayout | None']] = resolved.shard_batch

shard_layout

'ShardLayout | None' = resolved.shard_layout

nemo_automodel.components.distributed.context_parallel.sharder.ContextParallelSharder._indices(
    padded_seq_len: int,
    device
) -> torch.Tensor

nemo_automodel.components.distributed.context_parallel.sharder.ContextParallelSharder.gather_token_tensor(
    tensor: torch.Tensor,
    seq_dim: int = 1,
    trim: bool = False,
    fill: float | int | None = None
) -> torch.Tensor

Differentiably gather a token-aligned local shard to global order.

With trim=True the result is returned in the caller’s original coordinates using the shard layout: sliced back to original_seq_len, un-flattened to input_row_shape (THD), or mapped through the reported position map (fill for input positions whose tokens were dropped, e.g. re-padded pack slots). Raises when no layout is present (nothing to trim to).

nemo_automodel.components.distributed.context_parallel.sharder.ContextParallelSharder.shard(
    batch: dict[str, typing.Any]
) -> tuple[collections.abc.Callable, dict[str, typing.Any]]

Shard a batch and retain its layout for token-aligned tensors.

nemo_automodel.components.distributed.context_parallel.sharder.ContextParallelSharder.shard_token_tensor(
    tensor: torch.Tensor,
    seq_dim: int = 1,
    fill: float | int | None = None
) -> torch.Tensor

Shard a full-length token-aligned tensor exactly like the model inputs.

When shard layout are present (after the first shard_batch), the caller may pass tensors in its own coordinates and the verb applies the same transform the batch went through:

[B, S_in] tensors on a repositioned-row layout (reported position map, e.g. DSV4 packed repad) are scattered into the padded rows, fill filling the pad slots;
tensors matching the reported pre-flatten input_row_shape on a flat-stream (THD) layout are flattened first (the returned shard is in the model’s local stream coordinate);
tensors of original_seq_len are right-padded to padded_seq_len with the explicit fill value;
tensors already at padded_seq_len shard directly.

Any other length raises instead of silently sharding the wrong slice.

class nemo_automodel.components.distributed.context_parallel.sharder.ShardLayout(
    local_token_global_indices: torch.Tensor | None = None,
    original_seq_len: int | None = None,
    padded_seq_len: int | None = None,
    input_row_shape: tuple[int, ...] | None = None,
    input_token_stream_positions: torch.Tensor | None = None
)

Dataclass

What a shard_batch learned about the layout it just applied.

Returned as the third element of shard_batch and stored on the sharder by the dispatch (sharder.shard_layout = layout), so the token verbs can accept and return tensors in the CALLER’s coordinates: pad is an internal detail of the CP layout.

input_row_shape

tuple[int, ...] | None = None

input_token_stream_positions

Tensor | None = None

local_token_global_indices

Tensor | None = None

original_seq_len

int | None = None

padded_seq_len

int | None = None

nemo_automodel.components.distributed.context_parallel.sharder._cp_rank(
    cp_mesh
) -> int

Resolve this rank’s index within the CP submesh (0 without distributed).

nemo_automodel.components.distributed.context_parallel.sharder._normalize_cp_primary_and_positions(
    batch
) -> tuple[str, torch.Tensor, int, int, torch.Tensor, int]

Resolve the primary stream and normalize position_ids for a CP shard.

Pops nothing; determines whether the batch carries inputs_embeds ([batch, sequence, hidden]) or input_ids ([batch, sequence]), injects a 1-D position_ids arange when absent, and expands a shared [1, sequence] position row to the batch size.

Parameters:

batch

The full-sequence batch; position_ids is added/expanded in place.

Returns: str

“(primary_key, primary_seq_tensor, seq_len, batch_size, position_ids,

nemo_automodel.components.distributed.context_parallel.sharder._pad_cp_buffers_to_divisor(
    cp_buffers,
    cp_seq_dims,
    cp_no_restore_buffers,
    batch_buffer_keys,
    seq_len,
    cp_divisor,
    batch
)

Right-pad every CP buffer on its sequence dim to cp_divisor in place.

Pads sequence length to be divisible by 2 * cp_size (required by context_parallel load balancing); the inputs_embeds path can hit arbitrary seq lengths from the VLM collator, so padding happens here rather than relying on dataset-side padding. Each padded, batch-sourced buffer is mirrored back into batch so downstream dict readers see the padded shape.

Parameters:

cp_buffers

Sequence-aligned tensors, each [..., sequence, ...], padded and replaced in place.

cp_seq_dims

Per-buffer sequence axis.

cp_no_restore_buffers

Set of buffers the CP context must not restore; rebuilt to reference the padded tensors.

batch_buffer_keys

Map from buffer index to its batch key (buffers not sourced from the batch are omitted).

seq_len

The pre-pad sequence length.

cp_divisor

2 * cp_size.

batch

The batch dict; padded batch-sourced buffers are mirrored back.

Returns:

(cp_buffers, cp_no_restore_buffers) referencing the padded tensors.

nemo_automodel.components.distributed.context_parallel.sharder._pad_position_ids_seq_dim_(
    position_ids: torch.Tensor,
    seq_dim: int,
    pad_len: int
) -> torch.Tensor

nemo_automodel.components.distributed.context_parallel.sharder._pad_tensor_seq_dim_(
    tensor: torch.Tensor,
    seq_dim: int,
    pad_len: int,
    value: float | int = 0
) -> torch.Tensor

nemo_automodel.components.distributed.context_parallel.sharder._reorder_gathered_token_tensor(
    parts: list[torch.Tensor],
    index_parts: list[torch.Tensor],
    seq_dim: int = 1
) -> torch.Tensor

Reassemble rank-order gathered shards into global sequence order.

Parameters:

parts

list[torch.Tensor]

Per-rank local shards, in rank order.

index_parts

list[torch.Tensor]

Per-rank global token positions, in the same order.

seq_dim

intDefaults to 1

The sequence dimension of the shards.

Returns: torch.Tensor

The full-sequence tensor with position i holding the token whose

nemo_automodel.components.distributed.context_parallel.sharder.contiguous_local_indices(
    cp_mesh,
    padded_seq_len: int,
    device: torch.device | None = None
) -> torch.Tensor

Global positions owned by this rank under contiguous rank-ordered sharding.

Rank r owns [r * L, (r + 1) * L) with L = padded_seq_len // cp_size.

Parameters:

cp_mesh

The context-parallel device (sub)mesh.

padded_seq_len

int

Global sequence length after CP padding; must be divisible by the CP size.

device

torch.device | NoneDefaults to None

Device for the returned index tensor.

Returns: torch.Tensor

A 1-D int64 tensor of this rank’s global token positions.

nemo_automodel.components.distributed.context_parallel.sharder.convert_attention_mask_to_padding_mask(
    batch: dict
) -> None

Pop attention_mask and derive padding_mask (True == pad) in place.

Preserves padding semantics for modules such as MoE routers after the attention mask is stripped for CP. Idempotent: a no-op when padding_mask already exists or there is no attention_mask.

nemo_automodel.components.distributed.context_parallel.sharder.gather_token_tensor_by_indices(
    cp_mesh,
    tensor: torch.Tensor,
    local_indices: torch.Tensor,
    seq_dim: int = 1
) -> torch.Tensor

Differentiably gather a token-aligned local shard back to the full sequence.

Uses torch.distributed.nn.functional.all_gather so gradients route back to the owning ranks, then restores global sequence order from the gathered per-rank indices. Identity when CP is inactive.

Parameters:

cp_mesh

The context-parallel device (sub)mesh.

tensor

torch.Tensor

This rank’s local shard, e.g. [B, S/cp] token logprobs.

local_indices

torch.Tensor

This rank’s global token positions.

seq_dim

intDefaults to 1

The sequence dimension of tensor.

Returns: torch.Tensor

The full-sequence tensor in global order (padding not trimmed).

nemo_automodel.components.distributed.context_parallel.sharder.identity_local_indices(
    cp_mesh,
    padded_seq_len: int,
    device: torch.device | None = None
) -> torch.Tensor

Global positions owned by this rank when CP is inactive: all of them.

Index map of the identity sharder, so consumers run the same token-verb code path at cp_size <= 1 (shard/gather become identities).

nemo_automodel.components.distributed.context_parallel.sharder.round_robin_local_indices(
    cp_mesh,
    padded_seq_len: int,
    device: torch.device | None = None
) -> torch.Tensor

Global positions owned by this rank under torch context_parallel load balancing.

The sequence splits into 2 * cp_size chunks and rank r owns chunks r and 2 * cp_size - 1 - r (head-tail pairing, so causal-attention work stays even across ranks).

Parameters:

cp_mesh

The context-parallel device (sub)mesh.

padded_seq_len

int

Global sequence length after CP padding; must be divisible by 2 * cp_size.

device

torch.device | NoneDefaults to None

Device for the returned index tensor.

Returns: torch.Tensor

A 1-D int64 tensor of this rank’s global token positions, in local

nemo_automodel.components.distributed.context_parallel.sharder.shard_batch_aux_only(
    cp_mesh,
    tp_mesh,
    batch,
    loss_mask = None,
    padding_token_id: int = 0,
    extra_seq_buffers: dict[str, int] | None = None
)

Round-robin CP shard of the no-grad AUXILIARY streams only.

The framework-owned peer of :func:shard_batch_load_balanced for models that embed and sequence-shard their primary stream (input_ids / inputs_embeds and any pixel_* keys) inside their own forward (Megatron-style per-microbatch CP). Pads and round-robin-shards labels/position_ids/loss_mask/padding_mask (and any extra_seq_buffers) exactly like the load-balanced path, installs the same ring-SDPA context_parallel context, and leaves the primary stream full-length in the batch for the model to embed and shard via :func:shard_sequence_for_cp_round_robin.

Because the primary stream never enters the context_parallel buffer list, the grad-carrying-buffer constraint (resize_ rejects tensors requiring grad) does not apply: the model shards its embeddings with a differentiable index_select and gradients route to the embeddings / vision tower.

Parameters:

cp_mesh

The context-parallel device (sub)mesh (size > 1).

tp_mesh

The tensor-parallel device (sub)mesh (or None); unused here.

batch

The full-sequence batch; aux tensors are padded/mirrored in place and the primary stream (input_ids/inputs_embeds) is left full-length. labels is required.

loss_mask

Defaults to None

Optional per-token loss mask [batch, sequence], sharded alongside the batch.

padding_token_id

intDefaults to 0

Accepted per the contract but unused (the primary stream is not sharded here).

extra_seq_buffers

dict[str, int] | NoneDefaults to None

Additional batch keys mapped to their sequence axis, padded and sharded alongside the aux tensors.

Returns:

(ctx_factory, batch, ShardLayout) where entering ctx_factory()

nemo_automodel.components.distributed.context_parallel.sharder.shard_batch_contiguous(
    cp_mesh,
    tp_mesh,
    batch,
    loss_mask = None,
    padding_token_id: int = 0,
    pad_multiple: int = 1,
    extra_seq_keys: dict[str, int] | None = None,
    extra_pad_values: dict[str, typing.Any] | None = None,
    shard_primary: bool = True
)

Prepare and contiguously shard a batch for model-owned CP.

Normalizes the batch (mask conversion, position_ids, labels), pads the sequence to cp_size * max(pad_multiple, 2), then keeps one contiguous seq_start:seq_end slice per CP rank.

shard_primary is the only behavioral switch. When True (default) the primary stream (input_ids / inputs_embeds) is padded and sliced with the aux streams — the dispatch-level shard (e.g. DSV4). When False the primary and pixel streams are left FULL-length for a model that embeds and slices them inside its own forward per microbatch (e.g. Gemma4; see :func:shard_sequence_for_cp_contiguous). The padded layout is identical either way, so the two paths stay bit-for-bit slice-equivalent.

Parameters:

cp_mesh

The context-parallel device (sub)mesh.

tp_mesh

The tensor-parallel device (sub)mesh (or None).

batch

The full-sequence batch; mutated and sharded in place.

loss_mask

Defaults to None

Optional per-token loss mask; used as labels when the batch has none, otherwise sharded alongside the batch.

padding_token_id

intDefaults to 0

Pad sentinel for input_ids.

pad_multiple

intDefaults to 1

Required per-CP-rank shard length multiple (e.g. DSV4’s compress-ratio LCM). The effective divisor is cp_size * max(pad_multiple, 2).

extra_seq_keys

dict[str, int] | NoneDefaults to None

Model-specific per-token batch keys to pad and shard, mapped to their sequence dim (e.g. Gemma4 vision group ids).

extra_pad_values

dict[str, Any] | NoneDefaults to None

Pad sentinels for extra_seq_keys (default 0).

shard_primary

boolDefaults to True

When False, leave the primary/pixel streams full-length for in-forward slicing (aux-only shard); ShardLayout.padded_seq_len is then the length the model must pad its primary to before slicing.

Returns:

(contextlib.nullcontext, batch, ShardLayout) — transport lives in the

nemo_automodel.components.distributed.context_parallel.sharder.shard_batch_identity(
    cp_mesh,
    tp_mesh,
    batch,
    loss_mask = None,
    padding_token_id: int = 0
)

shard_batch of the identity sharder: a no-op that returns a trivial shard layout.

Returned by the dispatch when no CP prep applies, so callers hold working token verbs at every cp_size (nothing was padded: original == padded).

nemo_automodel.components.distributed.context_parallel.sharder.shard_batch_load_balanced(
    cp_mesh,
    tp_mesh,
    batch,
    loss_mask = None,
    padding_token_id: int = 0,
    extra_seq_buffers: dict[str, int] | None = None
)

Shard a batch with torch context_parallel round-robin load balancing.

ContextParallelSharder.shard_batch implementation for the default framework-owned CP path (layout "round_robin", indices from :func:round_robin_local_indices). Assumes an active CP mesh (size > 1). padding_token_id is accepted per the contract but unused: CP-pad slots are zero-filled, sit after every real token under the causal mask, and carry -100 labels.

Returns:

(ctx_factory, batch, ShardLayout) where entering ctx_factory()

nemo_automodel.components.distributed.context_parallel.sharder.shard_sequence_for_cp_contiguous(
    cp_mesh,
    tensor: torch.Tensor,
    seq_dim: int = 1,
    pad_value: float | int = 0,
    pad_multiple: int = 1
) -> tuple[torch.Tensor, torch.Tensor, int]

Pad and contiguously shard a full-length sequence tensor inside a model forward.

The contiguous peer of :func:shard_sequence_for_cp_round_robin: a model that keeps one contiguous seq_start:seq_end slice per CP rank (model-owned p2p ring, e.g. Gemma4) and leaves its primary stream full-length (see :func:shard_batch_contiguous with shard_primary=False) calls this on its embedded, spliced hidden states (and its 4-D per_layer_inputs) to keep this rank’s contiguous shard. It pads the sequence axis to cp_size * max(pad_multiple, 2) — Gemma4’s pad_multiple divisor, NOT the round-robin 2 * cp_size — so the slice aligns with the aux streams the contiguous aux-only sharder sliced. Differentiable (index_select), so gradients route to the embeddings / vision tower.

Parameters:

cp_mesh

The context-parallel (sub)mesh; None or size <= 1 is an identity.

tensor

torch.Tensor

Full-length sequence tensor, e.g. inputs_embeds [B, S, H] or per_layer_inputs [B, S, L, H] (seq axis 1).

seq_dim

intDefaults to 1

The sequence axis of tensor.

pad_value

float | intDefaults to 0

Fill for the CP-padding slots appended on seq_dim.

pad_multiple

intDefaults to 1

Per-CP-rank shard length multiple (the effective divisor is cp_size * max(pad_multiple, 2)), matching :func:shard_batch_contiguous.

Returns: torch.Tensor

(local, local_indices, padded_seq_len): this rank’s contiguous shard

nemo_automodel.components.distributed.context_parallel.sharder.shard_sequence_for_cp_round_robin(
    cp_mesh,
    tensor: torch.Tensor,
    seq_dim: int = 1,
    pad_value: float | int = 0
) -> tuple[torch.Tensor, torch.Tensor, int]

Pad and round-robin shard a full-length sequence tensor inside a model forward.

The in-forward peer of :func:shard_batch_load_balanced: a model that leaves its primary stream full-length (see :func:shard_batch_aux_only) calls this on its embedded, spliced hidden states to keep this CP rank’s head-tail chunk pair, matching the layout the ring-SDPA context and the aux streams were sharded to. Differentiable (index_select), so gradients route to the embeddings / vision tower.

Parameters:

cp_mesh

The context-parallel device (sub)mesh; None or size <= 1 makes this an identity (no pad, arange indices).

tensor

torch.Tensor

Full-length sequence tensor, e.g. inputs_embeds of shape [batch, sequence, hidden]; the axis given by seq_dim is padded to 2 * cp_size and sharded.

seq_dim

intDefaults to 1

The sequence axis of tensor.

pad_value

float | intDefaults to 0

Fill for the CP-padding slots appended on seq_dim (0 for embeddings, matching the zero-padding the load-balanced path applies to floating buffers).

Returns: torch.Tensor

(local, local_indices, padded_seq_len): this rank’s shard laid out

nemo_automodel.components.distributed.context_parallel.sharder.shard_token_tensor_by_indices(
    tensor: torch.Tensor,
    local_indices: torch.Tensor,
    seq_dim: int = 1
) -> torch.Tensor

Keep this rank’s slice of a full-length token-aligned tensor.

Parameters:

tensor

torch.Tensor

Full (padded) sequence tensor, e.g. [B, S] advantages.

local_indices

torch.Tensor

This rank’s global token positions.

seq_dim

intDefaults to 1

The sequence dimension of tensor.

Returns: torch.Tensor

The local shard, laid out identically to the model inputs.

nemo_automodel.components.distributed.context_parallel.sharder._NO_SHARD_LAYOUT = ShardLayout()

nemo_automodel.components.distributed.context_parallel.sharder._ROUND_ROBIN_PAD_FILL = {'labels': -100, 'padding_mask': True, 'attention_mask': False}