feat(sft): DSv4 Flash offline-packed SFT with THD + CP by weijiac0619 · Pull Request #5126 · NVIDIA-NeMo/Megatron-Bridge

weijiac0619 · 2026-07-27T18:00:07Z

Adds offline-packed SFT support for DSv4-Flash with contiguous Context Parallel, required by the DSv4 CSA compressor which exchanges boundary hidden states between adjacent CP ranks. Extends _partition_packed_batch_for_cp with a contiguous mode that slices a consecutive token window per CP rank, clips cu_seqlens to the local window, and injects cp_partition_mode into packed sequence metadata so MCore attention receives the correct partition mode. Also adds deepseek_v4_flash_sft_openmath_thinking_packed_config following the GPT-OSS two-step workflow (pre-pack with prepare_gpt_sft_packed_data.py, then train). All changes are gated behind has_packed and cp_size > 1 so existing users are unaffected.

based on NVIDIA/Megatron-LM#6158

copy-pr-bot · 2026-07-27T18:00:10Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

claude · 2026-07-27T20:32:43Z

test line one
test line two

claude · 2026-07-27T20:33:08Z

Review Summary

The change adds a contiguous context-parallel partition mode for packed THD batches (needed for DSv4 CSA hybrid attention), threads cp_partition_mode/cp_group/local_cp_size through to PackedSeqParams when the installed MCore supports them, and exposes a new deepseek_v4_flash_sft_openmath_thinking_packed_config recipe. Overall the direction is sound and the contiguous slicing math is careful. A few points to confirm:

Correctness

gpt_step.py:245-258 — the "sentinel-free trimmed" comment only holds for the legacy cu_seqlens key. cu_seqlens_q/cu_seqlens_kv/cu_seqlens_unpadded map to raw (untrimmed) batch tensors, so any -1 sentinel there would be clamped to 0, the exact bug the comment says it avoids. Please confirm those never carry -1 padding in the contiguous path (see inline comment).
gpt_step.py:412 — cp_partition_mode is set on the batch but is only in _LEGACY_PACKED_SEQ_PARAM_KEYS, not _CURRENT_PACKED_SEQ_PARAM_KEYS. Current-format (cu_seqlens_q) batches therefore never propagate the mode to get_packed_seq_params and silently fall back to zigzag. Confirm DSv4 only uses the legacy cu_seqlens path (see inline comment).
cp_group/local_cp_size are listed in _LEGACY_PACKED_SEQ_PARAM_KEYS but are never set on the batch anywhere in gpt_step.py, so they always resolve to None in get_packed_seq_params. If MCore expects a real cp_group/local_cp_size for contiguous partitioning, this plumbing is incomplete; if the dev-MCore default is fine, consider dropping the unused keys.

Recipe

deepseek_v4.py:69 — the new function is annotated to return ConfigContainer, but ConfigContainer is not imported in this module. It only works because of the future annotations import (deferred/string annotation). Fine at runtime, but get_type_hints() / tooling would fail to resolve it; consider importing it for consistency with sibling modules.

Tests

The new contiguous branch in _partition_packed_batch_for_cp (the bulk of the diff) has no unit coverage. TestPartitionPackedBatchForCp only exercises the default zigzag path. Given the intricate argmin/max_seqlen/clip recomputation, a targeted unit test for the contiguous mode is strongly warranted (unit-preferred per CONTRIBUTING).
deepseek_v4_flash_sft_openmath_thinking_packed_config is not added to DEEPSEEK_V4_SFT_RECIPES in tests/functional_tests/test_groups/recipes/test_deepseek_recipes_finetune.py, so the new recipe is never built/smoke-tested.

Suggested test cases

