> This page is for version Nightly (default).
> 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.
> 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.models.mistral_large4.context_parallel

Contiguous CP batch preparation and differentiable key transport for Large 4.

## Module Contents

### Functions

| Name                                                                                                                     | Description                                                                      |
| ------------------------------------------------------------------------------------------------------------------------ | -------------------------------------------------------------------------------- |
| [`gather_tokens`](#nemo_automodel-components-models-mistral_large4-context_parallel-gather_tokens)                       | Gather \[local_tokens, ...] tensors to \[CP \* local_tokens, ...] in rank order. |
| [`prepare_cp_inputs`](#nemo_automodel-components-models-mistral_large4-context_parallel-prepare_cp_inputs)               | Prepare contiguous CP shards with positions computed before partitioning.        |
| [`prepare_packed_cp_inputs`](#nemo_automodel-components-models-mistral_large4-context_parallel-prepare_packed_cp_inputs) | Preserve packed documents while preparing contiguous model-owned CP.             |
| [`shifted_cp_loss`](#nemo_automodel-components-models-mistral_large4-context_parallel-shifted_cp_loss)                   | Compute HF shifted CE across shard boundaries, scaled for FSDP CP averaging.     |

### Data

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

### API

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

Gather \[local_tokens, ...] tensors to \[CP \* local_tokens, ...] in rank order.

Local token counts must agree. Differentiable transport sums remote-query
contributions back to each key owner; parameter averaging remains FSDP-owned.

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

Prepare contiguous CP shards with positions computed before partitioning.

**Parameters:**

**`batch`** `CPBatch`

Global BSHD token tensors \[batch, sequence, ...] and optional
complete image tensors \[channels, height, width]. THD uses one packed
row per microbatch, with seq_lens/seq_lens_padded document lengths.

---

**`num_chunks`** `int` — default: 1

Positive number of PP microbatches along the batch axis.

---

**`shard_primary`** `bool` — default: True

Whether to shard IDs/embeddings here. False keeps them
global for image splicing and per-microbatch sharding in forward.

---

**Returns:** `dict[str, ContextParallelSharder]`

Mapping containing the model-owned ContextParallelSharder. Labels,

```python
nemo_automodel.components.models.mistral_large4.context_parallel.prepare_packed_cp_inputs(
    batch: nemo_automodel.components.models.mistral_large4.context_parallel.CPBatch,
    num_chunks: int,
    shard_primary: bool
) -> dict[str, nemo_automodel.components.distributed.context_parallel.sharder.ContextParallelSharder]
```

Preserve packed documents while preparing contiguous model-owned CP.

**Parameters:**

**`batch`** `CPBatch`

Global IDs/labels/positions \[chunks, sequence], document lengths
\[chunks, documents], and optional complete images \[channels, height, width].
Exactly one packed row belongs to each PP microbatch.

---

**`num_chunks`** `int`

Positive number of independent pipeline microbatches; one without PP.

---

**`shard_primary`** `bool`

False retains global IDs until complete-image splicing
and/or stage-zero embedding; True shards text IDs here.

---

**Returns:** `dict[str, ContextParallelSharder]`

Sharder yielding local IDs/labels/positions/padding/document IDs

```python
nemo_automodel.components.models.mistral_large4.context_parallel.shifted_cp_loss(
    logits: torch.Tensor,
    labels: torch.Tensor,
    valid: torch.Tensor,
    mesh: torch.distributed.device_mesh.DeviceMesh
) -> torch.Tensor
```

Compute HF shifted CE across shard boundaries, scaled for FSDP CP averaging.

**Parameters:**

**`logits`** `torch.Tensor`

FP32 capped logits \[batch, local_sequence, vocab].

---

**`labels`** `torch.Tensor`

Integer unshifted labels \[batch, local_sequence].

---

**`valid`** `torch.Tensor`

Boolean flattened local \[batch \* local_sequence] keep mask.

---

**`mesh`** `DeviceMesh`

CP mesh; peers must hold shards of the same padded global batch.

---

**Returns:** `torch.Tensor`

Local loss sum times CP / global valid-target count, a differentiable

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