> This page is for version 0.7.0 · 26.10.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

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

# nemo_automodel.components.models.inkling2.context_parallel

Contiguous Inkling2 context shards, differentiable KV and convolution halos.

## Module Contents

### Functions

| Name                                                                                                                     | Description                                                                          |
| ------------------------------------------------------------------------------------------------------------------------ | ------------------------------------------------------------------------------------ |
| [`convolution_history`](#nemo_automodel-components-models-inkling2-context_parallel-convolution_history)                 | Prepend causal history to masked \[batch, local_sequence, channels].                 |
| [`gather_sequence`](#nemo_automodel-components-models-inkling2-context_parallel-gather_sequence)                         | Gather equal-length \[batch, local_sequence, ...] tensors in CP rank order.          |
| [`prepare_cp_inputs`](#nemo_automodel-components-models-inkling2-context_parallel-prepare_cp_inputs)                     | Return a sharder for packed or unpacked global \[batch, sequence] text/image inputs. |
| [`unpack_row_layout`](#nemo_automodel-components-models-inkling2-context_parallel-unpack_row_layout)                     | Normalize shared THD packs to document positions and validity, in place.             |
| [`validate_document_positions`](#nemo_automodel-components-models-inkling2-context_parallel-validate_document_positions) | Require unit increments inside documents and zero at a new document.                 |

### Data

[`CPBatch`](#nemo_automodel-components-models-inkling2-context_parallel-CPBatch)

### API

```python
nemo_automodel.components.models.inkling2.context_parallel.convolution_history(
    masked: torch.Tensor,
    mesh: torch.distributed.device_mesh.DeviceMesh | None,
    width: int
) -> torch.Tensor
```

Prepend causal history to masked \[batch, local_sequence, channels].

**Parameters:**

**`masked`** `torch.Tensor`

This rank's shard \[batch, local_sequence, channels] with invalid positions
already zeroed.

---

**`mesh`** `DeviceMesh | None`

One-dimensional CP mesh; None or a single rank pads with zeros only.

---

**`width`** `int`

Convolution taps; width - 1 history positions are prepended.

---

**Returns:** `torch.Tensor`

New \[batch, local_sequence + width - 1, channels] tensor. Rank zero and positions

```python
nemo_automodel.components.models.inkling2.context_parallel.gather_sequence(
    value: torch.Tensor,
    mesh: torch.distributed.device_mesh.DeviceMesh | None,
    differentiable: bool
) -> torch.Tensor
```

Gather equal-length \[batch, local_sequence, ...] tensors in CP rank order.

**Parameters:**

**`value`** `torch.Tensor`

This rank's shard \[batch, local_sequence, ...]; every rank supplies the same
shape and dtype.

---

**`mesh`** `DeviceMesh | None`

One-dimensional CP mesh; None or a single rank returns the input itself.

---

**`differentiable`** `bool`

Route gradients back to each shard's source rank.

---

**Returns:** `torch.Tensor`

\[batch, global_sequence, ...] with the shards concatenated in rank order, newly

```python
nemo_automodel.components.models.inkling2.context_parallel.prepare_cp_inputs(
    batch: nemo_automodel.components.models.inkling2.context_parallel.CPBatch,
    num_chunks: int = 1
) -> dict[str, nemo_automodel.components.distributed.context_parallel.sharder.ContextParallelSharder]
```

Return a sharder for packed or unpacked global \[batch, sequence] text/image inputs.

IDs, image masks and images remain global for per-microbatch image splicing
inside forward. Labels, positions and padding become local. The shared
recipes own target alignment and loss normalization; no weights are touched.

```python
nemo_automodel.components.models.inkling2.context_parallel.unpack_row_layout(
    batch: nemo_automodel.components.models.inkling2.context_parallel.CPBatch
) -> None
```

Normalize shared THD packs to document positions and validity, in place.

The shared packer supplies real/padded lengths \[batch, documents], padded
with -1000. Padding is determined from those lengths, never from token IDs:
a real EOS remains valid even when the tokenizer also uses it as PAD.

**Parameters:**

**`batch`** `CPBatch`

Mutable mapping with input_ids \[batch, sequence], seq_lens and
seq_lens_padded \[batch, documents], optional position_ids and labels
\[batch, sequence]. Replaces metadata with boolean padding_mask and
integer position_ids \[batch, sequence]; token IDs and labels are unchanged.

---

```python
nemo_automodel.components.models.inkling2.context_parallel.validate_document_positions(
    keep: torch.Tensor,
    positions: torch.Tensor
) -> None
```

Require unit increments inside documents and zero at a new document.

A contiguous CP shard may begin partway through a document. Later resets
explicitly identify new documents; a gap cannot silently resume an old one.

**Parameters:**

**`keep`** `torch.Tensor`

Boolean validity \[batch, sequence].

---

**`positions`** `torch.Tensor`

Integer logical positions \[batch, sequence]. The first real
span may start at a nonzero offset, including after left padding.

---

```python
nemo_automodel.components.models.inkling2.context_parallel.CPBatch: TypeAlias = dict[str, torch.Tensor | list[torch.Tensor] | str | None]
```