tests/unit_tests/training/test_gpt_step.py::TestPartitionPackedBatchForCp — add a contiguous-mode case asserting the local slice, clipped cu_seqlens, and recomputed cu_seqlens_argmin / max_seqlen.
tests/unit_tests/training/test_gpt_step.py::TestPartitionPackedBatchForCp — add a contiguous case that fails divisibility (total_tokens not divisible by cp_size) and asserts the RuntimeError.
tests/unit_tests/training/test_gpt_step.py::TestPartitionPackedBatchForCp — add a contiguous middle-PP-stage case (only seqlen keys, no data tensors) asserting the batch is returned unchanged.
tests/unit_tests/training/utils/test_packed_seq_utils.py::TestGetPackedSeqParams — add a case passing cp_partition_mode/cp_group/local_cp_size through the legacy cu_seqlens path, asserting they reach PackedSeqParams when the field exists (and are dropped when it does not).
tests/functional_tests/test_groups/recipes/test_deepseek_recipes_finetune.py::TestDeepSeekV4FinetuneRecipes::test_deepseek_v4_sft_recipes — add deepseek_v4_flash_sft_openmath_thinking_packed_config to DEEPSEEK_V4_SFT_RECIPES.

No perf/recipe performance configs are touched (scripts/performance/configs/ has no DeepSeek entries) — no perf tests impacted.

claude · 2026-07-27T20:33:28Z

+        _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


The comment says the sentinel-free trimmed version is used "to avoid clamping -1 padding to 0", but that only holds for the legacy cu_seqlens key — _SEQLEN_MAP["cu_seqlens"] points at the trimmed cu_seqlens local, whereas cu_seqlens_q, cu_seqlens_kv, and cu_seqlens_unpadded map to the raw batch.get(...) tensors, which are not trimmed. If any of those carry -1 sentinels, src.clamp(min=start) promotes -1 to start and the subsequent - start yields 0, i.e. the exact failure the comment claims to avoid. Please confirm the current-format cu_seqlens_q/kv batches never contain -1 padding in the contiguous path, or trim them like cu_seqlens.

Good catch. For DSv4 offline SFT (the only contiguous-CP user today), batches always take the legacy cu_seqlens path so cu_seqlens_q/kv/unpadded are None and never processed. However the comment was misleading — updated it to clarify that cu_seqlens_q/kv/unpadded in the current format do not use -1 sentinels, so they are safe to clip directly without trimming.

claude · 2026-07-27T20:33:29Z

-        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


batch["cp_partition_mode"] is set to a plain str here, and it flows into _packed_metadata_for_forward via _LEGACY_PACKED_SEQ_PARAM_KEYS → get_packed_seq_params. But cp_partition_mode is only added to _LEGACY_PACKED_SEQ_PARAM_KEYS, not _CURRENT_PACKED_SEQ_PARAM_KEYS. For current-format batches (those with cu_seqlens_q), the cp_partition_mode / cp_group / local_cp_size set on the batch will therefore never reach get_packed_seq_params, so the hasattr(PackedSeqParams, "cp_partition_mode") block in the cu_seqlens_q branch always falls back to "zigzag". If DSv4 contiguous CP is expected to run through the current-format path, this is a silent correctness gap. If DSv4 only uses the legacy cu_seqlens path, please confirm.

Fixed — added cp_partition_mode, cp_group, local_cp_size to _CURRENT_PACKED_SEQ_PARAM_KEYS as well (a6cc865). For DSv4 offline SFT the legacy path is always taken, but this ensures correctness for current-format batches too.

cuichenx · 2026-07-27T20:53:48Z

+    cfg.dataset = default_openmathinstruct2_thinking_config(
+        seq_length=cfg.model.seq_length,
+        enable_offline_packing=True,
+        pad_seq_to_mult=4,  # 2 * cp_size when CP=2; divisible by any CP up to 4


this should be dynamic to cfg.model.context_parallel_size

good catch. forgot to set that after validation. resolving now

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) <noreply@anthropic.com>

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) <noreply@anthropic.com>

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) <noreply@anthropic.com>

cuichenx · 2026-07-31T21:32:42Z

/ok to test 92b27ab

weijiac0619 · 2026-07-31T21:45:18Z

/ok to test b44a2ab

weijiac0619 · 2026-08-01T00:58:56Z

/ok to test 4d3c463

weijiac0619 force-pushed the weijia/dsv4-sft-offline-thd-cp branch 7 times, most recently from 3138f2a to fb283af Compare July 27, 2026 20:20

weijiac0619 marked this pull request as ready for review July 27, 2026 20:25

weijiac0619 requested a review from cuichenx July 27, 2026 20:27

claude Bot reviewed Jul 27, 2026

View reviewed changes

