bridge.models.deepseek.deepseek_v4_step

DSv4-specific training step with contiguous CP partition support.

DSv4 hybrid attention uses a CSA (Compressed Sparse Attention) compressor that exchanges boundary hidden states between adjacent CP ranks. This requires contiguous token assignment (each rank gets a consecutive slice), unlike the default zigzag interleaved assignment used by standard causal models.

MCore enforces cp_partition_mode=’contiguous’ is only valid with dsv4_hybrid attention (see TransformerConfig validation). Use –step_func dsv4_step for DSv4 SFT/pretrain with CP > 1.

Module Contents

Functions

_packed_metadata_for_forward	Extract packed-sequence metadata for DSv4, including CP partition fields.
_partition_packed_batch_contiguous	Slice a consecutive [start, end) token window for this CP rank.
get_batch	get_batch with DSv4 contiguous CP partition support.
forward_step	Forward training step for DSv4 with contiguous CP partition support.

Data

logger
_DSV4_CURRENT_PACKED_SEQ_PARAM_KEYS
_DSV4_LEGACY_PACKED_SEQ_PARAM_KEYS
_SEQLEN_KEYS

API

bridge.models.deepseek.deepseek_v4_step.logger

‘getLogger(…)’

bridge.models.deepseek.deepseek_v4_step._DSV4_CURRENT_PACKED_SEQ_PARAM_KEYS

(‘cu_seqlens_q’, ‘cu_seqlens_kv’, ‘cu_seqlens_q_padded’, ‘cu_seqlens_kv_padded’, ‘max_seqlen_q’, ‘ma…

bridge.models.deepseek.deepseek_v4_step._DSV4_LEGACY_PACKED_SEQ_PARAM_KEYS

(‘cu_seqlens’, ‘cu_seqlens_unpadded’, ‘cu_seqlens_argmin’, ‘max_seqlen’, ‘cu_seqlens_unpadded_argmin…

bridge.models.deepseek.deepseek_v4_step._packed_metadata_for_forward(batch: dict) → dict | None

Extract packed-sequence metadata for DSv4, including CP partition fields.

bridge.models.deepseek.deepseek_v4_step._SEQLEN_KEYS

‘frozenset(…)’

bridge.models.deepseek.deepseek_v4_step._partition_packed_batch_contiguous(

batch: dict[str, torch.Tensor],

cp_size: int,

) → dict[str, torch.Tensor]

Slice a consecutive [start, end) token window for this CP rank.

Only data tensors (tokens, labels, loss_mask, position_ids, etc.) are sliced. Sequence-length metadata (cu_seqlens, max_seqlen, …) is intentionally kept at global values — the DSv4 CSA compressor needs global sequence boundaries to correctly exchange boundary hidden states between adjacent CP ranks. This mirrors how zigzag mode leaves cu_seqlens untouched.

The packed sequence length must be divisible by cp_size — ensure packed_sequence_size = N * cp_size when running pack_sft_data.

bridge.models.deepseek.deepseek_v4_step.get_batch(

data_iterator: Iterable,

cfg: megatron.bridge.training.config.ConfigContainer,

use_mtp: bool = False,

*,

pg_collection,

vp_stage: int | None = None,

)

get_batch with DSv4 contiguous CP partition support.

Identical to gpt_step.get_batch but dispatches to contiguous partitioning when cfg.model.cp_partition_mode == ‘contiguous’, and injects cp_partition_mode into the batch so get_packed_seq_params can forward it to PackedSeqParams.

bridge.models.deepseek.deepseek_v4_step.forward_step(

state: megatron.bridge.training.state.GlobalState,

data_iterator: Iterable,

model: megatron.core.models.gpt.GPTModel,

return_schedule_plan: bool = False,

)

Forward training step for DSv4 with contiguous CP partition support.