don't merge - #5284
Draft
wplf wants to merge 8 commits into
Draft
Conversation
… 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. (cherry picked from commit dc835c2)
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.
(cherry picked from commit b6a0063)
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. (cherry picked from commit 25cc5a3)
Squash of: Support discard-output recompute for MoE shared experts (overlap + non-overlap) + configure shared experts in test_fsdp_1f1b_memory_opt. Adds output-discarding selective recompute for the MoE shared-expert path (both A2A-overlap and non-overlap), wiring through fine_grained_callables, moe_layer, transformer_layer and the RNG-state handling in tensor_parallel/random.py. Includes shared-expert config in the FSDP 1F1B overlap memory test.
Squash of the fused-mRoPE work (Add Qwen3.5 MRoPE fusion benchmark support; Fix THD mRoPE CP fallback consistency; mRoPE THD review cleanup; enforce per-sequence CP divisibility on the fused THD launch path; unit-test coverage for real Qwen3.5-VL shapes). Adds a fused mRoPE kernel (megatron/core/fusions/fused_mrope.py) with an is_fused_mrope_available() gate, raw-mrope-freqs plumbing through rope_utils / rotary_pos_embedding / gpt_model / attention, the transformer_config + arguments toggles, and tests/unit_tests/fusions/ test_fused_mrope.py. Core only: the examples/multimodal_dev integration is intentionally dropped because that example is already upstream (NVIDIA#4751) and has diverged from this branch's copy. Co-Authored-By: Li Tao <litao@nvidia.com>
Squash/port of: Add fused pre-gated-delta-rule (GDN conv fusion). Adds fused_streamed / fused_mega pre-gated-delta-rule kernels (megatron/core/fusions/fused_pre_gated_delta_rule.py, fused_mega_pre_gated_delta_rule.py) and a pre_gated_delta_rule_impl config/arg (unfused|fused_streamed|fused_mega). Refactors the unfused conv1d -> _prepare_qkv -> g/beta block into GatedDeltaNet.pre_gated_delta_rule and dispatches fused vs unfused inside _compute_qkv_for_gated_delta_rule (stacks on the P7 recompute extraction). Port note: cp_group/cp_size are threaded into pre_gated_delta_rule from the caller (instead of relying on forward-local vars) to fit dev's resolve_cp_group structure.
Squash/port of: Add optional mcore GDN optimized wrapper; update install docs; Integrate GDN optimized kernels; bump mcore_gdn_opt submodule 9121702 -> 12605c5; fix GDN-opt test dispatch-validation call. Adds the third_party/mcore_gdn_opt submodule and an opt-in import path (MCORE_GDN_USE_OPT_WRAPPER=1) that swaps fla's chunk_gated_delta_rule for the optimized wrapper, with fla fallback. Includes the CUDA-opt benchmark/tests and the reproduction doc. WARNING: the mcore_gdn_opt SM100 backward (DHU+WY) kernels produce incorrect gradients — safe for inference, NOT safe for training. The optimized path is gated behind MCORE_GDN_USE_OPT_WRAPPER and defaults off. Co-Authored-By: bhsueh <byshiue@users.noreply.github.com>
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
2 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.
This branch rebases net-new work from
jinliangl/gdn-conv-fusion-and-optonto currentdev, one squashed commit per feature. The original branch was 241 commits behinddev; the Qwen3.5-VL multimodal_dev example it carried is already upstream (#4751) and is intentionally excluded.Commits (8)
feat(moe): add NVTX ranges to SharedExpertMLP— shared-expert overlap profilingfix(parallel_state): skip NCCL.Options under --fake-process-groupfeat(memory_snapshot): enable event recording (frames + device_traces)feat(moe): discard-output recompute for shared experts(overlap + non-overlap)feat(fusions): fused mRoPE for Qwen3.5-VL (core)— kernel + core integration only; the multimodal_dev example integration is dropped (already upstream, diverged)GDN selective recompute (norm_out + qkv discard-output)feat(ssm): fused pre-gated-delta-rule (GDN conv fusion)feat(ssm): optional mcore_gdn_opt optimized kernels for GatedDeltaNetGatedDeltaNet.forward(recompute hooks +cp_group/cp_sizethreading). Structure/AST verified, but numerical correctness is NOT GPU-tested —tests/unit_tests/ssm/test_gated_delta_net.pymust be run on GPU.mcore_gdn_opt(submodule) SM100 backward kernels produce incorrect gradients — safe for inference, NOT safe for training. Gated behindMCORE_GDN_USE_OPT_WRAPPER=1, default off.isortnot yet run;.gitignorecarries a couple of incidental entries (AGENTS.md,*.pyc) that should be trimmed before any real submission.