nemo_automodel.components.models.deepseek_v4.cp

Context-parallel helpers for the DeepSeek V4 custom model.

This implements the Miles-style training path: each CP rank owns a contiguous query shard, while K/V and compressed K/V are all-gathered with autograd-aware collectives before DSV4 sparse attention consumes them.

Module Contents

Functions

Name	Description
`_lcm`	-
`_pad_1d`	-
`_pad_length`	-
`_repad_dsv4_packed_batch`	Insert DSV4 compression-safe padding into packed BSHD rows before CP slicing.
`_valid_packed_lengths`	-
`build_dsv4_cp_causal_padding_mask`	Build local-query/global-key additive mask for Miles-style DSV4 CP.
`build_dsv4_cp_packed_causal_padding_mask`	Build local-query/global-key additive mask for packed DSV4 CP.
`build_packed_seq_ids`	Build per-token packed sequence IDs from padded packed lengths.
`dsv4_cp_all_gather`	All-gather activation tensors across CP ranks and concatenate on `dim`.
`dsv4_cp_all_gather_metadata`	All-gather non-differentiable metadata such as padding masks.
`dsv4_cp_enabled`	Return whether a real CP process group is active.
`dsv4_cp_local_seq_multiple`	Required per-CP-rank sequence-length multiple for DSV4 Miles-style CP.
`dsv4_cp_rank`	Return this rank’s index in the DSV4 CP group, or 0 without CP.
`dsv4_cp_size`	Return the DSV4 CP group size, or 1 without CP.
`make_dsv4_contiguous_shard_cp_batch_and_ctx`	Contiguously shard a batch for DeepSeek V4 Miles-style context parallelism.

Data

_SEQ_LENS_PADDING_VALUE

API

nemo_automodel.components.models.deepseek_v4.cp._lcm(
    a: int,
    b: int
) -> int

nemo_automodel.components.models.deepseek_v4.cp._pad_1d(
    values: list[int],
    width: int,
    padding_value: int = _SEQ_LENS_PADDING_VALUE
) -> torch.Tensor

nemo_automodel.components.models.deepseek_v4.cp._pad_length(
    length: int,
    multiple: int
) -> int

nemo_automodel.components.models.deepseek_v4.cp._repad_dsv4_packed_batch(
    batch: dict,
    cp_size: int,
    pad_multiple: int,
    padding_token_id: int,
    sync_packed_length: bool = False,
    loss_mask: torch.Tensor | None = None
) -> tuple[dict, torch.Tensor | None, torch.Tensor]

Insert DSV4 compression-safe padding into packed BSHD rows before CP slicing.

The generic packed dataset may pad each packed sequence only for TE CP. DSV4 compression additionally needs document boundaries to align to compressor windows; CSA then uses packed_seq_ids to reset the previous-window overlap. This routine rebuilds each row from real sequence spans, pads every span to pad_multiple, and appends row-level pack padding with sequence ID 0. When requested for HybridEP, the final physical length is max-reduced across ranks before CP slicing so every rank in a flattened DP x CP expert group is uniform.

nemo_automodel.components.models.deepseek_v4.cp._valid_packed_lengths(
    row: torch.Tensor,
    padding_value: int = _SEQ_LENS_PADDING_VALUE
) -> list[int]

nemo_automodel.components.models.deepseek_v4.cp.build_dsv4_cp_causal_padding_mask(
    position_ids: torch.Tensor,
    key_len: int,
    dtype: torch.dtype,
    device: torch.device,
    cp_group,
    padding_mask: torch.Tensor | None = None,
    sliding_window: int | None = None
) -> torch.Tensor

Build local-query/global-key additive mask for Miles-style DSV4 CP.

position_ids are the local query positions after contiguous CP slicing. Keys are in global sequence order because DSV4 gathers K/V along sequence. padding_mask follows the internal convention True=padding.

nemo_automodel.components.models.deepseek_v4.cp.build_dsv4_cp_packed_causal_padding_mask(
    position_ids: torch.Tensor,
    packed_seq_ids: torch.Tensor,
    dtype: torch.dtype,
    device: torch.device,
    cp_group,
    padding_mask: torch.Tensor | None = None,
    sliding_window: int | None = None
) -> torch.Tensor

Build local-query/global-key additive mask for packed DSV4 CP.

packed_seq_ids is local to the CP rank. The function all-gathers sequence IDs and document-local positions for the key side, then applies same-document, causal, padding, and optional sliding-window constraints.

