nemo_automodel.components.models.deepseek_v41.cp

Contiguous context parallelism for the V4.1 text backbone.

Module Contents

Functions

Name	Description
`gather_sequence`	Gather rank-ordered sequence shards, summing remote uses in backward.
`shard_cp_batch`	Prepare packed boundaries once, then keep one contiguous query shard.

API

nemo_automodel.components.models.deepseek_v41.cp.gather_sequence(
    tensor: torch.Tensor,
    group: torch.distributed.ProcessGroup | None
) -> torch.Tensor

Gather rank-ordered sequence shards, summing remote uses in backward.

Parameters:

tensor

torch.Tensor

Tensor of shape [batch, local_sequence, …], with arbitrary trailing dimensions. Equal sequence lengths are required on every rank.

group

dist.ProcessGroup | None

CP process group, or None for an identity operation.

Returns: torch.Tensor

Tensor of shape [batch, global_sequence, …]. Floating activations retain

nemo_automodel.components.models.deepseek_v41.cp.shard_cp_batch(
    cp_mesh: torch.distributed.device_mesh.DeviceMesh | None,
    tp_mesh: torch.distributed.device_mesh.DeviceMesh | None,
    batch: dict[str, typing.Any],
    loss_mask: torch.Tensor | None = None,
    padding_token_id: int | None = 0,
    pad_multiple: int = 1,
    packed_alignment: int = 1,
    sync_packed_length: bool = False
) -> tuple[collections.abc.Callable, dict[str, typing.Any], nemo_automodel.components.distributed.context_parallel.sharder.ShardLayout]

Prepare packed boundaries once, then keep one contiguous query shard.

Parameters:

cp_mesh

DeviceMesh | None

Context-parallel mesh, or None for a local packed forward.

tp_mesh

DeviceMesh | None

Optional tensor-parallel mesh.

batch

dict[str, Any]

Text tensors input_ids, labels, position_ids and optional binary attention_mask [batch, global_sequence]. Packed input adds seq_lens and optional seq_lens_padded [batch, documents]. Labels must already be shifted independently within each document. Replaced in place with local tensors and document IDs [batch, local_sequence].

loss_mask

torch.Tensor | NoneDefaults to None

Optional tensor of shape [batch, global_sequence].

padding_token_id

int | NoneDefaults to 0

Input token used for padding; None uses zero. Validity comes from metadata.

pad_multiple

intDefaults to 1

LCM of this model’s active compression ratios.

packed_alignment

intDefaults to 1

Document alignment for compression groups, independent of CP size.

sync_packed_length

boolDefaults to False

Synchronize the physical packed length across WORLD for HybridEP’s uniform input.

Returns: Callable

Null context factory, the local batch with runtime cp_group, and its ShardLayout.