fix(megatron): ROCm-safe attention_backend + Mamba/SFT E2E, dead-code & CI cleanup#870
Merged
Merged
Conversation
…narrowing - Runtime summary now reads step timing live from the Actions API instead of hand-wrapped start/end timers, so untracked stages (e.g. a ~19min git clone) show up automatically instead of staying invisible. - coverage_summary.py's classify() now separates "counts toward a total" from "gets its own row": a bare file with no sub-package of its own (e.g. core/base_module.py) folds into its group's total without faking a sub-package row. junit_summary's Time column now matches runtime's h/m/s. - select_tests.py no longer narrows unit tests (the full suite is only ~5min, so narrowing risked under-selection for little gain), and now treats a bare examples/<file> (e.g. run_pretrain.sh, which MaxText's E2E runs directly) as blast-radius-unknown instead of silently ignoring it.
Neither patch runs in any other case in this file: SDMA param all-gather only activates behind ENABLE_SDMA_ALLGATHER=1 (env, not a --arg) + a distributed optimizer, and fused residual+RMSNorm needs use_turbo_rms_norm + PRIMUS_FUSED_RESIDUAL_NORM_V2. Confirmed on real MI300X hardware with coverage instrumentation that both patches install and actually execute (not just get registered and skipped), at ~60s added runtime. Also looked at two bigger targets and ruled both out for now, for different reasons: - Muon optimizer: third_party/Emerging-Optimizers is already a populated submodule (just needs PYTHONPATH, not a new dependency as first assumed), and the vendor import chain works. But coverage shows the dispatch wrapper in muon_optimizer_patches.py installs and then never actually runs during training, so --optimizer muon isn't exercising the muon path yet -- needs tracing why megatron.training.training's optimizer setup isn't calling through the patched reference before this is worth an E2E case. - Legacy ZeroBubble PP scheduler (core/pipeline_parallel/zerobubble/): has genuinely bit-rotted since being disabled -- it now throws "SchedNodeRuntime.prepare() got an unexpected keyword argument 'force_all_reduce'", a real incompatibility with the current megatron-core version that needs its own fix, not a config tweak.
The fast-exit path only flushed stdout/stderr before os._exit(0), which skips atexit and silently drops any process coverage data that hadn't been saved yet. Route through the existing flush_before_hard_exit() (handles both the output flush and coverage.save()) instead of a partial inline reimplementation.
…m-safe Primus patches Megatron's parse_args to return a namespace built straight from YAML/CLI, bypassing argparse's type coercion. Enum-typed args such as attention_backend therefore arrived as plain strings and lost every downstream `== AttnBackend.member` comparison in LanguageModule._set_attention_backend(), so `--attention_backend fused` (and friends) were silently no-ops. - MegatronArgBuilder now coerces enum-typed overrides through Megatron's own argparse `type` converter (reusing the parser's definitions, no hand-kept mapping), so attention_backend resolves to the real AttnBackend enum. - Add a before_train patch (ROCm-gated) that reconciles attention_backend with the image's baked NVTE_*_ATTN. Stock megatron asserts the vars are unset-or-equal to what the backend wants and crashes on ROCm images that intentionally bake NVTE_FLASH_ATTN=0 (TE flash attn is unavailable on ROCm). The patch instead lets the selected backend win: `auto` enables only the backends the platform hasn't disabled (respecting the baked FLASH=0), and an explicit backend forces exactly its combination, overriding the baked defaults. Unit-tested the coercion (mocked and against the real Megatron parser) and the patched reconciliation.
Verified on a16-19 (MI300X x8). - test_mamba_370M / test_zebra_llama_1B_hybrid: cover the MambaStack and HybridStack code paths, which (unlike the GPT configs) go through megatron-core's own attention probe on the default `auto` backend -- now ROCm-safe via the attention_backend patch, so no NVTE workaround is needed. - test_qwen2_sft_lora: first E2E case for the posttrain suite (MegatronSFTTrainer) and primus/backends/megatron/peft/*.py. Uses a dedicated tiny config (reusing sft_trainer.yaml + qwen2.5_7B.yaml, pointing tokenizer_model at a real-but-tiny Qwen2 checkpoint so the runner hook's HF->Megatron conversion matches the model shape) plus a local jsonl SFT fixture, so the test is fast and offline. Pins the unfused backend via --attention_backend (its head_dim=2 is below the fused CK kernel minimum). - 01_convert_checkpoints.py: a general posttrain fix, not test-only -- the conversion hook runs in a separate subprocess against Megatron-Bridge's bundled Megatron-LM, which never sees the attention_backend patch. Since the conversion computes no attention, neutralize (unset + restore) the NVTE_*_ATTN vars around AutoBridge.import_ckpt() to get past stock megatron's `auto` assertion on ROCm's baked NVTE_FLASH_ATTN=0.
…entation Vendored Megatron-LM's setup_model_and_optimizer() dispatches to its own get_megatron_muon_optimizer whenever "muon" is in config.optimizer, before Primus's patched get_megatron_optimizer ever runs. muon_optimizer_patches.py only patches the latter, so it can never be reached, and the backported moun.py / moun_optimizer_config.py / layer_wise_optimizer.py behind it is unreachable too. Remove the patch, its unit test, and that implementation. --optimizer muon continues to work via Megatron's native path. Keep the remaining muon_* fields in trainer_base.yaml (drop only muon_weight_decay / muon_weight_decay_method, which nothing reads): Megatron-LM's native get_megatron_optimizer_config() opportunistically copies any args attr matching an AdamOptimizerConfig field name, and reads these; since most have no CLI flag, this yaml is the only way to tune Muon's internals via Primus config. Also drop the now-stale muon references in the patch-explorer skill doc and the coverage summary notes.
…h_wgrad Superseded by te_wgrad_store.py's WeightGradStore-based approach (introduced in the same wgrad-split migration that should have retired these but didn't). No file imports either module or the classes they define (_LinearWithWGradSplit / _GroupedLinearWithWGradSplit), and there is no dedicated unit test for them.
Cut PR churn: run the GitHub Actions version-update check monthly instead of weekly, and ignore minor/patch bumps so only major version updates open a PR. Dependabot security updates are unaffected and still land regardless.
WangLingxun
force-pushed
the
cicd/coverage-runtime-e2e-fixes
branch
from
July 13, 2026 09:13
7d72911 to
0d7f5b1
Compare
Wrap an argparse add_argument() call that exceeded the 110-char line limit, so `pre-commit run black --all-files` (the code-lint CI job) passes. Formatting only; no behavior change.
WangLingxun
marked this pull request as ready for review
July 13, 2026 09:40
WangLingxun
requested review from
Xiaoming-AMD,
limou102 and
wenxie-amd
as code owners
July 13, 2026 09:40
Add a first-class Mamba2 130M pretrain example (model config + MI300X/MI355X example yamls) and an E2E case wiring it into the megatron suite -- the only E2E covering the megatron_bridge backend (trainer, load_recipe_config, MegatronBridgeArgBuilder, bridge patches). The test pins a tiny shape via extra_args so it is independent of the example yaml's sizes; seq_length is left to the yaml since the recipe feeds it to both model and dataset and Bridge asserts they match. Verified on MI300X. coverage_summary.py: rename the "(top-level)" row to "primus (top-level)" with an accurate note (loose primus/ modules folded by path depth), wrap NOTES with <br> so a long note can't stretch the column, move the tier dot to the module-name column, and drop the stale megatron_bridge "no E2E" note. Also trims the Mamba/Hybrid/SFT-LoRA E2E comments to the code-relevant rationale, and sets the new attention_backend_patches files' copyright to 2026.
WangLingxun
force-pushed
the
cicd/coverage-runtime-e2e-fixes
branch
from
July 14, 2026 07:06
2f056f4 to
cad4eef
Compare
Add a PRIORITY map (lower sorts first) so important groups lead the table regardless of coverage; ties keep coverage-descending order. core (Primus core, imported by every run) ranks highest. Ordering only: no target or gate.
WangLingxun
force-pushed
the
cicd/coverage-runtime-e2e-fixes
branch
from
July 15, 2026 05:14
dd53264 to
edc07be
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
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.
Summary
attention_backendenum args + make it ROCm-safe. Primus builds Megatron's arg namespace straight from YAML/CLI, bypassing argparse type coercion, so enum args likeattention_backendarrived as plainstrand lost every downstream== AttnBackend.membercheck —--attention_backend fusedwas a silent no-op.MegatronArgBuildernow coerces enum args through Megatron's own argparse converter. A ROCm-gatedbefore_trainpatch also lets the selected backend win over the image's bakedNVTE_*_ATTN(ROCm images bakeNVTE_FLASH_ATTN=0, which makes stock megatron'sauto/explicit-backend assertion crash):autorespects the bakedFLASH=0, explicit backends force their combination.test_mamba_370M/test_zebra_llama_1B_hybridcover the MambaStack/HybridStackautoattention path;test_qwen2_sft_lorais the first posttrain (MegatronSFTTrainer) +peft/*E2E, using a tiny offline config/fixture. The conversion hook (01_convert_checkpoints.py) neutralizesNVTE_*_ATTNaroundAutoBridge.import_ckpt()since it runs against Bridge's bundled megatron that never sees the patch.te_gemm_patch_wgrad/te_group_gemm_patch_wgrad.