nemo_automodel.components.models.deepseek_v4.cp.build_packed_seq_ids(
    seq_lens_padded: torch.Tensor,
    seq_len: int,
    device: torch.device,
    padding_value: int = _SEQ_LENS_PADDING_VALUE
) -> torch.Tensor

Build per-token packed sequence IDs from padded packed lengths.

IDs are 1-based within each batch row; 0 marks trailing pack padding that belongs to no sequence. seq_lens_padded may be right-padded with padding_value to make rows rectangular.

nemo_automodel.components.models.deepseek_v4.cp.dsv4_cp_all_gather(
    tensor: torch.Tensor,
    dim: int,
    cp_group
) -> torch.Tensor

All-gather activation tensors across CP ranks and concatenate on dim.

The distributed.nn functional collective preserves autograd, so backward routes gradients for gathered remote slices back to their owning ranks.

nemo_automodel.components.models.deepseek_v4.cp.dsv4_cp_all_gather_metadata(
    tensor: torch.Tensor | None,
    dim: int,
    cp_group
) -> torch.Tensor | None

All-gather non-differentiable metadata such as padding masks.

nemo_automodel.components.models.deepseek_v4.cp.dsv4_cp_enabled(
    cp_group
) -> bool

Return whether a real CP process group is active.

nemo_automodel.components.models.deepseek_v4.cp.dsv4_cp_local_seq_multiple(
    model_or_config
) -> int

Required per-CP-rank sequence-length multiple for DSV4 Miles-style CP.

Compress-ratio layers constrain how the sequence may be split across CP ranks: a ratio-R layer needs each local shard divisible by R, and ratio-4 layers use cross-window overlap so they need 2*R. The returned value is the LCM across all configured compress_ratios (1 when none are configured).

nemo_automodel.components.models.deepseek_v4.cp.dsv4_cp_rank(
    cp_group
) -> int

Return this rank’s index in the DSV4 CP group, or 0 without CP.

nemo_automodel.components.models.deepseek_v4.cp.dsv4_cp_size(
    cp_group
) -> int

Return the DSV4 CP group size, or 1 without CP.

nemo_automodel.components.models.deepseek_v4.cp.make_dsv4_contiguous_shard_cp_batch_and_ctx(
    cp_mesh,
    tp_mesh,
    batch,
    loss_mask = None,
    padding_token_id: int = 0,
    pad_multiple: int | None = None,
    sync_packed_length: bool = False
)

Contiguously shard a batch for DeepSeek V4 Miles-style context parallelism.

Exposed as ContextParallelSharder.shard_batch (via functools.partial to bind pad_multiple) and invoked by the CP dispatch. HybridEP can first max-reduce packed lengths so every rank contributes a uniform token count. Each CP rank then keeps one seq_start:seq_end slice; DSV4 attention all-gathers K/V across CP ranks during forward.

pad_multiple is the required per-CP-rank shard multiple (from dsv4_cp_local_seq_multiple); the global sequence is padded so it is divisible by cp_size and each local shard is divisible by pad_multiple (>= 2). At CP size one, the native THD route keeps the original tensors and packing metadata and exposes physical document IDs to the generic loss contract.

Parameters:

cp_mesh

Context-parallel mesh; size one selects native preparation.

tp_mesh

Tensor-parallel mesh passed to the shared sharder.

batch

Mutable mapping with input_ids [batch, global_sequence] or inputs_embeds [batch, global_sequence, hidden], labels and position_ids [batch, global_sequence], and optional attention_mask [batch, global_sequence] or [batch, 1, global_sequence, global_sequence]. Native packed seq_lens and seq_lens_padded are integer tensors [batch, documents] of real and physical lengths; -1000 pads unused document entries. Existing _packed_seq_ids [batch, global_sequence] take precedence. At CP1, physical lengths produce row-local IDs starting at one, with zero for trailing pack padding.

loss_mask

Defaults to None

Optional tensor [batch, global_sequence], used as labels if absent, otherwise kept as a separate sharded loss mask.

padding_token_id

intDefaults to 0

Token value used for added input padding.

pad_multiple

int | NoneDefaults to None

Required per-rank sequence-length multiple.

sync_packed_length

boolDefaults to False

Whether to synchronize packed lengths before sharding.

Returns:

Context-manager factory, the mutated batch mapping, and a ShardLayout

nemo_automodel.components.models.deepseek_v4.cp._SEQ_LENS_PADDING_VALUE = -1000