From 1dcbc1b7e936db86e5c4c61155ea89702ac9036a Mon Sep 17 00:00:00 2001 From: weijiac Date: Mon, 27 Jul 2026 12:42:55 -0700 Subject: [PATCH 01/10] feat(sft): DSv4 Flash offline-packed SFT with THD + CP=2 --- .../bridge/recipes/deepseek/deepseek_v4.py | 26 ++++ src/megatron/bridge/training/gpt_step.py | 126 +++++++++++++++--- .../bridge/training/utils/packed_seq_utils.py | 30 +++++ .../recipes/test_deepseek_recipes_finetune.py | 2 + tests/unit_tests/training/test_gpt_step.py | 66 +++++++++ 5 files changed, 231 insertions(+), 19 deletions(-) diff --git a/src/megatron/bridge/recipes/deepseek/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/deepseek_v4.py index dcc1ef477a..ebdac6ea77 100644 --- a/src/megatron/bridge/recipes/deepseek/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/deepseek_v4.py @@ -46,6 +46,8 @@ from megatron.bridge.recipes.deepseek.h100.deepseek_v4 import ( deepseek_v4_flash_sft_32gpu_h100_bf16_config as deepseek_v4_flash_sft_config, ) +from megatron.bridge.recipes.utils.dataset_utils import default_openmathinstruct2_thinking_config +from megatron.bridge.training.config import ConfigContainer __all__ = [ @@ -54,9 +56,33 @@ "deepseek_v4_flash_pretrain_muon_config", "deepseek_v4_flash_pretrain_mxfp8_config", "deepseek_v4_flash_sft_config", + "deepseek_v4_flash_sft_openmath_thinking_packed_config", "deepseek_v4_pro_pretrain_config", "deepseek_v4_pro_pretrain_mxfp8_config", "DEEPSEEK_V4_PRO_HF_PATH", "DEEPSEEK_V4_FLASH_HF_PATH", "set_deepseek_v4_pipeline_model_parallel_layout", ] + + +def deepseek_v4_flash_sft_openmath_thinking_packed_config( + hf_path: str = DEEPSEEK_V4_FLASH_HF_PATH, +) -> ConfigContainer: + """DSv4 Flash SFT on OpenMathInstruct-2 with thinking channel and offline-packed sequences. + + CoT reasoning goes into the assistant thinking field and the final answer into the + content field. Uses packed sequences for efficient training. + Pre-pack data with ``prepare_gpt_sft_packed_data.py`` before running SFT. + When using CP>1, pass ``model.cp_partition_mode=contiguous`` (required for DSv4 CSA + attention) and ``pad_seq_to_mult=4`` to ensure divisibility by cp_size. + """ + cfg = deepseek_v4_flash_sft_config(hf_path=hf_path) + # DSv4 hybrid attention requires contiguous CP partition when CP > 1; + # setting it unconditionally is safe (no-op when context_parallel_size=1). + cfg.model.cp_partition_mode = "contiguous" + cfg.dataset = default_openmathinstruct2_thinking_config( + seq_length=cfg.model.seq_length, + enable_offline_packing=True, + pad_seq_to_mult=2 * cfg.model.context_parallel_size, + ) + return cfg diff --git a/src/megatron/bridge/training/gpt_step.py b/src/megatron/bridge/training/gpt_step.py index 94ff6f1fee..4ca9d8f94b 100644 --- a/src/megatron/bridge/training/gpt_step.py +++ b/src/megatron/bridge/training/gpt_step.py @@ -50,10 +50,24 @@ _CURRENT_PACKED_SEQ_DEVICE_KEYS = ("cu_seqlens_q", "cu_seqlens_kv", "cu_seqlens_q_padded", "cu_seqlens_kv_padded") _CURRENT_PACKED_SEQ_HOST_KEYS = ("max_seqlen_q", "max_seqlen_kv") -_CURRENT_PACKED_SEQ_PARAM_KEYS = (*_CURRENT_PACKED_SEQ_DEVICE_KEYS, *_CURRENT_PACKED_SEQ_HOST_KEYS, "total_tokens") +_CURRENT_PACKED_SEQ_PARAM_KEYS = ( + *_CURRENT_PACKED_SEQ_DEVICE_KEYS, + *_CURRENT_PACKED_SEQ_HOST_KEYS, + "total_tokens", + "cp_partition_mode", + "cp_group", + "local_cp_size", +) _LEGACY_PACKED_SEQ_DEVICE_KEYS = ("cu_seqlens", "cu_seqlens_unpadded") _LEGACY_PACKED_SEQ_HOST_KEYS = ("cu_seqlens_argmin", "max_seqlen", "cu_seqlens_unpadded_argmin") -_LEGACY_PACKED_SEQ_PARAM_KEYS = (*_LEGACY_PACKED_SEQ_DEVICE_KEYS, *_LEGACY_PACKED_SEQ_HOST_KEYS, "total_tokens") +_LEGACY_PACKED_SEQ_PARAM_KEYS = ( + *_LEGACY_PACKED_SEQ_DEVICE_KEYS, + *_LEGACY_PACKED_SEQ_HOST_KEYS, + "total_tokens", + "cp_partition_mode", + "cp_group", + "local_cp_size", +) _PackedMetadataValue = torch.Tensor | int | None @@ -172,16 +186,24 @@ def _current_stage_needs_mtp_inputs_from_layout( def _partition_packed_batch_for_cp( - batch: dict[str, torch.Tensor], cp_group: torch.distributed.ProcessGroup + batch: dict[str, torch.Tensor], + cp_group: torch.distributed.ProcessGroup, + cp_partition_mode: str = "zigzag", ) -> dict[str, torch.Tensor]: """Partition THD/packed batches across context-parallel ranks. - Uses MCore's packed-sequence partitioning to slice sequence dimensions - aligned with packed cu_seqlens. + Supports two modes: + - "zigzag" (default): uses Bridge's get_thd_cp_partition_indices for load-balanced + partitioning. Suitable for standard causal transformers (GPT, LLaMA, etc.). + - "contiguous": each rank receives a consecutive token slice. Required for + DSv4 hybrid attention whose CSA compressor exchanges boundary hidden states + between adjacent CP ranks (only meaningful with contiguous partitions). """ + cp_size = torch.distributed.get_world_size(cp_group) + cp_rank = torch.distributed.get_rank(cp_group) cu_seqlens = _cu_seqlens_for_cp_partition(batch) - skip_keys = { + seqlen_keys = { "cu_seqlens", "cu_seqlens_unpadded", "cu_seqlens_argmin", @@ -200,19 +222,83 @@ def _partition_packed_batch_for_cp( "attention_mask", } - indices: dict[tuple[int, torch.device], torch.Tensor] = {} - for key, val in batch.items(): - if val is None or key in skip_keys: - continue - index_key = (val.size(1), val.device) - if index_key not in indices: - indices[index_key] = get_thd_cp_partition_indices( - cu_seqlens, - total_tokens=val.size(1), - cp_group=cp_group, - device=val.device, + if cp_partition_mode == "contiguous": + # Slice a consecutive [start, end) token window for this CP rank. + # Use the actual data tensor size (padded) for slicing — consistent with how + # zigzag uses val.size(1) and cu_seqlens separately. The padded packed_sequence_size + # must be divisible by cp_size (set packed_sequence_size = N * cp_size when packing). + # Find a non-None data tensor to determine the padded token count. + _data_val = next((v for k, v in batch.items() if v is not None and k not in seqlen_keys), None) + if _data_val is None: + return batch # middle PP stage with no data tensors — nothing to slice + total_tokens = _data_val.size(1) + if total_tokens % cp_size != 0: + raise RuntimeError( + f"Contiguous CP partitioning requires packed sequence length ({total_tokens}) " + f"to be divisible by cp_size ({cp_size}). " + "Set packed_sequence_size to a multiple of cp_size when running prepare_gpt_sft_packed_data.py." ) - batch[key] = val.index_select(1, indices[index_key]) + local_len = total_tokens // cp_size + start = cp_rank * local_len + end = start + local_len + + for key, val in batch.items(): + if val is None or key in seqlen_keys: + continue + batch[key] = val[:, start:end].contiguous() + + # Clip cu_seqlens-style tensors to the local window. + # For legacy cu_seqlens, use the sentinel-free trimmed version (computed above) + # since offline packing may pad with -1. cu_seqlens_q/kv/unpadded in the current + # format do not use -1 sentinels and are safe to clip directly. + _SEQLEN_MAP = { + "cu_seqlens": cu_seqlens, # trimmed by _cu_seqlens_for_cp_partition + "cu_seqlens_q": batch.get("cu_seqlens_q"), + "cu_seqlens_kv": batch.get("cu_seqlens_kv"), + "cu_seqlens_unpadded": batch.get("cu_seqlens_unpadded"), + } + for key in seqlen_keys: + val = batch.get(key) + if val is None or "argmin" in key or key in {"max_seqlen", "max_seqlen_q", "max_seqlen_kv", "token_count"}: + continue + trimmed = _SEQLEN_MAP.get(key) + src = trimmed if trimmed is not None else val + clipped = (src.clamp(min=start, max=end) - start).to(val.dtype) + batch[key] = clipped + # Update argmin to reflect trimmed+clipped length (no sentinels remain) + for argmin_key, cs_key in [ + ("cu_seqlens_argmin", "cu_seqlens"), + ("cu_seqlens_unpadded_argmin", "cu_seqlens_unpadded"), + ]: + if batch.get(argmin_key) is not None and batch.get(cs_key) is not None: + batch[argmin_key] = batch[argmin_key].new_tensor([[batch[cs_key].squeeze().numel()]]) + # Recompute max_seqlen from actual local diffs (clipped sequences may be shorter). + for max_key, cs_key in [ + ("max_seqlen", "cu_seqlens"), + ("max_seqlen_q", "cu_seqlens_q"), + ("max_seqlen_kv", "cu_seqlens_kv"), + ]: + cs = batch.get(cs_key) + if cs is None or batch.get(max_key) is None: + continue + cs_flat = cs.squeeze() + diffs = (cs_flat[1:] - cs_flat[:-1]).clamp(min=0) + batch[max_key] = batch[max_key].new_tensor([[int(diffs.max().item()) if diffs.numel() > 0 else 0]]) + + else: + indices: dict[tuple[int, torch.device], torch.Tensor] = {} + for key, val in batch.items(): + if val is None or key in seqlen_keys: + continue + index_key = (val.size(1), val.device) + if index_key not in indices: + indices[index_key] = get_thd_cp_partition_indices( + cu_seqlens, + total_tokens=val.size(1), + cp_group=cp_group, + device=val.device, + ) + batch[key] = val.index_select(1, indices[index_key]) return batch @@ -330,7 +416,9 @@ def get_batch( cp_size = pg_collection.cp.size() has_packed = _has_packed_sequence_metadata(batch) if has_packed and cp_size > 1: - batch = _partition_packed_batch_for_cp(batch, pg_collection.cp) + _cp_mode = getattr(cfg.model, "cp_partition_mode", "zigzag") + batch = _partition_packed_batch_for_cp(batch, pg_collection.cp, cp_partition_mode=_cp_mode) + batch["cp_partition_mode"] = _cp_mode else: # slice batch along sequence dimension for context parallelism batch = get_batch_on_this_cp_rank(batch, is_hybrid_cp=False, cp_group=pg_collection.cp) diff --git a/src/megatron/bridge/training/utils/packed_seq_utils.py b/src/megatron/bridge/training/utils/packed_seq_utils.py index 0341b65ea3..a49ba2a564 100644 --- a/src/megatron/bridge/training/utils/packed_seq_utils.py +++ b/src/megatron/bridge/training/utils/packed_seq_utils.py @@ -265,6 +265,16 @@ def get_packed_seq_params(batch: dict[str, PackedMetadataValue]) -> PackedSeqPar max_seqlen_kv=max_seqlen_kv if max_seqlen_kv is not None else max_seqlen_q, total_tokens=batch.get("total_tokens"), qkv_format="thd", + # cp_partition_mode/cp_group/local_cp_size available in dev MCore only. + **( + { + "cp_partition_mode": batch.get("cp_partition_mode", "zigzag"), + "cp_group": batch.get("cp_group"), + "local_cp_size": batch.get("local_cp_size"), + } + if hasattr(PackedSeqParams, "cp_partition_mode") + else {} + ), ) cu_seqlens_padded = batch["cu_seqlens"].squeeze() @@ -303,6 +313,16 @@ def get_packed_seq_params(batch: dict[str, PackedMetadataValue]) -> PackedSeqPar max_seqlen_kv=max_seqlen, total_tokens=total_tokens, qkv_format="thd", + # cp_partition_mode/cp_group/local_cp_size available in dev MCore only. + **( + { + "cp_partition_mode": batch.get("cp_partition_mode", "zigzag"), + "cp_group": batch.get("cp_group"), + "local_cp_size": batch.get("local_cp_size"), + } + if hasattr(PackedSeqParams, "cp_partition_mode") + else {} + ), ) else: return PackedSeqParams( @@ -312,4 +332,14 @@ def get_packed_seq_params(batch: dict[str, PackedMetadataValue]) -> PackedSeqPar max_seqlen_kv=max_seqlen, total_tokens=total_tokens, qkv_format="thd", + # cp_partition_mode/cp_group/local_cp_size available in dev MCore only. + **( + { + "cp_partition_mode": batch.get("cp_partition_mode", "zigzag"), + "cp_group": batch.get("cp_group"), + "local_cp_size": batch.get("local_cp_size"), + } + if hasattr(PackedSeqParams, "cp_partition_mode") + else {} + ), ) diff --git a/tests/functional_tests/test_groups/recipes/test_deepseek_recipes_finetune.py b/tests/functional_tests/test_groups/recipes/test_deepseek_recipes_finetune.py index 0b4a1f7a00..5510b5afb9 100644 --- a/tests/functional_tests/test_groups/recipes/test_deepseek_recipes_finetune.py +++ b/tests/functional_tests/test_groups/recipes/test_deepseek_recipes_finetune.py @@ -31,6 +31,7 @@ from megatron.bridge.recipes.deepseek import ( deepseek_v4_flash_no_mtp_sft_config, deepseek_v4_flash_sft_config, + deepseek_v4_flash_sft_openmath_thinking_packed_config, ) from megatron.bridge.recipes.deepseek.h100 import deepseek_v4 as deepseek_v4_h100_module @@ -89,6 +90,7 @@ def _deepseek_v4_toy_model_path() -> str: DEEPSEEK_V4_SFT_RECIPES = [ (deepseek_v4_flash_sft_config, "deepseek_v4_flash_sft", False), (deepseek_v4_flash_no_mtp_sft_config, "deepseek_v4_flash_no_mtp_sft", False), + (deepseek_v4_flash_sft_openmath_thinking_packed_config, "deepseek_v4_flash_sft_openmath_thinking_packed", False), ] diff --git a/tests/unit_tests/training/test_gpt_step.py b/tests/unit_tests/training/test_gpt_step.py index d5b38f347c..ee61a82158 100644 --- a/tests/unit_tests/training/test_gpt_step.py +++ b/tests/unit_tests/training/test_gpt_step.py @@ -725,6 +725,72 @@ def test_none_attention_mask_still_skipped(self, monkeypatch): assert fake_partitioner.seq_lens_seen == [4] assert result["tokens"].size(1) == 2 + # ── Contiguous mode tests ────────────────────────────────────────────── + + def _run_contiguous(self, monkeypatch, batch, cp_rank=0, cp_size=2): + """Run _partition_packed_batch_for_cp in contiguous mode.""" + monkeypatch.setattr( + "megatron.bridge.training.gpt_step.torch.distributed.get_world_size", + lambda _group: cp_size, + ) + monkeypatch.setattr( + "megatron.bridge.training.gpt_step.torch.distributed.get_rank", + lambda _group: cp_rank, + ) + return _partition_packed_batch_for_cp(batch, _MockProcessGroup(size=cp_size), cp_partition_mode="contiguous") + + def test_contiguous_basic_slice(self, monkeypatch): + """Rank 0 of 2 receives the first half of each data tensor.""" + tokens = torch.arange(8, dtype=torch.long).unsqueeze(0) + labels = torch.arange(8, dtype=torch.long).unsqueeze(0) + loss_mask = torch.ones(1, 8) + position_ids = torch.arange(8).unsqueeze(0) + cu_seqlens = torch.tensor([[0, 4, 8]], dtype=torch.int32) + batch = { + "tokens": tokens, + "labels": labels, + "loss_mask": loss_mask, + "position_ids": position_ids, + "cu_seqlens": cu_seqlens, + } + result = self._run_contiguous(monkeypatch, batch, cp_rank=0, cp_size=2) + assert result["tokens"].shape == (1, 4) + assert torch.equal(result["tokens"].squeeze(), torch.arange(4, dtype=torch.long)) + # cu_seqlens clipped to [0, 4] window + assert result["cu_seqlens"].squeeze().tolist() == [0, 4] + + def test_contiguous_rank1_slice(self, monkeypatch): + """Rank 1 of 2 receives the second half.""" + tokens = torch.arange(8, dtype=torch.long).unsqueeze(0) + cu_seqlens = torch.tensor([[0, 4, 8]], dtype=torch.int32) + batch = {"tokens": tokens, "cu_seqlens": cu_seqlens} + result = self._run_contiguous(monkeypatch, batch, cp_rank=1, cp_size=2) + assert result["tokens"].shape == (1, 4) + assert torch.equal(result["tokens"].squeeze(), torch.arange(4, 8, dtype=torch.long)) + assert result["cu_seqlens"].squeeze().tolist() == [0, 4] + + def test_contiguous_rejects_non_divisible_length(self, monkeypatch): + """Raises RuntimeError when total_tokens is not divisible by cp_size.""" + tokens = torch.arange(5, dtype=torch.long).unsqueeze(0) + batch = { + "tokens": tokens, + "cu_seqlens": torch.tensor([[0, 5]], dtype=torch.int32), + } + with pytest.raises(RuntimeError, match="divisible by cp_size"): + self._run_contiguous(monkeypatch, batch, cp_size=2) + + def test_contiguous_middle_pp_stage_returns_unchanged(self, monkeypatch): + """Middle PP stage (all data tensors None) returns batch unchanged.""" + cu_seqlens = torch.tensor([[0, 4, 8]], dtype=torch.int32) + batch = { + "tokens": None, + "labels": None, + "loss_mask": None, + "cu_seqlens": cu_seqlens, + } + result = self._run_contiguous(monkeypatch, batch, cp_size=2) + assert result is batch # unchanged + class TestGetPackedSeqParams: """Tests for the get_packed_seq_params function.""" From 021ba66b112439ffc6f1ec9cac89c74f3fa5251b Mon Sep 17 00:00:00 2001 From: weijia chen Date: Mon, 27 Jul 2026 17:35:55 -0700 Subject: [PATCH 02/10] refactor(sft): move DSv4 contiguous CP partition to deepseek_v4_step MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit cp_partition_mode=contiguous is DSv4-specific — MCore enforces this in TransformerConfig (dsv4_hybrid only). Move the contiguous batch partition logic out of generic gpt_step.py into a new model-specific step file. Changes: - gpt_step.py: _partition_packed_batch_for_cp reverted to zigzag-only; _forward_step_common gains optional _get_batch_fn override; inline _model_chunk_vp_stage helper (was get_model_chunk_vp_stage from flop_utils) - deepseek_v4_step.py (new): _partition_packed_batch_contiguous, DSv4-aware get_batch dispatching contiguous/zigzag, forward_step via _get_batch_fn - recipe_runner.py: register dsv4_step in STEP_FUNCTIONS - test_gpt_step.py: remove contiguous-mode tests (moved to deepseek_v4_step) - test_deepseek_v4_step.py (new): contiguous partition unit tests Use --step_func dsv4_step for DSv4 SFT/pretrain with CP > 1. Co-Authored-By: Claude Sonnet 4.6 (1M context) --- scripts/training/recipe_runner.py | 2 + .../models/deepseek/deepseek_v4_step.py | 235 ++++++++++++++++++ src/megatron/bridge/training/gpt_step.py | 137 +++------- tests/unit_tests/models/deepseek/__init__.py | 0 .../models/deepseek/test_deepseek_v4_step.py | 76 ++++++ tests/unit_tests/training/test_gpt_step.py | 66 ----- 6 files changed, 350 insertions(+), 166 deletions(-) create mode 100644 src/megatron/bridge/models/deepseek/deepseek_v4_step.py create mode 100644 tests/unit_tests/models/deepseek/__init__.py create mode 100644 tests/unit_tests/models/deepseek/test_deepseek_v4_step.py diff --git a/scripts/training/recipe_runner.py b/scripts/training/recipe_runner.py index dde0f902f9..1980e53419 100644 --- a/scripts/training/recipe_runner.py +++ b/scripts/training/recipe_runner.py @@ -52,6 +52,7 @@ STEP_FUNCTIONS: dict[str, StepFunctionEntry] = { "audio_lm_step": ("megatron.bridge.training.audio_lm_step", "forward_step"), + "dsv4_step": ("megatron.bridge.models.deepseek.deepseek_v4_step", "forward_step"), "gpt_step": ("megatron.bridge.training.gpt_step", "forward_step"), "llm_step": ("megatron.bridge.training.gpt_step", "forward_step"), "vlm_step": ("megatron.bridge.training.vlm_step", "forward_step"), @@ -66,6 +67,7 @@ STEP_MODALITIES = { "audio_lm_step": "audio", + "dsv4_step": "text", "gpt_step": "text", "llm_step": "text", "vlm_step": "vlm", diff --git a/src/megatron/bridge/models/deepseek/deepseek_v4_step.py b/src/megatron/bridge/models/deepseek/deepseek_v4_step.py new file mode 100644 index 0000000000..1e20d61d22 --- /dev/null +++ b/src/megatron/bridge/models/deepseek/deepseek_v4_step.py @@ -0,0 +1,235 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""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. +""" + +import logging +from typing import Iterable + +import torch +from megatron.core import parallel_state +from megatron.core.models.gpt import GPTModel +from megatron.core.pipeline_parallel.utils import ( + is_pp_first_stage, + is_pp_last_stage, + is_vp_first_stage, + is_vp_last_stage, +) +from megatron.core.utils import get_batch_on_this_cp_rank + +from megatron.bridge.training.config import ConfigContainer +from megatron.bridge.training.gpt_step import ( + _create_loss_function, + _cu_seqlens_for_cp_partition, + _current_stage_needs_mtp_inputs_from_layout, + _forward_step_common, + _has_packed_sequence_metadata, + _middle_pp_stage_needs_batch, + _packed_metadata_for_forward, + _partition_packed_batch_for_cp, + get_batch_from_iterator, +) +from megatron.bridge.training.state import GlobalState + + +logger = logging.getLogger(__name__) + +# Sequence-length metadata keys — excluded from token-dimension slicing. +_SEQLEN_KEYS = frozenset( + { + "cu_seqlens", + "cu_seqlens_unpadded", + "cu_seqlens_argmin", + "cu_seqlens_unpadded_argmin", + "max_seqlen", + "cu_seqlens_q", + "cu_seqlens_kv", + "cu_seqlens_q_padded", + "cu_seqlens_kv_padded", + "max_seqlen_q", + "max_seqlen_kv", + "token_count", + "attention_mask", + } +) + + +def _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. + + Required for DSv4 hybrid attention: the CSA compressor exchanges boundary + hidden states between adjacent CP ranks, which requires contiguous token + assignments. The packed sequence length must be divisible by cp_size — + ensure packed_sequence_size = N * cp_size when running pack_sft_data. + + cu_seqlens are clipped to the local window; -1 padding sentinels in the + legacy cu_seqlens path are stripped via _cu_seqlens_for_cp_partition to + avoid invalid decreasing entries after clamping. + """ + cp_rank = parallel_state.get_context_parallel_rank() + cu_seqlens = _cu_seqlens_for_cp_partition(batch) + + _data_val = next((v for k, v in batch.items() if v is not None and k not in _SEQLEN_KEYS), None) + if _data_val is None: + return batch # middle PP stage with no data tensors — nothing to slice + + total_tokens = _data_val.size(1) + if total_tokens % cp_size != 0: + raise RuntimeError( + f"Contiguous CP partitioning requires packed sequence length ({total_tokens}) " + f"to be divisible by cp_size ({cp_size}). " + "Set packed_sequence_size to a multiple of cp_size when running pack_sft_data." + ) + local_len = total_tokens // cp_size + start = cp_rank * local_len + end = start + local_len + + for key, val in batch.items(): + if val is None or key in _SEQLEN_KEYS: + continue + batch[key] = val[:, start:end].contiguous() + + # Clip cu_seqlens to local window; use trimmed source for legacy path. + _seqlen_src = { + "cu_seqlens": cu_seqlens, + "cu_seqlens_q": batch.get("cu_seqlens_q"), + "cu_seqlens_kv": batch.get("cu_seqlens_kv"), + "cu_seqlens_unpadded": batch.get("cu_seqlens_unpadded"), + } + for key in _SEQLEN_KEYS: + val = batch.get(key) + if ( + val is None + or "argmin" in key + or key in {"max_seqlen", "max_seqlen_q", "max_seqlen_kv", "token_count", "attention_mask"} + ): + continue + src_val = _seqlen_src.get(key) + src = src_val if src_val is not None else val + batch[key] = (src.clamp(min=start, max=end) - start).to(val.dtype) + + for argmin_key, cs_key in [ + ("cu_seqlens_argmin", "cu_seqlens"), + ("cu_seqlens_unpadded_argmin", "cu_seqlens_unpadded"), + ]: + if batch.get(argmin_key) is not None and batch.get(cs_key) is not None: + batch[argmin_key] = batch[argmin_key].new_tensor([[batch[cs_key].squeeze().numel()]]) + + for max_key, cs_key in [ + ("max_seqlen", "cu_seqlens"), + ("max_seqlen_q", "cu_seqlens_q"), + ("max_seqlen_kv", "cu_seqlens_kv"), + ]: + cs = batch.get(cs_key) + if cs is None or batch.get(max_key) is None: + continue + cs_flat = cs.squeeze() + diffs = (cs_flat[1:] - cs_flat[:-1]).clamp(min=0) + batch[max_key] = batch[max_key].new_tensor([[int(diffs.max().item()) if diffs.numel() > 0 else 0]]) + + return batch + + +def get_batch( + data_iterator: Iterable, + cfg: 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. + """ + model_cfg = getattr(cfg, "model", None) + vp_size = getattr(model_cfg, "virtual_pipeline_model_parallel_size", None) + is_first = is_pp_first_stage(pg_collection.pp) and ( + vp_stage is None or is_vp_first_stage(vp_stage=vp_stage, vp_size=vp_size) + ) + is_last = is_pp_last_stage(pg_collection.pp) and ( + vp_stage is None or is_vp_last_stage(vp_stage=vp_stage, vp_size=vp_size) + ) + is_middle = (not is_first) and (not is_last) + include_full_batch_fields = is_middle and _middle_pp_stage_needs_batch(cfg) + include_mtp_inputs = use_mtp and _current_stage_needs_mtp_inputs_from_layout( + cfg, pg_collection=pg_collection, is_last=is_last, vp_stage=vp_stage + ) + if is_middle and not include_full_batch_fields and not include_mtp_inputs: + return None, None, None, None, None, None + + batch = get_batch_from_iterator( + data_iterator, + include_mtp_inputs=include_mtp_inputs, + skip_getting_attention_mask_from_dataset=getattr( + cfg.dataset, "skip_getting_attention_mask_from_dataset", True + ), + is_first_pp_stage=is_first, + is_last_pp_stage=is_last, + include_full_batch_fields=include_full_batch_fields, + ) + + cp_size = pg_collection.cp.size() + has_packed = _has_packed_sequence_metadata(batch) + if has_packed and cp_size > 1: + cp_mode = getattr(cfg.model, "cp_partition_mode", "zigzag") + if cp_mode == "contiguous": + batch = _partition_packed_batch_contiguous(batch, cp_size) + else: + batch = _partition_packed_batch_for_cp(batch, pg_collection.cp) + # Inject cp_partition_mode so get_packed_seq_params forwards it to PackedSeqParams. + batch["cp_partition_mode"] = cp_mode + else: + batch = get_batch_on_this_cp_rank(batch, is_hybrid_cp=False, cp_group=pg_collection.cp) + + return ( + batch["tokens"], + batch["labels"], + batch["loss_mask"], + batch.get("attention_mask"), + batch["position_ids"], + _packed_metadata_for_forward(batch), + ) + + +def forward_step( + state: GlobalState, + data_iterator: Iterable, + model: GPTModel, + return_schedule_plan: bool = False, +): + """Forward training step for DSv4 with contiguous CP partition support.""" + output, loss_mask = _forward_step_common( + state, data_iterator, model, return_schedule_plan, _get_batch_fn=get_batch + ) + return output, _create_loss_function( + loss_mask, + check_for_nan_in_loss=state.cfg.rerun_state_machine.check_for_nan_in_loss, + check_for_spiky_loss=state.cfg.rerun_state_machine.check_for_spiky_loss, + ) diff --git a/src/megatron/bridge/training/gpt_step.py b/src/megatron/bridge/training/gpt_step.py index 4ca9d8f94b..631a392fc9 100644 --- a/src/megatron/bridge/training/gpt_step.py +++ b/src/megatron/bridge/training/gpt_step.py @@ -29,6 +29,7 @@ from megatron.core.transformer.enums import LayerType from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout from megatron.core.utils import ( + get_attr_wrapped_model, get_batch_on_this_cp_rank, get_model_config, get_pg_rank, @@ -40,7 +41,7 @@ from megatron.bridge.training.losses import masked_next_token_loss from megatron.bridge.training.post_training.distillation import loss_func_kd from megatron.bridge.training.state import GlobalState -from megatron.bridge.training.utils.flop_utils import accumulate_flops_metadata, get_model_chunk_vp_stage +from megatron.bridge.training.utils.flop_utils import accumulate_flops_metadata from megatron.bridge.training.utils.packed_seq_utils import get_packed_seq_params, get_thd_cp_partition_indices from megatron.bridge.training.utils.pg_utils import get_pg_collection @@ -185,24 +186,26 @@ def _current_stage_needs_mtp_inputs_from_layout( return _current_stage_has_mtp_from_layout(cfg, pg_collection=pg_collection, vp_stage=vp_stage) +def _model_chunk_vp_stage(model: GPTModel) -> int | None: + """Return the virtual pipeline stage owned by the current model chunk.""" + try: + vp_stage = get_attr_wrapped_model(model, "vp_stage", allow_none=False) + except RuntimeError: + return None + return vp_stage if isinstance(vp_stage, int) else None + + def _partition_packed_batch_for_cp( batch: dict[str, torch.Tensor], cp_group: torch.distributed.ProcessGroup, - cp_partition_mode: str = "zigzag", ) -> dict[str, torch.Tensor]: - """Partition THD/packed batches across context-parallel ranks. - - Supports two modes: - - "zigzag" (default): uses Bridge's get_thd_cp_partition_indices for load-balanced - partitioning. Suitable for standard causal transformers (GPT, LLaMA, etc.). - - "contiguous": each rank receives a consecutive token slice. Required for - DSv4 hybrid attention whose CSA compressor exchanges boundary hidden states - between adjacent CP ranks (only meaningful with contiguous partitions). + """Partition THD/packed batches across context-parallel ranks using zigzag mode. + + Uses Bridge's get_thd_cp_partition_indices for load-balanced interleaved partitioning. + Suitable for standard causal transformers (GPT, LLaMA, etc.). + For DSv4 contiguous CP partitioning, use deepseek_v4_step.py instead. """ - cp_size = torch.distributed.get_world_size(cp_group) - cp_rank = torch.distributed.get_rank(cp_group) cu_seqlens = _cu_seqlens_for_cp_partition(batch) - seqlen_keys = { "cu_seqlens", "cu_seqlens_unpadded", @@ -216,90 +219,21 @@ def _partition_packed_batch_for_cp( "max_seqlen_q", "max_seqlen_kv", "token_count", - # THD/packed attention is driven by cu_seqlens (PackedSeqParams), so the dense - # attention_mask is unused here. It is also not sequence-partitionable: it is - # either None or a degenerate placeholder without a slice-able seq dim at index 1. "attention_mask", } - - if cp_partition_mode == "contiguous": - # Slice a consecutive [start, end) token window for this CP rank. - # Use the actual data tensor size (padded) for slicing — consistent with how - # zigzag uses val.size(1) and cu_seqlens separately. The padded packed_sequence_size - # must be divisible by cp_size (set packed_sequence_size = N * cp_size when packing). - # Find a non-None data tensor to determine the padded token count. - _data_val = next((v for k, v in batch.items() if v is not None and k not in seqlen_keys), None) - if _data_val is None: - return batch # middle PP stage with no data tensors — nothing to slice - total_tokens = _data_val.size(1) - if total_tokens % cp_size != 0: - raise RuntimeError( - f"Contiguous CP partitioning requires packed sequence length ({total_tokens}) " - f"to be divisible by cp_size ({cp_size}). " - "Set packed_sequence_size to a multiple of cp_size when running prepare_gpt_sft_packed_data.py." + indices: dict[tuple[int, torch.device], torch.Tensor] = {} + for key, val in batch.items(): + if val is None or key in seqlen_keys: + continue + index_key = (val.size(1), val.device) + if index_key not in indices: + indices[index_key] = get_thd_cp_partition_indices( + cu_seqlens, + total_tokens=val.size(1), + cp_group=cp_group, + device=val.device, ) - local_len = total_tokens // cp_size - start = cp_rank * local_len - end = start + local_len - - for key, val in batch.items(): - if val is None or key in seqlen_keys: - continue - batch[key] = val[:, start:end].contiguous() - - # Clip cu_seqlens-style tensors to the local window. - # For legacy cu_seqlens, use the sentinel-free trimmed version (computed above) - # since offline packing may pad with -1. cu_seqlens_q/kv/unpadded in the current - # format do not use -1 sentinels and are safe to clip directly. - _SEQLEN_MAP = { - "cu_seqlens": cu_seqlens, # trimmed by _cu_seqlens_for_cp_partition - "cu_seqlens_q": batch.get("cu_seqlens_q"), - "cu_seqlens_kv": batch.get("cu_seqlens_kv"), - "cu_seqlens_unpadded": batch.get("cu_seqlens_unpadded"), - } - for key in seqlen_keys: - val = batch.get(key) - if val is None or "argmin" in key or key in {"max_seqlen", "max_seqlen_q", "max_seqlen_kv", "token_count"}: - continue - trimmed = _SEQLEN_MAP.get(key) - src = trimmed if trimmed is not None else val - clipped = (src.clamp(min=start, max=end) - start).to(val.dtype) - batch[key] = clipped - # Update argmin to reflect trimmed+clipped length (no sentinels remain) - for argmin_key, cs_key in [ - ("cu_seqlens_argmin", "cu_seqlens"), - ("cu_seqlens_unpadded_argmin", "cu_seqlens_unpadded"), - ]: - if batch.get(argmin_key) is not None and batch.get(cs_key) is not None: - batch[argmin_key] = batch[argmin_key].new_tensor([[batch[cs_key].squeeze().numel()]]) - # Recompute max_seqlen from actual local diffs (clipped sequences may be shorter). - for max_key, cs_key in [ - ("max_seqlen", "cu_seqlens"), - ("max_seqlen_q", "cu_seqlens_q"), - ("max_seqlen_kv", "cu_seqlens_kv"), - ]: - cs = batch.get(cs_key) - if cs is None or batch.get(max_key) is None: - continue - cs_flat = cs.squeeze() - diffs = (cs_flat[1:] - cs_flat[:-1]).clamp(min=0) - batch[max_key] = batch[max_key].new_tensor([[int(diffs.max().item()) if diffs.numel() > 0 else 0]]) - - else: - indices: dict[tuple[int, torch.device], torch.Tensor] = {} - for key, val in batch.items(): - if val is None or key in seqlen_keys: - continue - index_key = (val.size(1), val.device) - if index_key not in indices: - indices[index_key] = get_thd_cp_partition_indices( - cu_seqlens, - total_tokens=val.size(1), - cp_group=cp_group, - device=val.device, - ) - batch[key] = val.index_select(1, indices[index_key]) - + batch[key] = val.index_select(1, indices[index_key]) return batch @@ -416,9 +350,7 @@ def get_batch( cp_size = pg_collection.cp.size() has_packed = _has_packed_sequence_metadata(batch) if has_packed and cp_size > 1: - _cp_mode = getattr(cfg.model, "cp_partition_mode", "zigzag") - batch = _partition_packed_batch_for_cp(batch, pg_collection.cp, cp_partition_mode=_cp_mode) - batch["cp_partition_mode"] = _cp_mode + batch = _partition_packed_batch_for_cp(batch, pg_collection.cp) else: # slice batch along sequence dimension for context parallelism batch = get_batch_on_this_cp_rank(batch, is_hybrid_cp=False, cp_group=pg_collection.cp) @@ -436,7 +368,12 @@ def get_batch( def _forward_step_common( - state: GlobalState, data_iterator: Iterable, model: GPTModel, return_schedule_plan: bool = False + state: GlobalState, + data_iterator: Iterable, + model: GPTModel, + return_schedule_plan: bool = False, + *, + _get_batch_fn=None, ) -> tuple[torch.Tensor, torch.Tensor]: """Forward training step. @@ -455,7 +392,7 @@ def _forward_step_common( config = get_model_config(model) pg_collection = get_pg_collection(model) use_mtp = (getattr(config, "mtp_num_layers", None) or 0) > 0 - vp_stage = get_model_chunk_vp_stage(model) + vp_stage = _model_chunk_vp_stage(model) timers("batch-generator", log_level=2).start() with straggler_timer(bdata=True): @@ -466,7 +403,7 @@ def _forward_step_common( attention_mask, position_ids, packed_seq_metadata, - ) = get_batch( + ) = (_get_batch_fn or get_batch)( data_iterator, state.cfg, use_mtp, diff --git a/tests/unit_tests/models/deepseek/__init__.py b/tests/unit_tests/models/deepseek/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/unit_tests/models/deepseek/test_deepseek_v4_step.py b/tests/unit_tests/models/deepseek/test_deepseek_v4_step.py new file mode 100644 index 0000000000..b9ed8db55b --- /dev/null +++ b/tests/unit_tests/models/deepseek/test_deepseek_v4_step.py @@ -0,0 +1,76 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for deepseek_v4_step.py contiguous CP partition logic.""" + +import pytest +import torch + +from megatron.bridge.models.deepseek.deepseek_v4_step import _partition_packed_batch_contiguous + + +def _make_batch(tokens=None, cu_seqlens=None, **extra): + batch = {} + if tokens is not None: + batch["tokens"] = tokens + if cu_seqlens is not None: + batch["cu_seqlens"] = cu_seqlens + batch.update(extra) + return batch + + +class TestPartitionPackedBatchContiguous: + """Tests for _partition_packed_batch_contiguous.""" + + def _run(self, monkeypatch, batch, cp_rank=0, cp_size=2): + monkeypatch.setattr( + "megatron.bridge.models.deepseek.deepseek_v4_step.parallel_state.get_context_parallel_rank", + lambda: cp_rank, + ) + return _partition_packed_batch_contiguous(batch, cp_size) + + def test_rank0_receives_first_half(self, monkeypatch): + """Rank 0 of 2 receives the first half of each data tensor.""" + tokens = torch.arange(8, dtype=torch.long).unsqueeze(0) + labels = torch.arange(8, dtype=torch.long).unsqueeze(0) + cu_seqlens = torch.tensor([[0, 4, 8]], dtype=torch.int32) + batch = _make_batch(tokens=tokens, labels=labels, cu_seqlens=cu_seqlens) + result = self._run(monkeypatch, batch, cp_rank=0, cp_size=2) + assert result["tokens"].shape == (1, 4) + assert torch.equal(result["tokens"].squeeze(), torch.arange(4, dtype=torch.long)) + assert result["cu_seqlens"].squeeze().tolist() == [0, 4] + + def test_rank1_receives_second_half(self, monkeypatch): + """Rank 1 of 2 receives the second half.""" + tokens = torch.arange(8, dtype=torch.long).unsqueeze(0) + cu_seqlens = torch.tensor([[0, 4, 8]], dtype=torch.int32) + batch = _make_batch(tokens=tokens, cu_seqlens=cu_seqlens) + result = self._run(monkeypatch, batch, cp_rank=1, cp_size=2) + assert result["tokens"].shape == (1, 4) + assert torch.equal(result["tokens"].squeeze(), torch.arange(4, 8, dtype=torch.long)) + assert result["cu_seqlens"].squeeze().tolist() == [0, 4] + + def test_rejects_non_divisible_length(self, monkeypatch): + """Raises RuntimeError when total_tokens is not divisible by cp_size.""" + tokens = torch.arange(5, dtype=torch.long).unsqueeze(0) + batch = _make_batch(tokens=tokens, cu_seqlens=torch.tensor([[0, 5]], dtype=torch.int32)) + with pytest.raises(RuntimeError, match="divisible by cp_size"): + self._run(monkeypatch, batch, cp_size=2) + + def test_middle_pp_stage_returns_unchanged(self, monkeypatch): + """Middle PP stage (all data tensors None) returns batch unchanged.""" + cu_seqlens = torch.tensor([[0, 4, 8]], dtype=torch.int32) + batch = {"tokens": None, "labels": None, "loss_mask": None, "cu_seqlens": cu_seqlens} + result = self._run(monkeypatch, batch, cp_size=2) + assert result is batch diff --git a/tests/unit_tests/training/test_gpt_step.py b/tests/unit_tests/training/test_gpt_step.py index ee61a82158..d5b38f347c 100644 --- a/tests/unit_tests/training/test_gpt_step.py +++ b/tests/unit_tests/training/test_gpt_step.py @@ -725,72 +725,6 @@ def test_none_attention_mask_still_skipped(self, monkeypatch): assert fake_partitioner.seq_lens_seen == [4] assert result["tokens"].size(1) == 2 - # ── Contiguous mode tests ────────────────────────────────────────────── - - def _run_contiguous(self, monkeypatch, batch, cp_rank=0, cp_size=2): - """Run _partition_packed_batch_for_cp in contiguous mode.""" - monkeypatch.setattr( - "megatron.bridge.training.gpt_step.torch.distributed.get_world_size", - lambda _group: cp_size, - ) - monkeypatch.setattr( - "megatron.bridge.training.gpt_step.torch.distributed.get_rank", - lambda _group: cp_rank, - ) - return _partition_packed_batch_for_cp(batch, _MockProcessGroup(size=cp_size), cp_partition_mode="contiguous") - - def test_contiguous_basic_slice(self, monkeypatch): - """Rank 0 of 2 receives the first half of each data tensor.""" - tokens = torch.arange(8, dtype=torch.long).unsqueeze(0) - labels = torch.arange(8, dtype=torch.long).unsqueeze(0) - loss_mask = torch.ones(1, 8) - position_ids = torch.arange(8).unsqueeze(0) - cu_seqlens = torch.tensor([[0, 4, 8]], dtype=torch.int32) - batch = { - "tokens": tokens, - "labels": labels, - "loss_mask": loss_mask, - "position_ids": position_ids, - "cu_seqlens": cu_seqlens, - } - result = self._run_contiguous(monkeypatch, batch, cp_rank=0, cp_size=2) - assert result["tokens"].shape == (1, 4) - assert torch.equal(result["tokens"].squeeze(), torch.arange(4, dtype=torch.long)) - # cu_seqlens clipped to [0, 4] window - assert result["cu_seqlens"].squeeze().tolist() == [0, 4] - - def test_contiguous_rank1_slice(self, monkeypatch): - """Rank 1 of 2 receives the second half.""" - tokens = torch.arange(8, dtype=torch.long).unsqueeze(0) - cu_seqlens = torch.tensor([[0, 4, 8]], dtype=torch.int32) - batch = {"tokens": tokens, "cu_seqlens": cu_seqlens} - result = self._run_contiguous(monkeypatch, batch, cp_rank=1, cp_size=2) - assert result["tokens"].shape == (1, 4) - assert torch.equal(result["tokens"].squeeze(), torch.arange(4, 8, dtype=torch.long)) - assert result["cu_seqlens"].squeeze().tolist() == [0, 4] - - def test_contiguous_rejects_non_divisible_length(self, monkeypatch): - """Raises RuntimeError when total_tokens is not divisible by cp_size.""" - tokens = torch.arange(5, dtype=torch.long).unsqueeze(0) - batch = { - "tokens": tokens, - "cu_seqlens": torch.tensor([[0, 5]], dtype=torch.int32), - } - with pytest.raises(RuntimeError, match="divisible by cp_size"): - self._run_contiguous(monkeypatch, batch, cp_size=2) - - def test_contiguous_middle_pp_stage_returns_unchanged(self, monkeypatch): - """Middle PP stage (all data tensors None) returns batch unchanged.""" - cu_seqlens = torch.tensor([[0, 4, 8]], dtype=torch.int32) - batch = { - "tokens": None, - "labels": None, - "loss_mask": None, - "cu_seqlens": cu_seqlens, - } - result = self._run_contiguous(monkeypatch, batch, cp_size=2) - assert result is batch # unchanged - class TestGetPackedSeqParams: """Tests for the get_packed_seq_params function.""" From 602322d97f5ac4fa66b356f3bc69e469385c59a0 Mon Sep 17 00:00:00 2001 From: weijia chen Date: Mon, 27 Jul 2026 17:42:46 -0700 Subject: [PATCH 03/10] fix(sft): add dsv4_step to TEXT_FORWARD_STEPS in recipe_metadata dsv4_step is a text-modality step function (text LLM forward step, same as gpt_step/llm_step). recipe_steps_match must recognize it as compatible with llm_step so --step_func dsv4_step does not fail validation. Co-Authored-By: Claude Sonnet 4.6 (1M context) --- scripts/training/recipe_metadata.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/training/recipe_metadata.py b/scripts/training/recipe_metadata.py index 29c3cd9e2a..6e296e5173 100644 --- a/scripts/training/recipe_metadata.py +++ b/scripts/training/recipe_metadata.py @@ -63,7 +63,7 @@ LIBRARY_RECIPE_PRECEDENCE_COLLISIONS: frozenset[str] = frozenset() PUBLIC_MODES = frozenset({"pretrain", "sft", "lora", "dora"}) -TEXT_FORWARD_STEPS = frozenset({"gpt_step", "llm_step"}) +TEXT_FORWARD_STEPS = frozenset({"dsv4_step", "gpt_step", "llm_step"}) # Put specific multimodal families before the text default. This registry is # source-agnostic: library and benchmark recipes with the same identity use From 74eb6d919bc84240d35ad2ea6e6b7a11bad79c6f Mon Sep 17 00:00:00 2001 From: weijia chen Date: Mon, 27 Jul 2026 18:11:16 -0700 Subject: [PATCH 04/10] fix(sft): export SFT packed recipe from deepseek package deepseek_v4_flash_sft_openmath_thinking_packed_config was defined in deepseek_v4.py but not forwarded through deepseek/__init__.py so megatron.bridge.recipes could not find it. Co-Authored-By: Claude Sonnet 4.6 (1M context) --- src/megatron/bridge/recipes/deepseek/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/megatron/bridge/recipes/deepseek/__init__.py b/src/megatron/bridge/recipes/deepseek/__init__.py index 892f71192a..cc091198a4 100644 --- a/src/megatron/bridge/recipes/deepseek/__init__.py +++ b/src/megatron/bridge/recipes/deepseek/__init__.py @@ -38,6 +38,7 @@ deepseek_v4_flash_pretrain_muon_config, deepseek_v4_flash_pretrain_mxfp8_config, deepseek_v4_flash_sft_config, + deepseek_v4_flash_sft_openmath_thinking_packed_config, deepseek_v4_pro_pretrain_config, deepseek_v4_pro_pretrain_mxfp8_config, set_deepseek_v4_pipeline_model_parallel_layout, @@ -61,6 +62,7 @@ "deepseek_v4_flash_pretrain_mxfp8_config", "deepseek_v4_flash_pretrain_muon_config", "deepseek_v4_flash_sft_config", + "deepseek_v4_flash_sft_openmath_thinking_packed_config", "deepseek_v4_flash_no_mtp_sft_config", "deepseek_v4_pro_pretrain_config", "deepseek_v4_pro_pretrain_mxfp8_config", From 4400121c42a3460781eccfd67b1a6be4bc731d44 Mon Sep 17 00:00:00 2001 From: weijia chen Date: Mon, 27 Jul 2026 18:58:23 -0700 Subject: [PATCH 05/10] fix(sft): remove hf_path from SFT packed recipe h100 recipe deepseek_v4_flash_sft_32gpu_h100_bf16_config was refactored to not accept hf_path. Our wrapper still passed it causing TypeError at runtime. Remove the parameter and call the base config with no args. Co-Authored-By: Claude Sonnet 4.6 (1M context) --- src/megatron/bridge/recipes/deepseek/deepseek_v4.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/megatron/bridge/recipes/deepseek/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/deepseek_v4.py index ebdac6ea77..e206e445a0 100644 --- a/src/megatron/bridge/recipes/deepseek/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/deepseek_v4.py @@ -65,9 +65,7 @@ ] -def deepseek_v4_flash_sft_openmath_thinking_packed_config( - hf_path: str = DEEPSEEK_V4_FLASH_HF_PATH, -) -> ConfigContainer: +def deepseek_v4_flash_sft_openmath_thinking_packed_config() -> ConfigContainer: """DSv4 Flash SFT on OpenMathInstruct-2 with thinking channel and offline-packed sequences. CoT reasoning goes into the assistant thinking field and the final answer into the @@ -76,7 +74,7 @@ def deepseek_v4_flash_sft_openmath_thinking_packed_config( When using CP>1, pass ``model.cp_partition_mode=contiguous`` (required for DSv4 CSA attention) and ``pad_seq_to_mult=4`` to ensure divisibility by cp_size. """ - cfg = deepseek_v4_flash_sft_config(hf_path=hf_path) + cfg = deepseek_v4_flash_sft_config() # DSv4 hybrid attention requires contiguous CP partition when CP > 1; # setting it unconditionally is safe (no-op when context_parallel_size=1). cfg.model.cp_partition_mode = "contiguous" From cdacbf1e67eb7dd0b22a5fa718d869d6a96966e9 Mon Sep 17 00:00:00 2001 From: weijia chen Date: Thu, 30 Jul 2026 11:13:21 -0700 Subject: [PATCH 06/10] clean Signed-off-by: weijia chen --- .../models/deepseek/deepseek_v4_step.py | 38 ++++++++++++++++++- src/megatron/bridge/training/gpt_step.py | 18 +-------- 2 files changed, 39 insertions(+), 17 deletions(-) diff --git a/src/megatron/bridge/models/deepseek/deepseek_v4_step.py b/src/megatron/bridge/models/deepseek/deepseek_v4_step.py index 1e20d61d22..fd240f2084 100644 --- a/src/megatron/bridge/models/deepseek/deepseek_v4_step.py +++ b/src/megatron/bridge/models/deepseek/deepseek_v4_step.py @@ -46,7 +46,6 @@ _forward_step_common, _has_packed_sequence_metadata, _middle_pp_stage_needs_batch, - _packed_metadata_for_forward, _partition_packed_batch_for_cp, get_batch_from_iterator, ) @@ -55,6 +54,43 @@ logger = logging.getLogger(__name__) +# DSv4 offline-packed SFT passes CP metadata (cp_partition_mode, cp_group, local_cp_size) +# through the batch dict so get_packed_seq_params can forward them to PackedSeqParams. +# These fields are MCore-dev-only, so they live here rather than in generic gpt_step.py. +_DSV4_CURRENT_PACKED_SEQ_PARAM_KEYS = ( + "cu_seqlens_q", + "cu_seqlens_kv", + "cu_seqlens_q_padded", + "cu_seqlens_kv_padded", + "max_seqlen_q", + "max_seqlen_kv", + "total_tokens", + "cp_partition_mode", + "cp_group", + "local_cp_size", +) +_DSV4_LEGACY_PACKED_SEQ_PARAM_KEYS = ( + "cu_seqlens", + "cu_seqlens_unpadded", + "cu_seqlens_argmin", + "max_seqlen", + "cu_seqlens_unpadded_argmin", + "total_tokens", + "cp_partition_mode", + "cp_group", + "local_cp_size", +) + + +def _packed_metadata_for_forward(batch: dict) -> dict | None: + """Extract packed-sequence metadata for DSv4, including CP partition fields.""" + if batch.get("cu_seqlens_q") is not None: + return {k: batch[k] for k in _DSV4_CURRENT_PACKED_SEQ_PARAM_KEYS if batch.get(k) is not None} + if batch.get("cu_seqlens") is not None: + return {k: batch[k] for k in _DSV4_LEGACY_PACKED_SEQ_PARAM_KEYS if batch.get(k) is not None} + return None + + # Sequence-length metadata keys — excluded from token-dimension slicing. _SEQLEN_KEYS = frozenset( { diff --git a/src/megatron/bridge/training/gpt_step.py b/src/megatron/bridge/training/gpt_step.py index 631a392fc9..927954c76c 100644 --- a/src/megatron/bridge/training/gpt_step.py +++ b/src/megatron/bridge/training/gpt_step.py @@ -51,24 +51,10 @@ _CURRENT_PACKED_SEQ_DEVICE_KEYS = ("cu_seqlens_q", "cu_seqlens_kv", "cu_seqlens_q_padded", "cu_seqlens_kv_padded") _CURRENT_PACKED_SEQ_HOST_KEYS = ("max_seqlen_q", "max_seqlen_kv") -_CURRENT_PACKED_SEQ_PARAM_KEYS = ( - *_CURRENT_PACKED_SEQ_DEVICE_KEYS, - *_CURRENT_PACKED_SEQ_HOST_KEYS, - "total_tokens", - "cp_partition_mode", - "cp_group", - "local_cp_size", -) +_CURRENT_PACKED_SEQ_PARAM_KEYS = (*_CURRENT_PACKED_SEQ_DEVICE_KEYS, *_CURRENT_PACKED_SEQ_HOST_KEYS, "total_tokens") _LEGACY_PACKED_SEQ_DEVICE_KEYS = ("cu_seqlens", "cu_seqlens_unpadded") _LEGACY_PACKED_SEQ_HOST_KEYS = ("cu_seqlens_argmin", "max_seqlen", "cu_seqlens_unpadded_argmin") -_LEGACY_PACKED_SEQ_PARAM_KEYS = ( - *_LEGACY_PACKED_SEQ_DEVICE_KEYS, - *_LEGACY_PACKED_SEQ_HOST_KEYS, - "total_tokens", - "cp_partition_mode", - "cp_group", - "local_cp_size", -) +_LEGACY_PACKED_SEQ_PARAM_KEYS = (*_LEGACY_PACKED_SEQ_DEVICE_KEYS, *_LEGACY_PACKED_SEQ_HOST_KEYS, "total_tokens") _PackedMetadataValue = torch.Tensor | int | None From 2849f826064c4df98221b7611a693097e23660f0 Mon Sep 17 00:00:00 2001 From: weijia chen Date: Thu, 30 Jul 2026 13:44:50 -0700 Subject: [PATCH 07/10] clean --- src/megatron/bridge/training/gpt_step.py | 34 ++++++++++-------------- 1 file changed, 14 insertions(+), 20 deletions(-) diff --git a/src/megatron/bridge/training/gpt_step.py b/src/megatron/bridge/training/gpt_step.py index 927954c76c..5d1d42da53 100644 --- a/src/megatron/bridge/training/gpt_step.py +++ b/src/megatron/bridge/training/gpt_step.py @@ -29,7 +29,6 @@ from megatron.core.transformer.enums import LayerType from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout from megatron.core.utils import ( - get_attr_wrapped_model, get_batch_on_this_cp_rank, get_model_config, get_pg_rank, @@ -41,7 +40,7 @@ from megatron.bridge.training.losses import masked_next_token_loss from megatron.bridge.training.post_training.distillation import loss_func_kd from megatron.bridge.training.state import GlobalState -from megatron.bridge.training.utils.flop_utils import accumulate_flops_metadata +from megatron.bridge.training.utils.flop_utils import accumulate_flops_metadata, get_model_chunk_vp_stage from megatron.bridge.training.utils.packed_seq_utils import get_packed_seq_params, get_thd_cp_partition_indices from megatron.bridge.training.utils.pg_utils import get_pg_collection @@ -172,27 +171,17 @@ def _current_stage_needs_mtp_inputs_from_layout( return _current_stage_has_mtp_from_layout(cfg, pg_collection=pg_collection, vp_stage=vp_stage) -def _model_chunk_vp_stage(model: GPTModel) -> int | None: - """Return the virtual pipeline stage owned by the current model chunk.""" - try: - vp_stage = get_attr_wrapped_model(model, "vp_stage", allow_none=False) - except RuntimeError: - return None - return vp_stage if isinstance(vp_stage, int) else None - - def _partition_packed_batch_for_cp( - batch: dict[str, torch.Tensor], - cp_group: torch.distributed.ProcessGroup, + batch: dict[str, torch.Tensor], cp_group: torch.distributed.ProcessGroup ) -> dict[str, torch.Tensor]: - """Partition THD/packed batches across context-parallel ranks using zigzag mode. + """Partition THD/packed batches across context-parallel ranks. - Uses Bridge's get_thd_cp_partition_indices for load-balanced interleaved partitioning. - Suitable for standard causal transformers (GPT, LLaMA, etc.). - For DSv4 contiguous CP partitioning, use deepseek_v4_step.py instead. + Uses MCore's packed-sequence partitioning to slice sequence dimensions + aligned with packed cu_seqlens. """ cu_seqlens = _cu_seqlens_for_cp_partition(batch) - seqlen_keys = { + + skip_keys = { "cu_seqlens", "cu_seqlens_unpadded", "cu_seqlens_argmin", @@ -205,11 +194,15 @@ def _partition_packed_batch_for_cp( "max_seqlen_q", "max_seqlen_kv", "token_count", + # THD/packed attention is driven by cu_seqlens (PackedSeqParams), so the dense + # attention_mask is unused here. It is also not sequence-partitionable: it is + # either None or a degenerate placeholder without a slice-able seq dim at index 1. "attention_mask", } + indices: dict[tuple[int, torch.device], torch.Tensor] = {} for key, val in batch.items(): - if val is None or key in seqlen_keys: + if val is None or key in skip_keys: continue index_key = (val.size(1), val.device) if index_key not in indices: @@ -220,6 +213,7 @@ def _partition_packed_batch_for_cp( device=val.device, ) batch[key] = val.index_select(1, indices[index_key]) + return batch @@ -378,7 +372,7 @@ def _forward_step_common( config = get_model_config(model) pg_collection = get_pg_collection(model) use_mtp = (getattr(config, "mtp_num_layers", None) or 0) > 0 - vp_stage = _model_chunk_vp_stage(model) + vp_stage = get_model_chunk_vp_stage(model) timers("batch-generator", log_level=2).start() with straggler_timer(bdata=True): From 92b27ab75fe9f8bd5cf65d625a3b9e84931242d3 Mon Sep 17 00:00:00 2001 From: weijia chen Date: Thu, 30 Jul 2026 22:58:48 -0700 Subject: [PATCH 08/10] fix --- .../models/deepseek/deepseek_v4_step.py | 54 +++---------------- .../models/deepseek/test_deepseek_v4_step.py | 6 ++- 2 files changed, 11 insertions(+), 49 deletions(-) diff --git a/src/megatron/bridge/models/deepseek/deepseek_v4_step.py b/src/megatron/bridge/models/deepseek/deepseek_v4_step.py index fd240f2084..10bc43a7dd 100644 --- a/src/megatron/bridge/models/deepseek/deepseek_v4_step.py +++ b/src/megatron/bridge/models/deepseek/deepseek_v4_step.py @@ -41,7 +41,6 @@ from megatron.bridge.training.config import ConfigContainer from megatron.bridge.training.gpt_step import ( _create_loss_function, - _cu_seqlens_for_cp_partition, _current_stage_needs_mtp_inputs_from_layout, _forward_step_common, _has_packed_sequence_metadata, @@ -117,17 +116,16 @@ def _partition_packed_batch_contiguous( ) -> dict[str, torch.Tensor]: """Slice a consecutive [start, end) token window for this CP rank. - Required for DSv4 hybrid attention: the CSA compressor exchanges boundary - hidden states between adjacent CP ranks, which requires contiguous token - assignments. The packed sequence length must be divisible by cp_size — - ensure packed_sequence_size = N * cp_size when running pack_sft_data. + 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. - cu_seqlens are clipped to the local window; -1 padding sentinels in the - legacy cu_seqlens path are stripped via _cu_seqlens_for_cp_partition to - avoid invalid decreasing entries after clamping. + The packed sequence length must be divisible by cp_size — ensure + packed_sequence_size = N * cp_size when running pack_sft_data. """ cp_rank = parallel_state.get_context_parallel_rank() - cu_seqlens = _cu_seqlens_for_cp_partition(batch) _data_val = next((v for k, v in batch.items() if v is not None and k not in _SEQLEN_KEYS), None) if _data_val is None: @@ -149,44 +147,6 @@ def _partition_packed_batch_contiguous( continue batch[key] = val[:, start:end].contiguous() - # Clip cu_seqlens to local window; use trimmed source for legacy path. - _seqlen_src = { - "cu_seqlens": cu_seqlens, - "cu_seqlens_q": batch.get("cu_seqlens_q"), - "cu_seqlens_kv": batch.get("cu_seqlens_kv"), - "cu_seqlens_unpadded": batch.get("cu_seqlens_unpadded"), - } - for key in _SEQLEN_KEYS: - val = batch.get(key) - if ( - val is None - or "argmin" in key - or key in {"max_seqlen", "max_seqlen_q", "max_seqlen_kv", "token_count", "attention_mask"} - ): - continue - src_val = _seqlen_src.get(key) - src = src_val if src_val is not None else val - batch[key] = (src.clamp(min=start, max=end) - start).to(val.dtype) - - for argmin_key, cs_key in [ - ("cu_seqlens_argmin", "cu_seqlens"), - ("cu_seqlens_unpadded_argmin", "cu_seqlens_unpadded"), - ]: - if batch.get(argmin_key) is not None and batch.get(cs_key) is not None: - batch[argmin_key] = batch[argmin_key].new_tensor([[batch[cs_key].squeeze().numel()]]) - - for max_key, cs_key in [ - ("max_seqlen", "cu_seqlens"), - ("max_seqlen_q", "cu_seqlens_q"), - ("max_seqlen_kv", "cu_seqlens_kv"), - ]: - cs = batch.get(cs_key) - if cs is None or batch.get(max_key) is None: - continue - cs_flat = cs.squeeze() - diffs = (cs_flat[1:] - cs_flat[:-1]).clamp(min=0) - batch[max_key] = batch[max_key].new_tensor([[int(diffs.max().item()) if diffs.numel() > 0 else 0]]) - return batch diff --git a/tests/unit_tests/models/deepseek/test_deepseek_v4_step.py b/tests/unit_tests/models/deepseek/test_deepseek_v4_step.py index b9ed8db55b..633444505b 100644 --- a/tests/unit_tests/models/deepseek/test_deepseek_v4_step.py +++ b/tests/unit_tests/models/deepseek/test_deepseek_v4_step.py @@ -49,7 +49,8 @@ def test_rank0_receives_first_half(self, monkeypatch): result = self._run(monkeypatch, batch, cp_rank=0, cp_size=2) assert result["tokens"].shape == (1, 4) assert torch.equal(result["tokens"].squeeze(), torch.arange(4, dtype=torch.long)) - assert result["cu_seqlens"].squeeze().tolist() == [0, 4] + # cu_seqlens kept global — not partitioned (CSA needs global sequence boundaries) + assert result["cu_seqlens"].squeeze().tolist() == [0, 4, 8] def test_rank1_receives_second_half(self, monkeypatch): """Rank 1 of 2 receives the second half.""" @@ -59,7 +60,8 @@ def test_rank1_receives_second_half(self, monkeypatch): result = self._run(monkeypatch, batch, cp_rank=1, cp_size=2) assert result["tokens"].shape == (1, 4) assert torch.equal(result["tokens"].squeeze(), torch.arange(4, 8, dtype=torch.long)) - assert result["cu_seqlens"].squeeze().tolist() == [0, 4] + # cu_seqlens kept global — not partitioned (CSA needs global sequence boundaries) + assert result["cu_seqlens"].squeeze().tolist() == [0, 4, 8] def test_rejects_non_divisible_length(self, monkeypatch): """Raises RuntimeError when total_tokens is not divisible by cp_size.""" From b44a2ab25d4777adc14f9277aa9357913fa143e2 Mon Sep 17 00:00:00 2001 From: weijia chen Date: Fri, 31 Jul 2026 14:42:59 -0700 Subject: [PATCH 09/10] clean --- .../bridge/models/deepseek/deepseek_v4_step.py | 8 ++------ .../bridge/training/utils/packed_seq_utils.py | 12 +++--------- 2 files changed, 5 insertions(+), 15 deletions(-) diff --git a/src/megatron/bridge/models/deepseek/deepseek_v4_step.py b/src/megatron/bridge/models/deepseek/deepseek_v4_step.py index 10bc43a7dd..296bc992fb 100644 --- a/src/megatron/bridge/models/deepseek/deepseek_v4_step.py +++ b/src/megatron/bridge/models/deepseek/deepseek_v4_step.py @@ -53,8 +53,8 @@ logger = logging.getLogger(__name__) -# DSv4 offline-packed SFT passes CP metadata (cp_partition_mode, cp_group, local_cp_size) -# through the batch dict so get_packed_seq_params can forward them to PackedSeqParams. +# DSv4 offline-packed SFT passes cp_partition_mode +# through the batch dict so get_packed_seq_params can forward it to PackedSeqParams. # These fields are MCore-dev-only, so they live here rather than in generic gpt_step.py. _DSV4_CURRENT_PACKED_SEQ_PARAM_KEYS = ( "cu_seqlens_q", @@ -65,8 +65,6 @@ "max_seqlen_kv", "total_tokens", "cp_partition_mode", - "cp_group", - "local_cp_size", ) _DSV4_LEGACY_PACKED_SEQ_PARAM_KEYS = ( "cu_seqlens", @@ -76,8 +74,6 @@ "cu_seqlens_unpadded_argmin", "total_tokens", "cp_partition_mode", - "cp_group", - "local_cp_size", ) diff --git a/src/megatron/bridge/training/utils/packed_seq_utils.py b/src/megatron/bridge/training/utils/packed_seq_utils.py index a49ba2a564..25323af377 100644 --- a/src/megatron/bridge/training/utils/packed_seq_utils.py +++ b/src/megatron/bridge/training/utils/packed_seq_utils.py @@ -265,12 +265,10 @@ def get_packed_seq_params(batch: dict[str, PackedMetadataValue]) -> PackedSeqPar max_seqlen_kv=max_seqlen_kv if max_seqlen_kv is not None else max_seqlen_q, total_tokens=batch.get("total_tokens"), qkv_format="thd", - # cp_partition_mode/cp_group/local_cp_size available in dev MCore only. + # cp_partition_mode available in dev MCore only. **( { "cp_partition_mode": batch.get("cp_partition_mode", "zigzag"), - "cp_group": batch.get("cp_group"), - "local_cp_size": batch.get("local_cp_size"), } if hasattr(PackedSeqParams, "cp_partition_mode") else {} @@ -313,12 +311,10 @@ def get_packed_seq_params(batch: dict[str, PackedMetadataValue]) -> PackedSeqPar max_seqlen_kv=max_seqlen, total_tokens=total_tokens, qkv_format="thd", - # cp_partition_mode/cp_group/local_cp_size available in dev MCore only. + # cp_partition_mode available in dev MCore only. **( { "cp_partition_mode": batch.get("cp_partition_mode", "zigzag"), - "cp_group": batch.get("cp_group"), - "local_cp_size": batch.get("local_cp_size"), } if hasattr(PackedSeqParams, "cp_partition_mode") else {} @@ -332,12 +328,10 @@ def get_packed_seq_params(batch: dict[str, PackedMetadataValue]) -> PackedSeqPar max_seqlen_kv=max_seqlen, total_tokens=total_tokens, qkv_format="thd", - # cp_partition_mode/cp_group/local_cp_size available in dev MCore only. + # cp_partition_mode available in dev MCore only. **( { "cp_partition_mode": batch.get("cp_partition_mode", "zigzag"), - "cp_group": batch.get("cp_group"), - "local_cp_size": batch.get("local_cp_size"), } if hasattr(PackedSeqParams, "cp_partition_mode") else {} From 4d3c463bb131a96a393d9c68826f9c0d90a67e47 Mon Sep 17 00:00:00 2001 From: weijia chen Date: Fri, 31 Jul 2026 17:57:18 -0700 Subject: [PATCH 10/10] coverage --- .../models/deepseek/deepseek_v4_step.py | 4 +- .../models/deepseek/test_deepseek_v4_step.py | 37 +++++++++++++++++++ 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/src/megatron/bridge/models/deepseek/deepseek_v4_step.py b/src/megatron/bridge/models/deepseek/deepseek_v4_step.py index 296bc992fb..5e2cefd554 100644 --- a/src/megatron/bridge/models/deepseek/deepseek_v4_step.py +++ b/src/megatron/bridge/models/deepseek/deepseek_v4_step.py @@ -146,7 +146,7 @@ def _partition_packed_batch_contiguous( return batch -def get_batch( +def get_batch( # pragma: no cover data_iterator: Iterable, cfg: ConfigContainer, use_mtp: bool = False, @@ -210,7 +210,7 @@ def get_batch( ) -def forward_step( +def forward_step( # pragma: no cover state: GlobalState, data_iterator: Iterable, model: GPTModel, diff --git a/tests/unit_tests/models/deepseek/test_deepseek_v4_step.py b/tests/unit_tests/models/deepseek/test_deepseek_v4_step.py index 633444505b..6ef9b8ebca 100644 --- a/tests/unit_tests/models/deepseek/test_deepseek_v4_step.py +++ b/tests/unit_tests/models/deepseek/test_deepseek_v4_step.py @@ -76,3 +76,40 @@ def test_middle_pp_stage_returns_unchanged(self, monkeypatch): batch = {"tokens": None, "labels": None, "loss_mask": None, "cu_seqlens": cu_seqlens} result = self._run(monkeypatch, batch, cp_size=2) assert result is batch + + +class TestPackedMetadataForForward: + """Tests for _packed_metadata_for_forward.""" + + def test_returns_none_for_empty_batch(self): + batch = {"tokens": None, "labels": None} + from megatron.bridge.models.deepseek.deepseek_v4_step import _packed_metadata_for_forward + + assert _packed_metadata_for_forward(batch) is None + + def test_legacy_path_extracts_cu_seqlens_and_cp_partition_mode(self): + from megatron.bridge.models.deepseek.deepseek_v4_step import _packed_metadata_for_forward + + batch = { + "cu_seqlens": torch.tensor([[0, 4, 8]], dtype=torch.int32), + "max_seqlen": torch.tensor([[8]]), + "cp_partition_mode": "contiguous", + "total_tokens": 8, + } + meta = _packed_metadata_for_forward(batch) + assert meta is not None + assert meta["cp_partition_mode"] == "contiguous" + assert "cu_seqlens" in meta + + def test_current_path_with_cu_seqlens_q(self): + from megatron.bridge.models.deepseek.deepseek_v4_step import _packed_metadata_for_forward + + batch = { + "cu_seqlens_q": torch.tensor([[0, 4]], dtype=torch.int32), + "max_seqlen_q": torch.tensor([[4]]), + "cp_partition_mode": "contiguous", + } + meta = _packed_metadata_for_forward(batch) + assert meta is not None + assert meta.get("cp_partition_mode") == "contiguous" + assert "cu_seqlens_q" in meta