Bump Megatron-LM to miles-main-20260622 (latest NVIDIA dev) - #1466
Open
yueming-yuan wants to merge 13 commits into
Open
Bump Megatron-LM to miles-main-20260622 (latest NVIDIA dev)#1466yueming-yuan wants to merge 13 commits into
yueming-yuan wants to merge 13 commits into
Conversation
Megatron dev reimplemented PR #6's MTP-in-RL support natively: - process_mtp_loss derives MTP labels from input_ids when labels is None (RL). - config.mtp_detach_heads detaches output head + MTP embedding gradients. So on the miles side: set config.mtp_detach_heads=True when enable_mtp_training, and stop passing the now-unsupported mtp_kwargs to GPTModel.forward (labels=None + input_ids derivation is equivalent to mtp_labels=batch['tokens']).
Run miles' custom dsv4 attention on the bumped Megatron (miles-main-20260622): - rename plugin/script config to dev-native names: csa_window_size, csa_compress_ratios, csa_compress_rotary_base, o_groups, o_lora_rank, moe_n_hash_layers (mbridge). Keep dsv4_hc_* (mHC precision-branched). Script uses --csa-compress-ratios "[..]" string form; drop --no-activation-func-clamp-shared-expert. - dev API drift: tokenizer _vocab_size_with_padding moved to megatron.core.tokenizers; enable_gloo_process_groups -> use_gloo_process_groups; DeepSeekV4Attention accepts name=. - run_megatron worker: build forward-only model without DDP (wrap_with_ddp=run_backward). Validated against the old-megatron baseline (run-megatron, 4-layer prune): mHC+attention numerically exact, logprob mean within 8e-4.
yueming-yuan
requested review from
Shi-Dong,
Zhichenzzz,
fzyzcjy,
guapisolo,
jybsuper,
maocheng23 and
yushengsu-thu
as code owners
June 23, 2026 04:39
Contributor
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
yueming-yuan
added
run-ci-image
run-ci-megatron
and removed
run-ci-megatron
run-ci-image
labels
… on new Megatron convert_hf_to_torch_dist auto-forces PP=world_size for >1 GPU; bumped Megatron asserts hash-MoE layers + PP>1 need an explicit pipeline_model_parallel_layout. The 4-layer prune converts fine at PP1 on 1 GPU (validated). Full Flash/Pro use explicit multi-PP convert configs (their PP>1 hash-MoE convert layout is a separate follow-up).
… lora test mock Matches the dev Megatron arg rename used in model.py (use_gloo_process_groups).
Mirror GLM5: register indexer_replay_manager on V4Indexer (stream_idx=layer_id) and route the indexer top-k through get_topk_fn (flattened to [n_tokens, n_kv]). No-op/transparent when the manager is disabled; DSv4 forward verified unchanged (loss 19.92). Full record/replay-overlap validation needs an RL run with --use-rollout-indexer-replay.
…ive=dsv4_hybrid) get_dsv4_spec picks the dsv4 attention module by env var: 'miles' (default) keeps miles' DeepSeekV4Attention (BSHD + sparse-CP + tilelang, validated); 'native' swaps in dev's dsv4_hybrid sparse attention. Forces miles when TP>1 or CP>1 (dev dsv4_hybrid asserts TP==1, no sparse-CP). Default path verified unchanged (loss 19.92); dev dsv4_hybrid validated via test_dsv4_hybrid_native_parity (unfused passes on H200/SM9; fused needs SM10+). Native end-to-end with a checkpoint additionally needs mbridge dev-name mapping + dsv4_hybrid config coordination (opt-in follow-up; miles is the production path).
…put_with_bias
New dev split the MLP into _forward_mlp_output_with_bias + _forward_post_mlp and dropped
_sglang_pre_mlp_residual (true-on-policy was deferred). Anchor pre_mlp_residual on the
unconditional nvtx_range_push("mlp") (residual set in all branches); keep pre_mlp_layernorm_output
and mlp_output_with_bias anchors. Will re-verify after the true-on-policy migration.
…bias The megatron true-on-policy merge relocated the _sglang_pre_mlp_residual extraction out of _forward_mlp into _forward_mlp_output_with_bias, so the BSHD source-patcher anchor 'residual = getattr(self, "_sglang_pre_mlp_residual", hidden_states)' no longer exists. Mirror the THD config: target _forward_mlp_output_with_bias and prepend the pre_mlp_residual dump at the mlp nvtx push (matches the new structure). Fixes test_run_megatron + test_miles_dumper source-patch failures.
…_module) New dev's build_module passes name= to every spec-built module (megatron base Attention.__init__ accepts it). miles' custom attention classes overriding __init__ must accept + thread it: DSAMultiLatentAttention/DSAMLASelfAttention (glm5, DeepSeek-V3.2), HuggingfaceAttention base, qwen3_5/qwen3_next Attention. Mirrors the DeepSeekV4Attention name fix. Fixes 'DSAMLASelfAttention.__init__() got an unexpected keyword argument name'.
…lders New dev removed the moe_use_legacy_grouped_gemm argument and the legacy grouped-GEMM concept (grouped_mlp_modules now takes only moe_use_grouped_gemm). miles' model_provider and glm4 spec builders read args.moe_use_legacy_grouped_gemm, which no longer exists -> AttributeError on the Namespace. Drop the kwarg (dev spec fns take **kwargs / the new grouped path is the only one).
New Megatron renamed MultiTokenPredictionLayerSubmodules.transformer_layer ->
mtp_model_layer, so MTP weight names are now mtp.layers.{i}.mtp_model_layer.*.
miles' MTP weight converters (mcore<->hf) hardcoded 'transformer_layer' ->
'Unknown MTP component' on rollout weight export (test_mimo_7B_mtp_only_grad).
Accept both names across export (mimo/qwen3_next/deepseekv3/glm4moe/qwen3_5 + fp8/
mxfp8/nvfp4 quantizers) and mbridge import (mimo/qwen3_next/qwen3_5/glm4moe_lite/glm4moe).
New dev rewrote MTPLossLoggingHelper: it stores loss_sums/num_tokens (or loss_values) and only computes tracker['values'] after reduce_loss_in_tracker() (which also does the cross-rank all-reduce). miles read tracker['values'] directly (never populated) -> 'values' not in tracker -> mtp_losses unbound -> UnboundLocalError at extra_metrics. Call reduce_loss_in_tracker(), read values/loss_values, init mtp_losses=None + guard the metric. Fixes test_mimo_7B_mtp_only_grad past weight-export.
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
1 participant
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.
Rebases miles' custom Megatron-LM changes onto the latest NVIDIA/Megatron
dev.Megatron side (
radixark/Megatron-LM:miles-main-20260622, offnvidia/dev)dsv4): re-applied PR Tiny add Sample.group_index #28's mHC 4-stream-residual slice (config + transformer_block/layer + p2p/schedules + mappings + Float16Module fp32 snapshot), reconciled with dev's evolved layer structure and its nativedsv4_hybrid. Native infra (hash routing, sqrtsoftplus, dsa, input_ids) reused.Miles side (this PR)
docker/Dockerfile:MEGATRON_BRANCH->miles-main-20260622.csa_*,o_*,moe_n_hash_layers); keptdsv4_hc_*(mHC precision-branched).enable_->use_gloo_process_groups,DeepSeekV4Attention(name=).Validation
DSv4 4-layer prune via run-megatron vs old-megatron baseline: mHC + attention numerically exact (attn_output / 4-stream residual rel-L2 = 0); logprob mean within 8e-4 (residual diffs confined to dev-native MoE nondeterminism).