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

Batch padding, sequential sharding, and block-diagonal CP context setup.

## Module Contents

### Functions

| Name                                                                                                                           | Description                                                             |
| ------------------------------------------------------------------------------------------------------------------------------ | ----------------------------------------------------------------------- |
| [`_cp_blockdiag_doc_ids`](#nemo_automodel-components-distributed-blockdiag_cp-batch-_cp_blockdiag_doc_ids)                     | Resolve per-position document ids `[B, S]` (0 == padding) for the mask. |
| [`make_cp_blockdiag_batch_and_ctx`](#nemo_automodel-components-distributed-blockdiag_cp-batch-make_cp_blockdiag_batch_and_ctx) | Sequentially shard a batch for block-diagonal CP.                       |

### API

```python
nemo_automodel.components.distributed.blockdiag_cp.batch._cp_blockdiag_doc_ids(
    batch: dict,
    seq_len: int,
    device,
    batch_size: int
) -> torch.Tensor
```

Resolve per-position document ids `[B, S]` (0 == padding) for the mask.

Prefers the collator's `_packed_seq_ids` (1-based document index per token,
present when a pack holds >1 document). Otherwise falls back to the 4-D
block-causal `attention_mask` diagonal (valid positions) or, lacking both,
treats the whole sequence as a single document.

**Parameters:**

**`batch`** `dict`

The training batch; may contain `_packed_seq_ids` `[B, S]`
(int document index per token) or `attention_mask` (`[B, 1, S, S]`
block-causal bool, or `[B, S]` validity/indexed mask).

---

**`seq_len`** `int`

`S`, the (unpadded) sequence length.

---

**`device`**

Device for the returned tensor.

---

**`batch_size`** `int`

`B`, used for the all-ones fallback.

---

**Returns:** `torch.Tensor`

Per-position document ids `[B, S]` (int64, 0 == padding).

```python
nemo_automodel.components.distributed.blockdiag_cp.batch.make_cp_blockdiag_batch_and_ctx(
    cp_mesh: torch.distributed.device_mesh.DeviceMesh,
    tp_mesh: torch.distributed.device_mesh.DeviceMesh | None,
    batch: dict[str, typing.Any],
    loss_mask: torch.Tensor | None = None,
    padding_token_id: int = 0,
    shard_primary: bool = True
) -> tuple[typing.Callable[[], typing.ContextManager], dict[str, typing.Any], nemo_automodel.components.distributed.context_parallel.sharder.ShardLayout | None]
```

Sequentially shard a batch for block-diagonal CP.

Pads the sequence to a multiple of twice the CP world size, slices each
selected sequence-aligned tensor to this rank's contiguous chunk, and
returns a context whose lifetime activates per-document CP SDPA state.
`shard_primary=False` leaves token ids or embeddings untouched for models
that embed multimodal inputs inside `forward`.

Softmax attention must route through
:func:`~nemo_automodel.components.distributed.blockdiag_cp.runtime.cp_blockdiag_sdpa`
while this context is active. A model opts in by returning a
:class:`~nemo_automodel.components.distributed.context_parallel.sharder.ContextParallelSharder`
whose batch verb is this callable.

**Parameters:**

**`cp_mesh`** `DeviceMesh`

The context-parallel device (sub)mesh.

---

**`tp_mesh`** `DeviceMesh | None`

Accepted for the shared sharder signature; unused
(block-diagonal CP shards only the sequence dimension).

---

**`batch`** `dict[str, Any]`

The training batch. Contains exactly one primary stream:
`inputs_embeds` of shape `[batch, sequence, hidden]` or
`input_ids` of shape `[batch, sequence]`. The batch is mutated in
place: `attention_mask` is dropped, `padding_mask` of shape
`[batch, sequence]` (bool, True == pad) is added, and auxiliary
sequence-aligned tensors are padded then sliced to this rank's
`[row_offset, row_offset + sequence/cp)` chunk.

---

**`loss_mask`** `torch.Tensor | None` — default: None

Optional per-token loss mask `[B, S]`; padded with 0 and
sharded like the other sequence-aligned tensors (stored back into
`batch["loss_mask"]`).

---

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

Fill value for `input_ids` padding when the primary
stream is sharded here.

---

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

Whether to pad and shard the primary stream. Leave False
for models that embed multimodal inputs and shard the resulting
embeddings inside `forward` so FSDP hooks own the vision/embedding
parameter lifecycle.

---

**Returns:** `Callable[[], ContextManager]`

`(train_ctx, batch, layout)`: a zero-arg callable returning the