Bump Megatron-LM to miles-main-20260819 (latest NVIDIA dev) - #2673
Bump Megatron-LM to miles-main-20260819 (latest NVIDIA dev)#2673yueming-yuan wants to merge 30 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']).
… lora test mock Matches the dev Megatron arg rename used in model.py (use_gloo_process_groups).
…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.
New Megatron reads disable_{param,grad}_buffers_cpu_backup from
DistributedDataParallelConfig; _ParamAndGradBuffer no longer takes them
as ctor kwargs.
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude Code Review
This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.
Tip: disable this comment in your organization's Code Review settings.
2284788 to
0280147
Compare
0280147 to
8bc58b9
Compare
…rd-only run-megatron skips DDP wrap
yueming-yuan
added
the
bypass-fastfail
Rebuilds were driven by whether the PR diff touched docker paths, so every rerun of a Dockerfile-touching PR paid for an identical multi-arch build. Hash the build inputs, stamp the hash on the published tag as a label, and compare against it: a PR keeps one pr-<num> tag and rebuilds it only when the inputs actually move, when the tag is missing, or on the one-shot rebuild-ci-image label. resolve-ci-image now selects the PR image whenever that tag is current rather than only when this run built it, so reusing an image does not silently fall back to dev.
Also scope pull-requests: write to the docker-build call instead of the whole workflow; only that job removes the one-shot label.
…' into megatron-bump-20260819
New dev returns functools.partial builders for the mlp and shared-experts submodules instead of ModuleSpec, and MoELayer now passes name= to the shared-experts builder.
Megatron now rejects the combination: fake QAT swaps in straight-through weight tensors while TE's fused wgrad accumulation writes main_grad onto the original ones, so the quantized weights' gradients would be dropped.
Megatron's dsa variant now builds AbsorbedMLASelfAttention, which reads linear_kv_up_proj.weight directly; the pinned Bridge's LoRALinear has no such property. Upstream Bridge added one that returns the LoRA-effective weight.
Megatron's Mamba mixer exposes conv1d as conv1d_weight/conv1d_bias parameters; the pinned Bridge maps only the conv1d.* module names, so the weight sync skips them silently. Upstream Bridge accepts both spellings.
New Megatron renamed the MTP submodule to mtp_model_layer, so MTP expert weights missed the expert-offset rewrite and ranks broadcast differing tensor sets during weight sync, deadlocking the EP collectives.
Megatron now gates the HF tokenizer on args.trust_remote_code, which miles never set, so checkpoints shipping custom tokenizer code (Kimi) failed to load. miles already passes trust_remote_code=True everywhere it loads a checkpoint itself.
The private _keep_fp32 attribute was only honoured by the fp32 snapshot that rode in on the DSv4 megatron commit, which the base bump defers. Upstream now owns this: mark_keep_in_fp32 exempts the tensor from Float16Module's cast, so the scale survives the bf16 round trip through the weight sync.
TopKRouter.forward gained packed_seq_params upstream, which MoELayer passes positionally. Separately, the plain bridge provider never forwarded gradient_accumulation_fusion, so --no-gradient-accumulation-fusion could not reach models built through it (the LoRA bridge path already synced it).
# Conflicts: # .github/workflows/_build-pr-ci-image.yml
Megatron now resolves te_general_gemm, which it could not on TE 2.17 before the bump, and TE multiplies in the operand dtype -- so --moe-router-dtype fp32 became a bf16 multiply. The last ulp is enough for top-k to pick a different expert, and on qwen3.5 that flipped one token in layer 0 and cascaded to a 15% logprob drift. 25 model scripts ask for fp32 routing, so default it here rather than per model.
|
@codex review |
|
@claude review |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6ddbd76bb4
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
|
||
| ARG MEGATRON_REPO=radixark/Megatron-LM | ||
| ARG MEGATRON_BRANCH=miles-main | ||
| ARG MEGATRON_BRANCH=miles-main-20260819 |
There was a problem hiding this comment.
Changing only the CUDA Dockerfile leaves several automated contexts on the incompatible miles-main revision: _run-ci.yml:140 and _run-cpu-ci.yml:80 default to it when no PR-body override exists, release-branch-cut.yml:101 locks releases from it, and Dockerfile.rocm:24 still builds it directly. This commit now imports the new tokenizer module and reads renamed fields such as args.use_gloo_process_groups, so nightly/post-merge CI, release images, and ROCm images will pair the new Miles code with the pre-bump Megatron API and fail during import or initialization. Update those dependency defaults and the release lock source, or retain compatibility with the old API.
Useful? React with 👍 / 👎.
|
@yueming-yuan did you test the dist_muon path for the new branch? worthy having it |
There was a problem hiding this comment.
I reviewed this PR and didn't find any bugs. Because it's a large Megatron-LM version bump with numerics-sensitive compat adaptations spanning the training loop, LoRA, MTP, and weight quantization/conversion, a human look would still be worthwhile.
What was reviewed: the forced moe_router_use_torch_mm=True/trust_remote_code=True and new true_on_policy_mode guard in arguments.py; the enable_gloo_process_groups → use_gloo_process_groups rename across model.py/initialize.py/multi_lora_optimizer.py; the ddp_config-based LoRA buffer patch in lora_utils.py; the MTP loss-tracker (reduce_loss_in_tracker, values/loss_values) and mtp_detach_heads changes in model.py/model_provider.py; the transformer_layer→mtp_model_layer renaming applied consistently across update_weight/common.py and the fp8/mxfp8/nvfp4 quantizers; and the inkling plugin's move to partial-based ModuleSpec construction and mark_keep_in_fp32.
Extended reasoning...
Overview
This PR bumps the vendored Megatron-LM pin from a February 2026 fork point to the August 2026 NVIDIA dev HEAD, and adapts miles' megatron_utils backend and several model plugins (glm4, glm5, qwen3_5, qwen3_next, mimo, inkling) to the new upstream API surface. Substantial additional compat work landed across ~30 commits after this bug-hunting system's earlier review of this PR (which occurred very early, right after the initial squash of ported commits): moved vocab_size_with_padding import location, forced trust_remote_code=True and moe_router_use_torch_mm=True (to counter a new TE bf16 router-gemm default that increases logprob diff vs. old fp32 gemm), an explicit NotImplementedError for true_on_policy_mode on the megatron backend, a config-flag rename (enable_gloo_process_groups→use_gloo_process_groups), a ddp_config-based rewrite of the _ParamAndGradBuffer LoRA monkeypatch, MTP loss-tracker API changes (reduce_loss_in_tracker, values/loss_values key), mtp_detach_heads wiring guarded with getattr for contexts lacking the miles-specific arg, MTP submodule renaming (transformer_layer→mtp_model_layer) applied consistently in weight-sync regex and quantizer name-stripping, and an inkling plugin rewrite to the new partial-based ModuleSpec construction with a mark_keep_in_fp32 helper.
Security risks
None identified. trust_remote_code=True is now forced when building the tokenizer, which executes remote code from the tokenizer repo, but this is an intentional, disclosed compatibility change for the new Megatron tokenizer loading path and consistent with existing patterns in ML training frameworks; there is no user-controlled input reaching this path beyond the already-trusted model/checkpoint config.
Level of scrutiny
This warrants above-baseline scrutiny: it is a large dependency bump (39 files) touching numerics-sensitive code paths (router gemm precision, MTP loss aggregation, distributed optimizer buffer flags, LoRA parameter buffer patches) where subtle behavioral regressions would not necessarily surface as hard failures but as silent accuracy/training-quality drift. The PR description itself documents several compensating fixes for exactly this kind of regression (the router-gemm precision fix), which suggests the space is fragile and worth a second set of human eyes, especially since correctness here is best validated via the e2e CI numeric-diff tests referenced in the PR (some of which — Nemotron-3 Ultra, GLM5.1/5.2 LoRA, DeepSeek-V4 4-layer — are explicitly disabled pending follow-up work) rather than through static code reading alone.
Other factors
The renames and pattern updates (gloo flag, MTP layer name, quantizer stripping) were applied consistently across all call sites I checked, and the changes for run_megatron debug worker contexts correctly guard newly-introduced arg reads with getattr(..., False) defaults to avoid AttributeError in contexts lacking miles-specific args. This run's automated bug hunt reported no findings and no candidate issues were specifically ruled out, so this defer note reflects the change's inherent scope/risk rather than any unresolved concern.
|
[P2] Avoid re-scaling per-token MTP loss by the microbatch count At With |
3 participants
ci-megatron-pr: miles-main-20260819
What
Bump the Megatron-LM that miles depends on from the stale 2026-02-14 fork point to latest
NVIDIA/Megatron-LM
dev(3f49de20d, 2026-08-19). Supersedes #1466 (June bump, never merged).Scope: base bump only. DeepSeek-V4 and true-on-policy are explicitly disabled and land in
dedicated follow-up PRs (DSv4 first, then true-on-policy).
Megatron side —
radixark/Megatron-LM:miles-main-20260819nvidia/dev HEAD + 20 commits: 14 rebased from the June bump branch (
miles-main-20260622) +6 newly ported miles-main commits (#21 witness, #63 NVMe streaming, #68 fp32 MoE, #74 ROCm
dist-ckpt, #72 base64 te-precision-config, #82 LayerWise expert ownership).
Dropped as upstreamed into dev since June (verified feature-by-feature):
transformer_engine_int4_fake_qat.py, same env-var interface)DistributedDataParallelConfigfields)models/hybrid/refactor: pattern-MTP,is_mtp_layer,mtp_use_repeated_layer, hybrid MTP block specs)Miles side (this PR)
moe_use_legacy_grouped_gemm,MTP
transformer_layer→mtp_model_layerin all converters, MTP loss trackerreduce_loss_in_tracker,name=kwarg in custom attention,mtp_detach_headsadaptation,dumper anchors →
_forward_mlp_output_with_bias, LoRA buffer patch →ddp_configflags.--true-on-policy-moderaises NotImplementedError;test_model_provider_true_on_policy.pyskipped; DSv4 4-layer e2e disabled viaregister_cuda_ci(disabled=...). fsdp true-on-policy is unaffected.docker/Dockerfile:MEGATRON_BRANCH=miles-main-20260819.Follow-ups (not this PR)
--dsv4-impl {miles,megatron}dual backend + R3threshold sign-off.
miles_megatron_pluginspackaging.