> 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.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`](#nemo_automodel-components-distributed-context_parallel-sharder-ContextParallelSharder) | CP backend description: how a batch is sharded and where local tokens live. |
| [`ShardLayout`](#nemo_automodel-components-distributed-context_parallel-sharder-ShardLayout)                       | What a `shard_batch` learned about the layout it just applied.              |

### Functions

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

### Data

[`_NO_SHARD_LAYOUT`](#nemo_automodel-components-distributed-context_parallel-sharder-_NO_SHARD_LAYOUT)

[`_ROUND_ROBIN_PAD_FILL`](#nemo_automodel-components-distributed-context_parallel-sharder-_ROUND_ROBIN_PAD_FILL)

### API

```python
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`

---

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

```python
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).

```python
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.

```python
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.

```python
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`

---

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

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

```python
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,

```python
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.

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

```python
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
```

```python
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`** `int` — default: 1

The sequence dimension of the shards.

---

**Returns:** `torch.Tensor`

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

```python
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 | None` — default: None

Device for the returned index tensor.

---

**Returns:** `torch.Tensor`

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

```python
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`.

```python
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`** `int` — default: 1

The sequence dimension of `tensor`.

---

**Returns:** `torch.Tensor`

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

```python
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).

```python
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 | None` — default: None

Device for the returned index tensor.

---

**Returns:** `torch.Tensor`

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

```python
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`** — default: None

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

---

**`padding_token_id`** `int` — default: 0

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

---

**`extra_seq_buffers`** `dict[str, int] | None` — default: 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()`

```python
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`** — default: None

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

---

**`padding_token_id`** `int` — default: 0

Pad sentinel for `input_ids`.

---

**`pad_multiple`** `int` — default: 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] | None` — default: 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] | None` — default: None

Pad sentinels for `extra_seq_keys` (default 0).

---

**`shard_primary`** `bool` — default: 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

```python
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).

```python
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()`

```python
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`** `int` — default: 1

The sequence axis of `tensor`.

---

**`pad_value`** `float | int` — default: 0

Fill for the CP-padding slots appended on `seq_dim`.

---

**`pad_multiple`** `int` — default: 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

```python
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`** `int` — default: 1

The sequence axis of `tensor`.

---

**`pad_value`** `float | int` — default: 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

```python
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`** `int` — default: 1

The sequence dimension of `tensor`.

---

**Returns:** `torch.Tensor`

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

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

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