> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/context_parallel/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

Context-parallel batch sharding.

## Submodules

* **[`nemo_automodel.components.distributed.context_parallel.magi`](/nemo-automodel/nemo_automodel/components/distributed/context_parallel/magi)**
* **[`nemo_automodel.components.distributed.context_parallel.mamba`](/nemo-automodel/nemo_automodel/components/distributed/context_parallel/mamba)**
* **[`nemo_automodel.components.distributed.context_parallel.sharder`](/nemo-automodel/nemo_automodel/components/distributed/context_parallel/sharder)**
* **[`nemo_automodel.components.distributed.context_parallel.utils`](/nemo-automodel/nemo_automodel/components/distributed/context_parallel/utils)**

## Package 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. |

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

`(cp_mesh, padded_seq_len, device) ->
LongTensor` with the global position of each local token —
closed-form for contiguous/round-robin layouts; None for
data-dependent layouts, whose partition arrives with
`shard_layout` (their token verbs raise before the first shard).

---

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

`(cp_mesh, tp_mesh, batch, *, loss_mask=None,
padding_token_id=0) -> (ctx_factory, batch, ShardLayout | None)`.
Pads and shards the batch, installs any backend-owned attention
transport, and reports the shard layout it computed; `shard`
stores it as `shard_layout` for the token verbs.

---

**`shard_layout`** `'ShardLayout | None' = resolved.shard_layout`

The `ShardLayout` of the last `shard_batch`, set
by `shard`. Sharders are built per resolution/hook call, so
the layout never leaks across steps.

---

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

## Docs

- [nemo_automodel.components.distributed.context_parallel.magi](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/context_parallel/magi.md)
- [nemo_automodel.components.distributed.context_parallel.mamba](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/context_parallel/mamba.md)
- [nemo_automodel.components.distributed.context_parallel.sharder](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/context_parallel/sharder.md)
- [nemo_automodel.components.distributed.context_parallel.utils](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/context_parallel/utils.md)