Add fused pre-gated-delta-rule (GDN conv fusion) - #14
Closed
wplf wants to merge 35 commits into
Closed
Conversation
Adds a standalone VLM training playground under ``examples/multimodal_dev/`` with Qwen3.5-VL end-to-end. Highlights - Model-agnostic entry point (``pretrain_multimodal.py``) with a ``MODEL_REGISTRY`` so adding a new architecture is just a registry entry plus a backing module. - Qwen3.5-VL model: vision encoder, MRoPE, decoder, factory, specs, configurations covering proxy / 9B / 397B-A17B variants. - Datasets: mock data and CORD-V2 VLM dataset, with THD pack/pad in the collate function. - THD + CP support consolidated in ``forward_step.py`` and the model layer (uses MRoPE THD pre-computation and ``cu_seqlens_q_padded`` CP partitioning). - Run script + README, plus tests for MRoPE parity, CP correctness, CP support, and THD correctness / e2e. Also gates the torch DataLoader vanilla-collate path on the new ``use_vanilla_collate_fn`` arg (one-line change to ``megatron/training/datasets/data_samplers.py``) so CORD-V2 works under BSHD. Functional dependency: the new model arch sets ``mrope_interleaved=True`` in its config and relies on the core MRoPE interleaved layout introduced in a separate PR. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-Authored-By: BestJuly <19769279+BestJuly@users.noreply.github.com>
… preprocessing
Fixes 8 issues in vlm_dataset.py found by review against Megatron-Bridge's
qwen2_5_collate_fn reference implementation.
- loss_mask off-by-one (Bug 1): the previous mask was built on input_ids
while labels were shifted, dropping the image->text supervision signal
at the boundary. Now masks structural tokens on the shifted labels and
also shifts loss_mask itself left by 1.
- missing SFT prompt masking (Bug 2): user-turn and chat-template tokens
were trained on. Now uses backward substring token search (mirroring
create_multiturn_loss_mask_by_search) to unmask only the assistant
answer span.
- seq_length not enforced (Bug 3): long CORD-V2 samples could overflow.
Now end-truncates input_ids in __getitem__ with a warning.
- unsafe pad_token_id fallback (Bug 4): falling back to 0 silently masked
a real vocab token. Now falls back to EOS and raises if neither is set.
- silent image_token_id miss (Bug 6): fallback could return None, causing
dataset / model disagreement. Now raises ValueError.
- stale docstrings (Bug 8): updated Qwen2.5-VL / --image-size references
to Qwen3.5-VL / --total-seq-length.
- narrow skipped_tokens set (Bug 14): vision_start/end, im_start/end,
video_pad, endoftext were not masked on labels. Now uses
tok.all_special_ids union {pad_id, image_token_id}.
- lost Qwen-VL dynamic resolution (Bugs 15/17/19): fixed-square resize
removed; conversation content carries the image object;
qwen_vl_utils.process_vision_info extracts images; processor is called
with min_pixels / max_pixels.
- pixel_values bf16 conversion (Bug 18): moved from forward_step into the
dataset so per-step dtype checks become no-ops.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- raise --manual-gc-interval 5 → 50 to cut GC pause frequency on long runs. - enable --moe-permute-fusion and --moe-router-fusion in the MoE branch (no-op for dense variants since MOE_ARGS is gated on NUM_EXPERTS>0). - enable grad-accumulation fusion under FSDP by dropping --no-gradient-accumulation-fusion from FSDP_ARGS. - add --log-timers-to-tensorboard and --log-params-norm to surface timer breakdown and parameter L2 norm in TB/wandb. - drop the hardcoded CKPT_LOAD path from the in-script example invocations so the comment reflects from-scratch CP correctness runs. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Update the 'Copyright (c) 2025, NVIDIA CORPORATION' line to 2026 across all newly-added Python files under examples/multimodal_dev/ for the Qwen3.5-VL training example. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…1 vs CP=4 correctness - test_thd_e2e.py: rewrite TestPackBatch -> TestPackOrPadBatchPacked/Padded using the per-sample dict input shape produced by the dataset; drop attention_mask / position_ids / MRoPE / user cu_seqlens cases that are no longer the helper's concern. Add TestPackOrPadBatchDivisibleBy4 covering per-sample alignment when cp_size=2 forces divisible_by=4 (via monkeypatched mpu). - test_thd_correctness.py: swap _pack_batch(batch_dict) for pack_or_pad_batch(per-sample list, use_packed_sequence=True); compute THD position_ids locally since the helper no longer carries them. - test_cp_thd_correctness.py (new): single-torchrun script comparing CP=1 and CP=4 in one process via destroy + re-initialize model_parallel, with weights pinned by a state_dict snapshot. Uses MultimodalModel with a stub vision encoder (vision branch skipped via pixel_values=None); loss aggregated by AllReduce-SUM of (num, den) on the CP group; grad_norm aggregated by AllReduce-SUM of gradients on the CP group then dividing by cp_size, so each rank holds the CP-mean gradient (equivalent to CP=1's backward on the full-batch mean loss). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…ess + small cleanups
- forward_step.py:pack_or_pad_batch — previously crashed for TP>1 because
the per-sample loop dereferenced data on non-source ranks (which receive
data=None from get_batch), and cu_seqlens/max_seqlen used to build
PackedSeqParams were local Python lists never broadcast. Now: gate the
build loop on TP rank 0, broadcast cu_seqlens / cu_seqlens_padded as
part of the data dict, and derive max_seqlen / total_tokens from the
(broadcast) cu_seqlens on every rank — no extra collective.
- models/base.py — add public MultimodalModel.cp_split_loss_mask
(staticmethod) so the post-forward loss path doesn't need to import the
module's private _cp_split_tensor / _thd_cp_partition_index.
forward_step.py uses it instead of duplicating the slicing logic.
- forward_step.py — replace bare `except Exception` around get_args() with
`getattr(get_args(), 'sequence_parallel', False)` + AssertionError-only
fallback (matches what megatron's get_args() actually raises when args
are uninitialised in tests). Strip three WHAT/TODO comments that
narrated intent rather than explaining a non-obvious why.
- tests/_helpers.py (new) — shared grad_norm / mean_loss helpers.
test_thd_correctness.py uses them in place of its local copies.
test_cp_thd_correctness.py keeps its CP-aware _global_loss /
_global_grad_norm (genuinely different — they add AllReduce on the CP
group); the duplication noted in review was overstated.
- tests/test_cp_thd_correctness.py — drop _StubVisionEncoder's dummy
nn.Linear(1,1); MegatronModule does not need a parameter for state_dict
round-tripping (verified by re-running the CP=1 vs CP=4 suite — numbers
identical to the previous commit).
Verified locally:
- test_thd_e2e.py: 20/20 passed
- test_thd_correctness.py: ALL PASSED (BSHD vs THD equal-length parity)
- test_cp_thd_correctness.py: ALL PASSED (CP=1 vs CP=4, same numbers as
previous commit ec6d2d3: BSHD/THD loss + grad_norm)
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…(pylint 10/10)
Run `bash tools/autoformat.sh` toolchain (black --skip-magic-trailing-comma
--skip-string-normalization, isort, ruff check, pylint, mypy) directly on
the changed files (autoformat.sh only scans megatron/core + tests/, so
these don't normally go through the gate):
- black: reformat 6 files for line length / wrapping
- pylint: drop unused imports (test_thd_correctness.py: parallel_state,
_build_packed_seq_params; test_cp_thd_correctness.py: parallel_state);
add docstrings to 9 small functions / methods (test methods,
_NoCPGroup.size/rank, _StubVisionEncoder.__init__/forward, main()
entrypoints); add module-level `# pylint: disable=bad-builtin` to the
two stdout-reporting standalone scripts (test_thd_correctness.py,
test_cp_thd_correctness.py) where the many `print()`s are intentional.
- mypy: replace implicit Optional defaults — `seq_length: int = None` →
`Optional[int] = None` in pack_or_pad_batch; same for `mrope_section`
and `mtp_block_spec` in MultimodalModel.__init__; tighten
`get_batch(data_iterator: Iterator[Dict[str, Any]])` to
`Iterator[list[Dict[str, Any]]]` so the call to pack_or_pad_batch
type-checks.
Remaining mypy diagnostic — `transformer_engine.pytorch` missing
library-stub marker — is repo-wide (also flagged on `megatron/core/` files
in the main run) and tolerated because autoformat.sh runs mypy with
`|| true`.
Verified locally:
- test_thd_e2e.py: 20/20 passed
- test_thd_correctness.py: ALL PASSED
- test_cp_thd_correctness.py: ALL PASSED (CP=1 vs CP=4, numbers
identical to ec6d2d3 / 4332813)
- pylint score: 10.00/10 (was 9.48/10 → 9.95/10 → 10.00/10)
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
One-line follow-up to 0211665 — `_helpers.py` was missed in the preceding bulk lint commit's `git add`. Black removes the spaces around `**` (`total ** 0.5` -> `total**0.5`). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
… overlap profiling Wrap the 6 execution entry points in SharedExpertMLP with nvtx_range_push/pop so the dedicated shared-expert stream shows up as named slices in nsys profiles. Useful for inspecting where shared-expert FFN sits relative to the hybrid-ep dispatch A2A when --moe-shared-expert-overlap is on.
The Blackwell-FSDP path in `validate_args` auto-appends 'dp_cp' (and
'ep_dp' when EP>1) to `args.high_priority_stream_groups` whenever
`use_megatron_fsdp` / `use_torch_fsdp2` is set on `sm_10x` devices.
Combined with `--fake-process-group`, this makes `initialize_model_parallel`
call `create_group(..., pg_options=get_nccl_options(...))` with a
`ProcessGroupNCCL.Options` object for those groups. PyTorch's
`FakeProcessGroup._create_internal` rejects non-fake options with:
TypeError: _create_internal(): incompatible function arguments.
Invoked with: 0, 64, <ProcessGroupNCCL.Options ...>
Short-circuit `get_nccl_options` to return `None` whenever the default
process group is the fake backend, so the fake sub-groups are created
without NCCL-specific options. The change is a no-op for real backends.
…nly path - MultimodalModel.build_schedule_plan: run vision encoder + embedding scatter eagerly (main path), then delegate decoder-layer schedule plan to inner GPTModel. Vision encoder intentionally NOT part of A2A overlap. - multimodal forward_step: handle return_schedule_plan=True by calling model.build_schedule_plan(...). - combined_1f1b: relax GPTModel isinstance assertion to also accept MultimodalModel (decoder-only A2A overlap).
torch.cuda.memory._snapshot() at dump time gave only segment state — empty frames in blocks and 0 events in device_traces — because Megatron never called torch.cuda.memory._record_memory_history() to start the recorder. Enable the recorder at pretrain() time (before model+optimizer init), gated on the existing --record-memory-history flag. Mode='all' captures both segment state AND the full allocation/free event timeline with python stacks, so a downstream mem-profile peak summary / pytorch memory_viz can attribute each tensor to its allocation site.
…aunch path The fused THD dispatch (rope_utils.apply_rotary_pos_emb -> fused_apply_mrope_thd) calls the kernel directly and only validated total seqlen % cp_size, not each packed sub-sequence. The unfused per-sequence check in _get_thd_cp_splits() is bypassed on the fused path, so for CP>1 variable-length packing where the total is divisible but an individual sub-sequence is not, the kernel would silently compute wrong local->global CP token indices (global_start // cp_size). Add the per-sequence guard in get_fused_mrope_thd_unavailable_reason (only on the cp_size>1 path); when it triggers, the dispatch falls back to the unfused path.
…ivisibility - Add fwd/bwd parity at the real deployment shape head_dim=256, rotary_dim=64 (rotary_percent=0.25, 75% pass-through) with mrope_section=[11,11,10], plus the non-interleaved and full-rotary variants, for both BSHD and THD. The existing parametrized tests only covered head_dim=16/20 with rotary_dim=16 (~80% rotated). - Add a regression test that get_fused_mrope_thd_unavailable_reason rejects a packed batch whose total length is CP-divisible but an individual sub-sequence is not (and accepts the all-divisible control).
Fused MRoPE
…n-overlap) Switch shared_experts selective recompute from a standard checkpoint (keeps the output) to CheckpointWithoutOutput: discard the shared-expert output in the forward and regenerate it in backward from a grad hook. The recompute is registered AFTER any pre_mlp_layernorm recompute (so the shared expert's input is restored first) and before its backward — via the moe node's expert_output in the A2A-overlap fine-grained callables, or via mlp_output_with_bias[0] in TransformerLayer._forward_post_mlp on the single-call path. Split CheckpointWithoutOutput.discard_output_and_register_recompute into discard_output() + register_recompute_hook() for the overlap path (which frees the output and registers the hook in different callables); disabled under MoE cudagraph partial capture. Validated on Qwen3.5-VL 397B proxy (8xGB200, EP=8, mbs=4): both paths run with stable loss / no NaN; controlled memory 85.94 -> 84.75 GB (-1.19 GB); 3-way numerical parity within the run-to-run FP-noise floor. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- pin checkout to 12605c5 (first correct + NaN-free DHU backward) - use pinned submodules (warn against --remote → old buggy kernel) - add nested-cutlass symlink fix + 4th fused extension build - add portable-wheel build + force-reinstall flow - document the fused dv_dhu (DV_DHU=1) wheel-import caveat + .so copy workaround - add DHU-backward kernel version history (3a371f5→297386a→5661cfa→12605c5) - add 397B-proxy 20-step validation (grad norm finite, matches FLA) Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Recompute the GatedDeltaNet QKV projection + preparation block (in_proj -> CP all-to-all -> conv1d -> _prepare_qkv -> g/beta) as a discard-output checkpoint, selected via recompute_modules="gdn_qkv". The block outputs (query/key/value/ g/beta/gate) are discarded in the forward and regenerated from a grad hook in the backward, freeing the large GDN QKV-prep activations. When gdn_norm_out recompute is also enabled, the two discard-output checkpoints have a forward-order data dependency (the QKV block output `gate` feeds the gated-norm block), so both are registered to a single CheckpointManager that replays their recompute in forward order (qkv -> norm_out) from one unified grad hook on the layer output. Adds "gdn_qkv" to the selective-recompute allowed_modules in TransformerConfig.
The memory-opt test lists "shared_experts" in recompute_modules but never set moe_shared_expert_intermediate_size, so the model had no shared experts and the shared-experts discard-output recompute path was never exercised. Set shared_expert_intermediate_size=512 so the recompute actually fires (in both the overlap and non-overlap paths compared by the test).
Adds test_selective_recompute_gdn_qkv (mirrors test_selective_recompute_norm_out) to TestGatedDeltaNet: builds a no-recompute baseline GatedDeltaNet and one with recompute_modules=["gdn_qkv"], runs forward+backward, and asserts the output, all parameter grads and the input grad match bit-for-bit. Verifies the QKV projection+prep discard-output recompute is numerically exact.
Add gdn_qkv whole-block recompute for GatedDeltaNet
Support discard-output recompute for MoE shared experts under A2A overlap
Add optional mcore GDN optimized wrapper
Squash of yuzhongw/gdn_conv_fusion (Yuzhong Wang, 11 commits b3c1c7e..05bd2da) onto qwen35-vl-central-dev. - megatron/core/fusions/fused_pre_gated_delta_rule.py: streamed fused pre-GDR forward/backward kernels (fused_streamed_pre_gated_delta_rule). - megatron/core/fusions/fused_mega_pre_gated_delta_rule.py: mega-fused pre-GDR forward/backward kernels (fused_mega_pre_gated_delta_rule). - transformer_config.py / arguments.py: --pre-gated-delta-rule-impl {unfused, fused_streamed, fused_mega} (auto-generated from config). - tests/unit_tests/ssm/test_gated_delta_net.py: packed fused GDN path tests. Integration note: central-dev already restructured the GDN forward for gdn_qkv discard-output recompute, so rather than the source branch's forward rewrite the fused dispatch is wired into _compute_qkv_for_gated_delta_rule: after the CP all-to-all it calls the fused wrapper (returning (q,k,v,gate,beta,g)) when pre_gated_delta_rule_impl selects a fused path, reordered to the method's (q,k,v,g,beta,gate) layout; the unfused path and the gdn_qkv/gdn_norm_out recompute structure are unchanged. Co-authored-by: Yuzhong Wang <yuzhongw@nvidia.com>
5 tasks
wplf
force-pushed
the
jinliangl/qwen35-vl-central-dev
branch
from
June 23, 2026 07:45
79f6060 to
ea678f6
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Integrates Yuzhong Wang's fused pre-gated-delta-rule (GDN conv fusion) onto
jinliangl/qwen35-vl-central-dev.Squash of the 11 GDN-fusion commits from
yuzhongw-nvidia:yuzhongw/gdn_conv_fusion(b3c1c7e..05bd2da), authored by Yuzhong Wang, without the unrelated upstreamdevsyncs that branch also carried.What it adds
megatron/core/fusions/fused_pre_gated_delta_rule.py— streamed fused pre-GDR fwd/bwd kernels.megatron/core/fusions/fused_mega_pre_gated_delta_rule.py— mega-fused pre-GDR fwd/bwd kernels.transformer_config.py/arguments.py—--pre-gated-delta-rule-impl {unfused, fused_streamed, fused_mega}.tests/unit_tests/ssm/test_gated_delta_net.py— fused-vs-unfused parity tests.Integration notes (differs from the source branch)
central-dev already restructured the GDN forward for gdn_qkv discard-output recompute, while the source branch restructured the same forward for the fused kernels — the two are not textually mergeable. So instead of taking the source branch's forward rewrite, the fused dispatch is wired into the existing
_compute_qkv_for_gated_delta_rule:pre_gated_delta_rule_implto_fused_streamed/_fused_mega(or the unfusedpre_gated_delta_rule), reordering the fused(q,k,v,gate,beta,g)to this method's(q,k,v,g,beta,gate)layout;pre_gated_delta_rule(qkvzba, ...);The external
causal_conv1dimport is also made fully optional (guarded) so importing GatedDeltaNet never fails when the package is absent; the fused backward raises a clear error only if actually invoked without it.Testing (oci-hsg GB200, container has
causal-conv1dinstalled)TestFusedPreGatedDeltaRule+ config tests: 14 passed (fused == unfused for forward / THD / padding / pre_gated_delta_rule / packed fwd+bwd / conv boundary).test_selective_recompute_gdn_qkv+test_selective_recompute_norm_out: 12 passed (gdn_qkv / gdn_norm_out unaffected across TP/SP/CP).🤖 Generated with Claude Code