> This page is for version 0.4.0 · 26.04.
> 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.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.step5.cp

Memory-bounded sequence transport for Step-5 context parallelism.

## Module Contents

### Classes

| Name                                                                            | Description                                                                |
| ------------------------------------------------------------------------------- | -------------------------------------------------------------------------- |
| [`_GatherHalo`](#nemo_automodel-components-models-step5-cp-_GatherHalo)         | Exchange only the preceding window and return its gradients to its owners. |
| [`_GatherSequence`](#nemo_automodel-components-models-step5-cp-_GatherSequence) | Gather directly into one allocation; reduce-scatter the backward gradient. |

### Functions

| Name                                                                            | Description                                                              |
| ------------------------------------------------------------------------------- | ------------------------------------------------------------------------ |
| [`gather_halo`](#nemo_automodel-components-models-step5-cp-gather_halo)         | Prepend only the preceding causal window using point-to-point transport. |
| [`gather_metadata`](#nemo_automodel-components-models-step5-cp-gather_metadata) | Gather non-differentiable positions, masks or indexer projections.       |
| [`gather_sequence`](#nemo_automodel-components-models-step5-cp-gather_sequence) | Gather contiguous sequence shards with their autograd edges.             |

### API

```python
class nemo_automodel.components.models.step5.cp._GatherHalo()
```

**Bases:** `Function`

Exchange only the preceding window and return its gradients to its owners.

```python
nemo_automodel.components.models.step5.cp._GatherHalo.backward(
    ctx: torch.autograd.function.FunctionCtx,
    grad_output: torch.Tensor
) -> tuple[torch.Tensor, None, None, None]
```

staticmethod

```python
nemo_automodel.components.models.step5.cp._GatherHalo.forward(
    ctx: torch.autograd.function.FunctionCtx,
    x: torch.Tensor,
    group: torch.distributed.ProcessGroup,
    halo_size: int,
    left_padding: int
) -> torch.Tensor
```

staticmethod

```python
class nemo_automodel.components.models.step5.cp._GatherSequence()
```

**Bases:** `Function`

Gather directly into one allocation; reduce-scatter the backward gradient.

```python
nemo_automodel.components.models.step5.cp._GatherSequence.backward(
    ctx: torch.autograd.function.FunctionCtx,
    grad_output: torch.Tensor
) -> tuple[torch.Tensor, None]
```

staticmethod

```python
nemo_automodel.components.models.step5.cp._GatherSequence.forward(
    ctx: torch.autograd.function.FunctionCtx,
    x: torch.Tensor,
    group: torch.distributed.ProcessGroup
) -> torch.Tensor
```

staticmethod

```python
nemo_automodel.components.models.step5.cp.gather_halo(
    x: torch.Tensor,
    group: torch.distributed.ProcessGroup | None,
    halo_size: int,
    left_padding: int = 0
) -> torch.Tensor
```

Prepend only the preceding causal window using point-to-point transport.

**Parameters:**

**`x`** `torch.Tensor`

Tensor \[batch, local_sequence, ...], with equal contiguous CP shards.

---

**`group`** `dist.ProcessGroup | None`

CP process group, or None for a single rank.

---

**`halo_size`** `int`

Maximum number of preceding tokens, normally window - 1.

---

**`left_padding`** `int` — default: 0

Local zero positions before the received halo, for aligning
an attention tile with global key coordinates. No padding is sent.

---

**Returns:** `torch.Tensor`

Tensor \[batch, left_padding + min(rank \* local_sequence, halo_size) + local_sequence, ...].

```python
nemo_automodel.components.models.step5.cp.gather_metadata(
    x: torch.Tensor,
    group: torch.distributed.ProcessGroup | None
) -> torch.Tensor
```

Gather non-differentiable positions, masks or indexer projections.

**Parameters:**

**`x`** `torch.Tensor`

Tensor of shape \[batch, local_sequence, ...] without gradients.

---

**`group`** `dist.ProcessGroup | None`

CP group, or None for a single rank.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, global_sequence, ...]. It aliases x without CP.

```python
nemo_automodel.components.models.step5.cp.gather_sequence(
    x: torch.Tensor,
    group: torch.distributed.ProcessGroup | None
) -> torch.Tensor
```

Gather contiguous sequence shards with their autograd edges.

**Parameters:**

**`x`** `torch.Tensor`

Tensor of shape \[batch, local_sequence, ...], contiguous by CP rank.

---

**`group`** `dist.ProcessGroup | None`

CP group, or None for a single rank.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, global_sequence, ...]. It aliases x without CP.