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 pre-embedded 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
) -> tuple[typing.Callable[[], typing.ContextManager], dict[str, typing.Any]]

Sequentially shard a pre-embedded batch for block-diagonal CP.

Pads the sequence to a multiple of the CP world size, slices each sequence-aligned tensor to this rank’s contiguous chunk (a differentiable slice for inputs_embeds, so gradients flow back to the trainable vision tower / embedding table), and returns (train_ctx, batch) where entering train_ctx activates the per-document CP SDPA state for the step.

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 attaching this callable to the batch as _cp_make_batch_fn (the model-owned CP hook honored by :func:nemo_automodel.components.distributed.cp_utils.make_cp_batch_and_ctx) and rebinding its attention’s SDPA call for the step.

Parameters:

cp_mesh

DeviceMesh

The context-parallel device (sub)mesh.

tp_mesh

DeviceMesh | None

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

batch

dict[str, Any]

The training batch. Must contain pre-embedded inputs_embeds [B, S, H] (multimodal token replacement happens pre-shard) and is mutated in place: attention_mask is dropped, padding_mask [B, S] (bool, True == pad) is added, and every sequence-aligned tensor is padded then sliced to this rank’s [row_offset, row_offset + S_full/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

Accepted for signature compatibility; unused (the batch is pre-embedded, so there are no token ids to pad).

Returns: Callable[[], ContextManager]

(train_ctx, batch): a zero-arg callable returning the per-step