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`	Resolve per-position document ids `[B, S]` (0 == padding) for the mask.
`make_cp_blockdiag_batch_and_ctx`	Sequentially shard a batch for block-diagonal CP.

API

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

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 | NoneDefaults to 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

intDefaults to 0

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

shard_primary

boolDefaults to 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