weijiac0619 force-pushed the weijia/dsv4-sft-offline-thd-cp branch 2 times, most recently from a6cc865 to 29e64cb Compare July 27, 2026 20:44

weijiac0619 changed the title ~~feat(sft): DSv4 Flash offline-packed SFT with THD + CP=2~~ feat(sft): DSv4 Flash offline-packed SFT with THD + C Jul 27, 2026

cuichenx reviewed Jul 27, 2026

View reviewed changes

weijiac0619 force-pushed the weijia/dsv4-sft-offline-thd-cp branch from 29e64cb to 2588a87 Compare July 27, 2026 20:56

feat(sft): DSv4 Flash offline-packed SFT with THD + CP=2

1dcbc1b

weijiac0619 force-pushed the weijia/dsv4-sft-offline-thd-cp branch from 2588a87 to 1dcbc1b Compare July 27, 2026 21:01

weijiac0619 changed the title ~~feat(sft): DSv4 Flash offline-packed SFT with THD + C~~ feat(sft): DSv4 Flash offline-packed SFT with THD + CP Jul 27, 2026

weijiac0619 and others added 3 commits July 27, 2026 17:35

cuichenx added high-priority and removed needs-more-tests Requires additional L0 and L1 test coverage before merge full-test-suite labels Jul 31, 2026

cuichenx previously approved these changes Jul 31, 2026

View reviewed changes

cuichenx added ready-to-merge PR is approved, current, and only waiting for CI to pass before merge and removed needs-review PR is ready for code review and waiting on a reviewer labels Jul 31, 2026

copy-pr-bot Bot temporarily deployed to public July 31, 2026 21:33 Inactive

copy-pr-bot Bot had a problem deploying to test July 31, 2026 21:33 Error

clean

b44a2ab

weijiac0619 dismissed cuichenx’s stale review via b44a2ab July 31, 2026 21:44

copy-pr-bot Bot temporarily deployed to test July 31, 2026 21:46 Inactive

copy-pr-bot Bot temporarily deployed to public July 31, 2026 21:46 Inactive

copy-pr-bot Bot temporarily deployed to public July 31, 2026 22:03 Inactive

copy-pr-bot Bot temporarily deployed to public July 31, 2026 22:33 Inactive

coverage

4d3c463

copy-pr-bot Bot temporarily deployed to public August 1, 2026 00:59 Inactive

copy-pr-bot Bot temporarily deployed to test August 1, 2026 00:59 Inactive

copy-pr-bot Bot temporarily deployed to public August 1, 2026 01:16 Inactive

copy-pr-bot Bot temporarily deployed to public August 1, 2026 01:17 Inactive

copy-pr-bot Bot temporarily deployed to public August 1, 2026 01:49 Inactive

yaoyu-33 added the needs-more-tests Requires additional L0 and L1 test coverage before merge label Aug 1, 2026

cuichenx approved these changes Aug 2, 2026

View reviewed changes

cuichenx merged commit d033420 into main Aug 2, 2026
84 checks passed

cuichenx deleted the weijia/dsv4-sft-offline-thd-cp branch August 2, 2026 18:19

Uh oh!

Conversation

weijiac0619 commented Jul 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

copy-pr-bot Bot commented Jul 27, 2026

Uh oh!

claude Bot commented Jul 27, 2026

Uh oh!

claude Bot commented Jul 27, 2026

Review Summary

Uh oh!

claude Bot Jul 27, 2026

Choose a reason for hiding this comment

Uh oh!

weijiac0619 Jul 27, 2026

Choose a reason for hiding this comment

Uh oh!

claude Bot Jul 27, 2026

Choose a reason for hiding this comment

Uh oh!

weijiac0619 Jul 27, 2026

Choose a reason for hiding this comment

Uh oh!

cuichenx Jul 27, 2026

Choose a reason for hiding this comment

Uh oh!

weijiac0619 Jul 27, 2026

Choose a reason for hiding this comment

Uh oh!

cuichenx commented Jul 31, 2026

Uh oh!

weijiac0619 commented Jul 31, 2026

Uh oh!

weijiac0619 commented Aug 1, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

weijiac0619 commented Jul 27, 2026 •

edited

Loading