Repository navigation
refactor(runner): retire the examples launchers in favour of primus-cli - #999
Merged
Merged
Conversation
4 tasks done
wenxie-amd
pushed a commit
that referenced
this pull request
Aug 21, 2026
## Summary Documentation-only. Rewrites every guide, example README and config header comment that still told users to `export EXP` and run `examples/run_*.sh`, so the whole repository describes launches in terms of `runner/primus-cli`. **No behaviour changes and no scripts deleted here.** `examples/run_pretrain.sh` and friends still exist on this branch and still work; this PR only stops the documentation from presenting them as the way to launch training. Retiring the scripts, migrating their downstream callers, and the fixes that came out of that migration are in the stacked follow-up **[#999](#999, which is based on this branch. - **Launch snippets rewritten** to the three runner modes — `primus-cli direct`, `primus-cli container`, `primus-cli slurm srun` — across `docs/`, `examples/`, and the header comments of the diffusion and SFT configs. Container environment options stay at the container layer rather than leaking past the `--` into the Primus CLI. - **`examples/README.md` gains a mode table** explaining when to use `direct` / `container` / `slurm`, the `--` argument shape, and where the environment defaults and setup hooks live (`runner/helpers/envs/`, `runner/helpers/hooks/`, `runner/.primus.yaml`). - **References repointed at the code that owns the behaviour** instead of at the launcher script: the HipBLASLt tuning master switch and its stage variable at `runner/helpers/hooks/train/pretrain/prepare_experiment.sh`, deterministic mode at `runner/helpers/hooks/05_deterministic.sh` and `runner/helpers/envs/base_env.sh`, and ANP / `NCCL_NET_PLUGIN` selection at `runner/helpers/hooks/03_enable_ainic.sh`. The HipBLASLt section previously omitted that the stage is ignored unless `PRIMUS_HIPBLASLT_TUNING=1`; that is now stated where users hit it. - **New guard test `tests/runner/test_docs_cli_examples.py`** fails when a Markdown command invokes a `.sh` that is not in the tree, so the next launcher move cannot leave stale instructions behind. It matches only `bash` / `sh` / `source` invocations, so prose that merely names a script stays out of scope. A second assertion pins the hybrid-models guide's `--volume` / `--env DATA_PATH` placement, which is the shape that is easy to get wrong when translating a snippet. ## Scope boundary Documentation whose statements only become true once the launchers are actually gone is deliberately **not** in this PR — it lives in #999 alongside the code that makes it true: - the Kubernetes section removal, which belongs with the deletion of `examples/run_k8s_pretrain.sh`; - the multi-node benchmark commands that name `runner/helpers/launch/slurm_pretrain.sh`, a helper #999 creates; - the auto-benchmark and ODC READMEs, whose scripts #999 repoints at `primus-cli`; - the data-preparation guide's `TOKENIZED_*_DATA_PATH` / split-seed / rank-coordination sections, which describe prepare-hook behaviour #999 restores. The rule applied throughout: this branch must be self-consistent on its own. Nothing here references a file that does not exist yet, and nothing contradicts the tree as it stands. ## Test plan - [x] `pre-commit run --from-ref origin/main --to-ref HEAD` — clean. - [x] `tests/runner/run_all_tests.sh` — 10/10 suites. - [x] `pytest tests/runner/` — 2 passed, including `test_docs_cli_examples.py`, which confirms every documented script invocation resolves against this branch's tree (the old `examples/run_*.sh` scripts are still present here and still referenced where they have not been rewritten). - [x] Stacking verified: this branch plus #999 reproduces the pre-split tree byte for byte (identical tree hash), so splitting lost nothing.
WangLingxun
changed the base branch from
refactor/primus-cli-launchers
to
main
September 8, 2026 07:19
The examples/ launch scripts had grown into a second, divergent copy of the launcher: they re-implemented the cluster env setup, cache layout, NCCL tuning and hook dispatch that runner/primus-cli already owns, and the parallel run_*_cli.sh variants meant every fix had to be made twice. Keeping the two copies in sync was a recurring source of drift, and the K8s submission client that justified keeping run_pretrain.sh alive is no longer needed. - Move the behaviour that only lived under examples/ into the runner (the persistent kernel/JIT cache layout, TRAIN_LOG honoured by direct mode) and delete the entry scripts and networking helpers outright, so there is one launcher interface to learn and one implementation to maintain. - Turn the srun + container assembly that the packaged launchers still share into an explicit helper under runner/helpers/launch/ rather than a user-facing entry point. The packaged customer/MoE launchers and the benchmark runners were relying on the EXP/NNODES contract for real value, so they call that helper instead of open-coding the same srun flags nine times; everything that only needed a single-node launch calls primus-cli directly. - Preserve the packaged-launcher contracts while consolidating: Slurm GPU counts remain scheduler-derived unless explicitly overridden, UCCL-EP and clean-container controls cross the container boundary, and PRIMUS_SKIP_PIP consistently skips per-run dependency hooks. - Drop the backend prepare.py copies under examples/: they were only ever invoked by the deleted dispatchers, and the hooks the CLI actually runs own that logic now. The two behaviours the megatron copy still had over its hook are restored in the hook itself in a follow-up. - Repoint the docstrings, hook headers and CI comments that named examples/run_pretrain.sh at primus-cli or at the runner hook that owns the behaviour, since they now describe a file that does not exist. - Remove the Kubernetes documentation together with the client it describes: the examples/README.md section for run_k8s_pretrain.sh and the deployment guide's Kubernetes chapter, which now states that no submission client is shipped and points at the distributed-variable mapping instead. Keeping the script deletion and its documentation apart would have left the tree claiming a capability it no longer has. - Point the remaining multi-node benchmark commands and the auto-benchmark and ODC docs at runner/helpers/launch/slurm_pretrain.sh and primus-cli, which is only correct once this commit creates the helper and rewrites the callers. - Probe the persistent cache root before exporting it. The /workspace default is only writable inside the container, so bare-metal direct runs, which now source the same base_env.sh, fall back to a per-user cache with a warning instead of failing at the first Triton JIT compile. An explicitly configured PRIMUS_CACHE_ROOT is still taken as-is. - Mount PATH_TO_BNXT_TAR_PACKAGE into the container the way the deleted launcher did. Forwarding the variable alone left the rebuild hook unable to see the tar it names, and that hook now warns rather than informs when it skips a rebuild the user explicitly asked for. - Keep the MLflow and Databricks credentials in the container passthrough allowlist. The deleted launcher forwarded them explicitly, the automatic prefix passthrough does not cover them, and the logging guide already documents them as allowlisted. - Take the launcher's exit code from PIPESTATUS[0] in the shared srun helper. tee was reporting success for failed training runs, which SLURM, CI and the benchmark runners all read as a passing job. cleanup_rs.sh matters beyond the rename: its pkill pattern matched the old script name, so it had stopped killing the process it is meant to reap. The container env passthrough list already carried HSA_KERNARG_POOL_SIZE twice before this change; the duplicate is dropped while the list is being edited here, leaving the order that the environment-variables reference documents.
Folding the examples/ launchers into primus-cli routed every run through
runner/helpers/hooks/train/pretrain/, and three preparation behaviours that
only the old dispatcher and its prepare.py had did not survive the move. Each
one fails late and quietly, so restore them where the CLI actually looks.
- Move the NeMo AutoModel prepare entry into the hook directory searched by
primus-cli and let it consume only its own arguments, so training overrides
continue to the trainer instead of being rejected by the hook parser.
- Return early for SFT, and only create a held-out validation split when
evaluation needs one and the config sets neither valid_data_path nor
test_data_path. Either one means the run already has its own evaluation
data, and the generated split would replace it.
- Advertise the generated split only through the extra.* keys the config
left unset. Those lines become CLI arguments, so emitting both keys
unconditionally overrode an explicitly configured test_data_path.
- Honor custom train/eval tokenized prefixes exactly. Key split caches by the
split seed and test size, validate their metadata, and coordinate nodes with
atomic completion/failure markers plus a bounded wait on shared storage.
- Keep exposing the resolved Megatron backend and generated split paths through
the hook extra-argument protocol used by primus-cli direct mode.
The split seed is spelled --split_seed, with a PRIMUS_SPLIT_SEED environment
entry, rather than --seed. primus-cli direct mode hands the same argument list
to the hooks and to training, and seed is already a Megatron training parameter
(trainer_base.yaml defaults it to 1234). Sharing one flag between two consumers
meant a --seed the user passed to reseed training also changed the split cache
key, which re-downloaded BookCorpus, re-tokenized the corpus and deleted the
tokenized files already on disk -- so a seed sweep paid for a full data
preparation per seed. The hook parser uses parse_known_args, so under the new
name --seed is left in the unknown arguments and reaches the config overrides
untouched. The default stays 42 and the on-disk payload keeps its "seed" key,
so existing bookcorpus_{train,valid}_<key>.json caches and .split.json metadata
remain valid. PRIMUS_SPLIT_SEED is the entry to prefer through primus-cli,
since it keeps the knob out of the argument list training also parses.
Document the restored contracts in the data-preparation guide (the exact
TOKENIZED_DATA_PATH / TOKENIZED_TRAIN_DATA_PATH / TOKENIZED_EVAL_DATA_PATH
prefixes, the split-cache metadata, the split seed being separate from the
Megatron training seed, and the bounded rank wait that requires shared
storage), and repoint the ODC README at the primus-cli invocation run_odc.sh
now issues.
AMD Triton cannot compile the num_stages=4 variants of flash-linear-attention's
KDA intra-chunk kernels. On rocm/primus:v26.3 (torch 2.10, ROCm 7.2, Triton
3.6.0) the first autotune sweep of a KDA run dies in the MLIR pass pipeline:
error: 'tt.load' op operation destroyed but still has uses
error: Failures have been detected while processing an MLIR pass pipeline
RuntimeError: PassManager::run failed
The failure is in AMD's software pipeliner --
tritonamdgpu-schedule-loops{num_stages=4} feeding tritonamdgpu-pipeline. On
Triton 3.7.0 and 3.7.1 the full sweep compiles and the autotuner settles on
num_stages=2. The [2, 3, 4] space is still declared in FLA 0.5.2 (the pinned
version) and in upstream main, so there is no fixed FLA release to move to.
The workaround used to be a prepare-time shell hook that shelled out to a sed
script and rewrote chunk_intra.py inside site-packages. That mutated a shared
image, needed a separate PrimusParser subprocess just to decide whether the run
used KDA at all, and left the installed package differing from its wheel. Do it
inside the training process instead: a before_train patch drops the unsafe
candidates from the two kernels' autotune space, keeping the num_warps sweep so
the autotuner still picks the best remaining config. Nothing on disk changes.
- Add megatron.fla.kda_safe_autotune. It walks Triton's decorator chain to the
object owning .configs -- @triton.heuristics wraps FLA's CachedAutotuner, so
the config list is one .fn hop below the module attribute -- and filters it.
- Also guard CachedAutotuner.maybe_load_cached_config, which writes a config
read from a best-config JSON straight into self.cache without consulting
.configs. FLA_CACHE_MODE plus FLA_CONFIG_DIR could otherwise reintroduce an
unsafe config behind the filter. The guard lowers num_stages on the entries it
finds rather than deleting them, because deleting loops under
FLA_CACHE_MODE=ALWAYS: should_check_fla_cache returns True unconditionally in
that mode, so run() reloads the JSON before every launch, and an emptied cache
sends Triton back to benchmarking only for the next reload to overwrite the
result with the unsafe config again -- a full autotune per kernel call for the
duration of the run. The other cache modes gate on `key not in self.cache` and
were never exposed to it. num_stages is adjusted in place so num_ctas,
maxnreg, pre_hook and ir_override -- which FLA populates on Triton >= 3.5.1 --
survive; rebuilding the config would silently reset them.
- Gate on ROCm, and on whether the run reaches FLA's KDA kernels at all --
use_fla_triton_kda, kda_backend and use_kimi_k3_attention_backend are all
merged onto args before before_train runs. The second condition is there to
avoid importing fla into a run that has no other use for it; trimming the
config list of a kernel that never launches would itself be harmless.
- Retire PRIMUS_FLA_TRITON_AUTOTUNE_WORKAROUND without a replacement knob,
rather than carrying a tri-state override for something this cheap.
- Fail soft everywhere: a missing FLA, a renamed kernel or a restructured
.configs logs and returns rather than blocking training.
There is deliberately no architecture gate. What breaks is a compiler pass,
which is essentially chip-independent, and gfx94* only ever described where the
failure happened to be seen first. 18071e1 records the same lesson from the
buffer-store miscompile: a gfx942-only gate was later found to skip real
kernels carrying the same pattern on gfx950, and "until someone remembers to,
that architecture is silently unprotected". kimi_k3 is KDA on MI355X (gfx950),
so that is a live configuration here rather than a hypothetical.
Over-applying it is cheap, and measured. Letting the full sweep run on gfx942
and comparing its winner against the narrowed space: zebra 300M (K=32, V=64)
and a k3 curve shape (K=V=128, B=2, T=2048, H=16) both pick num_warps=2,
num_stages=2, so narrowing is free; the k3 full shape (K=V=128, B=1, T=4096,
H=96) picks num_stages=3, where the narrowed kernel is 5.4% slower in isolation
and 0.96% slower at the KDA op level. An end-to-end A/B on 4x MI300X (60
iterations, four interleaved repeats each) stayed inside run-to-run noise, at
+0.10% and -0.83% in opposite directions. So ~1% of a single op bounds the cost
of applying this too widely, against a hard compile failure when it is missed.
chunk_kda_fwd_kernel_intra_sub_chunk is narrowed even though it is dead code
today: it is only reached with safe_gate=True, which nothing in Primus sets,
and training logs only ever show autotune records for the bwd kernel. It stays
in so the protection is already in place if safe_gate is ever enabled, with a
comment noting that it is also the one kernel where num_stages=4 actually wins
(by 4.9% at safe_gate shapes), so that call needs re-measuring at that point.
Delete this patch once the supported toolchain floor reaches Triton 3.7.
Validated on 8x MI300X with zebra_llama_300M_kda_pure: the patch narrows both
kernels from 12 to 4 configs and training runs to completion.
Latest main added MaxDiffusion quick starts through the launcher removed by this PR. Route those examples through the root primus-cli wrapper and teach the documentation guard that tools/fetch_primus.sh belongs to the external MAD checkout after `cd MAD`.
WangLingxun
force-pushed
the
refactor/retire-examples-launchers
branch
from
September 8, 2026 07:40
60883a8 to
b485d07
Compare
The ODC launcher now delegates to runner/primus-cli, so its fixture must stub that entry point rather than the removed examples/run_pretrain.sh. This keeps the environment-wiring test isolated from real training after the launcher migration.
This was referenced Sep 10, 2026
Xiaoming-AMD
marked this pull request as ready for review
September 10, 2026 03:32
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
The auto-benchmark runner currently masks primus-cli failures by forcing a zero exit code, so benchmark outcomes can be reported incorrectly.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Retires the legacy examples/run_* launcher scripts and standardizes training launches (direct/container/slurm) on runner/primus-cli, while restoring and testing previously-regressed prepare-hook behaviors and replacing the on-disk FLA Triton workaround with an in-process Megatron patch.
Changes:
- Removed
examples/run_*launchers and migrated downstream scripts/docs toprimus-cli(plus a shared SLURM helper for packaged launchers). - Restored/expanded prepare-hook behavior (split-seed decoupling, SFT short-circuit, held-out split coordination) with new regression tests.
- Added an in-process FLA KDA Triton autotune-space narrowing patch and corresponding unit tests; normalized Spur
srunflag handling viaprimus-cli slurm.
File summaries
| File | Description |
|---|---|
| tools/hybrid/launch_mamba_hybrid_300M.sh | Switch hybrid launcher to primus-cli direct. |
| tools/hybrid/launch_hybrid_faflag.sh | Switch hybrid launcher to primus-cli direct. |
| tools/hybrid/fla_order_dataset.py | Update usage example to primus-cli. |
| tools/ci/select_tests.py | Update CI path classification comment for primus-cli-based E2E. |
| tools/auto_benchmark/run_primus_autobenchmark.sh | Run benchmarks via primus-cli direct instead of removed launchers. |
| tools/auto_benchmark/Primus_Auto_Benchmark_README.md | Update autobenchmark docs to reference primus-cli. |
| tests/unit_tests/ci/test_select_tests.py | Adjust unit test to use a “bare examples file” sentinel post-launcher removal. |
| tests/unit_tests/backends/megatron/test_fla_kda_autotune_patches.py | New unit tests for the FLA KDA autotune narrowing patch. |
| tests/trainer/test_maxtext_trainer.py | Update E2E shell entry to runner/primus-cli direct. |
| tests/runner/test_requirement_hooks.py | New test ensuring PRIMUS_SKIP_PIP=1 makes requirement hooks a clean no-op. |
| tests/runner/test_primus_cli_slurm.sh | Add Spur-specific srun normalization coverage. |
| tests/runner/test_prepare_hooks.py | New tests for restored Megatron/NeMo prepare-hook behavior. |
| tests/runner/test_odc_run_launcher.sh | Update ODC launcher test stub to fake runner/primus-cli. |
| tests/runner/test_launcher_regressions.py | New regression tests for SLURM entry env precedence and shared launcher env forwarding. |
| tests/runner/test_launcher_migration.py | New migration tests for ODC status propagation and slurm node-flag parsing forms. |
| tests/runner/test_docs_cli_examples.py | Allowlist an external-repo script reference that appears in docs. |
| runner/primus-cli-slurm.sh | Normalize Spur-incompatible srun flags; improve CLI override parsing for --opt=val and -N2. |
| runner/primus-cli-slurm-entry.sh | Default GPUS_PER_NODE from SLURM_GPUS_ON_NODE when available. |
| runner/primus-cli-direct.sh | Honor TRAIN_LOG as a default log file path. |
| runner/primus-cli-container.sh | Ensure BNXT tar is volume-mounted when only provided via host path env var. |
| runner/helpers/launch/slurm_pretrain.sh | New shared SLURM launcher translating EXP/NNODES/DATA_PATH into primus-cli slurm srun .... |
| runner/helpers/hooks/train/pretrain/torchtitan/00_install_requirements.sh | Add PRIMUS_SKIP_PIP=1 early exit. |
| runner/helpers/hooks/train/pretrain/nemo_automodel/prepare.py | Move off deleted examples/scripts/utils.py; accept unknown args via parse_known_args. |
| runner/helpers/hooks/train/pretrain/megatron/prepare.py | Restore/extend dataset prep behaviors (split seed decoupling, held-out split, rank coordination) and argument parsing. |
| runner/helpers/hooks/train/pretrain/megatron/01_install_emerging_optimizers.sh | Add PRIMUS_SKIP_PIP=1 early exit. |
| runner/helpers/hooks/train/pretrain/megatron/00_install_requirements.sh | Add PRIMUS_SKIP_PIP=1 early exit. |
| runner/helpers/hooks/train/pretrain/megatron_bridge/00_install_requirements.sh | Add PRIMUS_SKIP_PIP=1 early exit. |
| runner/helpers/hooks/train/pretrain/maxtext/00_install_requirements.sh | Add PRIMUS_SKIP_PIP=1 early exit. |
| runner/helpers/hooks/train/pretrain/diffusion/00_install_requirements.sh | Add PRIMUS_SKIP_PIP=1 early exit. |
| runner/helpers/hooks/02_rebuild_bnxt.sh | Improve skip logging when BNXT tar is not readable at runtime. |
| runner/helpers/hooks/01_patch_te_flash_attn_max_version.sh | Update comment to reflect current implementation approach. |
| runner/helpers/envs/base_env.sh | Add writable cache-root fallback; default GPUS_PER_NODE from SLURM_GPUS_ON_NODE. |
| runner/.primus.yaml | Expand forwarded env allowlist (Databricks/MLflow/UEP) and adjust forwarded keys list. |
| primus/core/odc/rocshmem_runtime/scripts/run_odc.sh | Route ODC runs through primus-cli direct and propagate exit status with clearer DONE/FAILED logging. |
| primus/core/odc/rocshmem_runtime/scripts/cleanup_rs.sh | Update process-kill patterns to match primus-cli. |
| primus/backends/megatron/sft/packing.py | Update comment to reflect new cache/root semantics. |
| primus/backends/megatron/patches/fla_kda_autotune_patches.py | New in-process patch narrowing FLA KDA autotune space on ROCm + FLA-KDA runs. |
| primus/backends/megatron/core/transformer/fla_flash_attention.py | Update comment referencing launcher tee pipeline. |
| examples/torchtitan/prepare.py | Delete legacy examples prepare script (moved to hooks). |
| examples/scripts/utils.py | Delete legacy examples utility module (replaced by runner hook utils). |
| examples/scripts/prepare_experiment.py | Delete legacy prepare orchestrator (superseded by primus-cli hook flow). |
| examples/scripts/get_nccl_ib_hca.sh | Delete legacy helper (superseded by runner env helpers). |
| examples/scripts/get_ip_interface.sh | Delete legacy helper (superseded by runner env helpers). |
| examples/scripts/detect_nccl_ib_tc.sh | Delete legacy helper (superseded by runner hook/env logic). |
| examples/run_slurm_pretrain.sh | Delete legacy SLURM launcher. |
| examples/run_slurm_pretrain_cli.sh | Delete legacy SLURM CLI shim. |
| examples/run_pretrain.sh | Delete legacy main launcher. |
| examples/run_pretrain_cli.sh | Delete legacy CLI shim. |
| examples/run_local_pretrain.sh | Delete legacy local-container launcher. |
| examples/run_local_pretrain_cli.sh | Delete legacy local CLI shim. |
| examples/run_k8s_pretrain.sh | Delete legacy Kubernetes submission client script. |
| examples/README.md | Remove Kubernetes launcher section; document primus-cli as the entry point and shared SLURM helper. |
| examples/odc/README.md | Update ODC docs to describe primus-cli-based flow and updated SFT/prepare-hook behavior. |
| examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain.yaml | Update header launch comments to primus-cli. |
| examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain-bs8-acoff.yaml | Update header launch comments to primus-cli. |
| examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain-bs32-acon.yaml | Update header launch comments to primus-cli. |
| examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-finetune.yaml | Update header launch comments to primus-cli. |
| examples/moe_package/run_qwen_235B_a22B_pretrain_mi355x.sh | Update comment to reference shared SLURM helper. |
| examples/moe_package/run_gpt_oss_120B_mi355x.sh | Migrate to shared SLURM helper. |
| examples/moe_package/run_deepseek_v2_pretrain_mi355x.sh | Migrate to shared SLURM helper. |
| examples/moe_package/run_deepseek_v2_lite_pretrain_mi355x.sh | Migrate to shared SLURM helper. |
| examples/deepseek-v4/run_deepseek_v4_pro_muon_local.sh | Update local docker wrapper to invoke primus-cli direct instead of removed launcher. |
| examples/customer_package/run_qwen3_30b_a3b_pretrain_mi355x.sh | Migrate to shared SLURM helper. |
| examples/customer_package/run_qwen3_235b_a22b_pretrain_mi355x.sh | Migrate to shared SLURM helper. |
| examples/customer_package/run_qwen_8b_pretrain_mi355x.sh | Migrate to shared SLURM helper. |
| examples/customer_package/run_llama3.1_8b_pretrain_mi355x.sh | Migrate to shared SLURM helper. |
| examples/customer_package/run_llama3.1_70b_pretrain_mi355x.sh | Migrate to shared SLURM helper. |
| docs/05-operations/deployment.md | Remove Kubernetes client references; clarify k8s support expectations and required env vars. |
| docs/04-technical-guides/hybrid-models/README.md | Document the in-process FLA KDA autotune narrowing patch. |
| docs/04-technical-guides/hybrid-models/kda-guide.md | Replace on-disk workaround reference with new patch; update patch inventory. |
| docs/04-technical-guides/hybrid-models/gdn-guide.md | Clarify the patch does not apply to GDN paths. |
| docs/04-technical-guides/data-preparation.md | Update references from examples prepare scripts to runner prepare hooks; document split-seed behavior. |
| docs/03-configuration-reference/environment-variables.md | Update references and semantics around tokenized path env vars and tuning hooks. |
| docs/02-user-guide/torchtitan-training.md | Update “legacy script” section to shared SLURM helper. |
| docs/02-user-guide/pretraining.md | Update MaxDiffusion launch documentation to primus-cli direct. |
| docs/02-user-guide/megatron-lm-training.md | Update multi-node helper references to runner/helpers/launch/slurm_pretrain.sh. |
| benchmark/megatron/model/benchmark_model.py | Switch benchmark subprocess to shared SLURM helper. |
| benchmark/megatron/checkpoint/ckpt_launch.py | Switch checkpoint benchmark subprocess to shared SLURM helper. |
Review details
- Files reviewed: 81/82 changed files
- Comments generated: 2
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Capture the primus-cli status from PIPESTATUS before tee can mask it, so failed benchmark runs are reported correctly. Also correct the packed-dataset cache docstring to match the launcher environment defaults and pin the exit-code behavior with a regression test.
Xiaoming-AMD
approved these changes
Sep 11, 2026
jianhan-amd
added a commit
that referenced
this pull request
Sep 24, 2026
main retired the examples/ launchers in favour of primus-cli (#999) and deleted examples/megatron/prepare.py, which this branch had taught to skip bookcorpus tokenisation for diffusion runs. Keep the deletion: the replacement, runner/helpers/hooks/train/pretrain/megatron/prepare.py, already skips any config with model_type: diffusion_model, which every WAN config sets. Co-authored-by: Cursor <cursoragent@cursor.com>
jianhan-amd
added a commit
that referenced
this pull request
Sep 24, 2026
main retired examples/run_pretrain.sh in favour of primus-cli (#999), so wan_train.sh, which exec'd it, no longer ran. No other model keeps a launcher at the repository root either: Flux's Megatron configs are started with primus-cli directly and carry any recipe-specific environment in a top-level env: block, which TrainRuntime applies after the launcher's defaults. The WAN example configs now do the same, with each variable only where it has an effect: - MIOPEN_FIND_MODE=5 in all three. base_env.sh still pins Fast mode, which sends the Conv3d patch-embed backward to a naive im2col path. - PRIMUS_TE_GEMM_WORKSPACE_MIB=128 in te_spec only. The local specs never reach a TransformerEngine GEMM, so the workspace patch has nothing to resize. - PRIMUS_TURBO_GEMM_BACKEND=fp4:AITER in local_spec_mxfp4 only, since only the FP4 dispatcher reads it. PYTORCH_CUDA_ALLOC_CONF is left out on purpose. A YAML env value wins over anything exported on the CLI, and expandable_segments:True deadlocks in the KFD VA-map ioctl on some ROCm builds, so it has to stay a per-environment choice. The WAN_DEBUG block goes too; those are plain NCCL and torch variables to export when needed. The usage comments now show the primus-cli command, and the trainer comment that cited examples/run_pretrain.sh as a source of the atomic-fp32 default now cites only base_env.sh, which still sets it. Co-authored-by: Cursor <cursoragent@cursor.com>
Xiaoming-AMD
pushed a commit
that referenced
this pull request
Oct 8, 2026
The CLI runs runner/helpers/hooks/train/pretrain/maxdiffusion/prepare.py. This copy was only reachable through examples/scripts/prepare_experiment.py, which #999 removed together with the other examples/<backend>/prepare.py.
Xiaoming-AMD
pushed a commit
that referenced
this pull request
Oct 8, 2026
The CLI runs runner/helpers/hooks/train/pretrain/maxdiffusion/prepare.py. This copy was only reachable through examples/scripts/prepare_experiment.py, which #999 removed together with the other examples/<backend>/prepare.py.
Xiaoming-AMD
pushed a commit
that referenced
this pull request
Oct 8, 2026
…hook path The script said itself it was used only by examples/run_pretrain.sh, which #999 removed; primus-cli never called it. Its work is already covered: 00_install_requirements.sh installs the requirements, ROCm torch and the two site-package fixes, prepare.py hands the vendored checkout to the adapter as --backend_path instead of an editable install, and the Shardy and TensorFlow preload fixes are runtime patches under primus/backends/maxdiffusion/patches. The pretraining guide and examples/README.md still said the prepare hooks ran the script and that it applied all four patches, so describe what the hooks actually do, and point the warm-an-image recipe at the hook.
This was referenced Oct 9, 2026
Merged
i-kosarev
added a commit
to ROCm/MAD
that referenced
this pull request
Oct 9, 2026
The Primus payload check in primus_maxtext_rccl_overlay asserts examples/run_pretrain.sh. AMD-AGI/Primus#999 retired that script and scripts/jax-maxtext/run.sh has driven `primus-cli direct` since, so the gate gets both directions wrong: it fails the build on a Primus ref new enough to have dropped the file, and it passes on a ref where the entry point we actually run could be absent. Checked against the two live refs: ref run_pretrain.sh primus-cli v26.7.0 present present jax-maxtext-v26.8 ABSENT present So the build is fine today only because tools/fetch_primus.sh defaults to v26.7.0 after #279. Pointing PRIMUS_REF at jax-maxtext-v26.8 -- which carries the MI300X/MI325X/MI355X configs -- fails at this layer with a bare `test` exit 1 and no indication why. requirements-jax.txt is left alone: the next layer pip-installs it, so that check still guards something real.
Xiaoming-AMD
added a commit
that referenced
this pull request
Oct 12, 2026
#1130) ## Why Every backend already owns a top-level directory named after the `framework` field in its experiment YAML, but model studies and walkthroughs for those backends sat as *siblings* of the backends. So `examples/` mixed backends with Megatron-only content, and nothing told you which backend `deepseek-v4/` or `moe_package/` belonged to. ## What moved Megatron-specific material now lives under `examples/megatron/`: | New location | Was | | --- | --- | | `megatron/models/deepseek-v4/` | `examples/deepseek-v4/` | | `megatron/models/kimi-k3/` | `examples/models/kimi-k3/` | | `megatron/guides/customer_package/` | `examples/customer_package/` | | `megatron/guides/moe_package/` | `examples/moe_package/` | | `megatron/guides/odc/` | `examples/odc/` | | `megatron/guides/offline_tune/` | `examples/offline_tune/` | | `megatron/guides/tuning_agent/` | `examples/agents/tuning_agent/` | `mlperf/` and `hardware_configs/` deliberately stay at the top level: MLPerf submissions span Megatron, Megatron-Bridge and diffusion, and the hardware configs feed `primus projection` rather than any one backend. `examples/README.md` gains a Directory Layout section stating the rule. Backend directories keep their top-level position, which is what makes this safe: the two places that resolve a backend directory by name at runtime — `examples/scripts/prepare_experiment.py` building `examples/<framework>/prepare.py`, and `SCAN_ROOTS` in `primus/core/config/schema_check.py` scanning `examples/<backend>/configs` — are untouched, as are the `from examples.scripts.utils import ...` imports. ## Beyond renaming Ten scripts derive the repo root from their own location, e.g. ```bash REPO="$(cd "${HERE}/../../.." && pwd)" ``` They are two levels deeper now, so those `../` chains were extended and each one was verified to still resolve to the repo root. `.gitignore` (the negation that un-ignores the projection site's JSON), the `deploy-backend-gap-dashboard` path filter, `tools/backend_gap_report/pages-sections.json`, and the HipBLASLt stage-2 call in `runner/helpers/hooks/train/pretrain/prepare_experiment.sh` all hard-code moved paths and were updated too. ## Verification - All 105 references to the moved paths were rewritten; zero occurrences of the old paths remain. - A checker resolved all 609 literal `examples/...` references in tracked files against `main` and against this branch. Unresolved paths went from 61 to 59 — no new breakage, and two fewer because the tuning-agent docs now point at `target_cluster_mi355x_4nodes.yaml`, the only cluster file that exists (they referenced `target_cluster.yaml` and `..._2nodes.yaml`, neither of which was ever added). The remaining 59 are pre-existing and untouched: synthetic fixtures in `test_config_schema_check.py`, upstream paths such as Megatron-LM's `examples/multimodal/`, and shell strings the scanner truncates at `+`-concatenation. - `pre-commit run` passes on all 98 changed files, including shellcheck. - `pytest tests/runner/test_docs_cli_examples.py tests/unit_tests/configs/test_config_schema_check.py` — 64 passed. ## One unrelated fix, needed to get CI green `test_documented_script_invocations_exist` already fails on `main` (verified on a clean checkout): `docs/02-user-guide/jax-maxtext-training.md` runs `bash tools/fetch_primus.sh`, which is not in this repository. That command runs after `cd MAD`, so it belongs to the external ROCm/MAD checkout — it is allow-listed here with that rationale. #999 carries the identical fix, so if that merges first this becomes a trivial duplicate-line conflict. ## Note for #999 `examples/maxdiffusion/prepare.py` should be deleted there: that PR migrates the per-backend `prepare.py` files into `runner/helpers/hooks/train/pretrain/`, but this one survived and still does `from examples.scripts.utils import ...`, a module the same PR deletes. Nothing references it and `runner/helpers/hooks/train/pretrain/maxdiffusion/prepare.py` supersedes it. ## Reviewing `git show --find-renames` (or `?w=1` with rename detection) keeps this to the ~40 content lines; the rest is pure movement. --------- Co-authored-by: xmpeng-dev <xmpeng.dev@gmail.com> Co-authored-by: xiaoming-amd <xiaompen@amd.com>
i-kosarev
added a commit
to ROCm/rocm-systems
that referenced
this pull request
Oct 13, 2026
…e gate Adds `llama-3.1-405b-maxtext` at 4 nodes (32 GPUs) as the first JAX workload in the gate, and the first whose base image is not rocm/primus. Stacked on Draft until #12899 and #13180 land. The enabling work is ROCm/MAD#228, which validated this workload at 4N. Model family ------------ Until now every workload was PyTorch/Megatron-LM out of MAD's scripts/primus_megatron-lm. A "family" key now dispatches the three places that differ: install_madengine() copies scripts/jax-maxtext and checks Primus out, generate_manifest() emits `--config_path <yaml>` plus a `distributed.primus` block instead of `--model_repo`, and the overlay builder uses a different recipe. Settings lifted from MAD#228's template are each a failure mode, not taste: XLA_GPU_AUTOTUNE_LEVEL=0 Primus defaults to level 4, which at 32 ranks presents as a dead hang -- no steps, no NCCL output -- rather than a slow start. The MAD-native XLA_AUTOTUNE_LEVEL is a silent no-op. per_device_batch_size=1 the shipped config OOMs at 241.25 GiB for jit_train_step at 4 nodes. --ulimit nofile=1048576 a 32-device clique exhausts the default and stalls silently at "Initialize clique". They ride the per-workload env_vars key #13180 introduced. Why this family needs its own A/B baseline ------------------------------------------ The Megatron baseline is the stock registry image. That cannot work here: rocm/jax-training:maxtext-* ships MaxText but no Primus, so scripts/jax-maxtext/run.sh would not find primus-cli and the baseline phase would die before step 0. So the baseline is built: base image plus the Primus payload, stock RCCL. The candidate is that same Dockerfile plus the librccl swap -- literally a prefix, so the payload layers come out of the cache and only the swap builds twice. The pair therefore differs by the RCCL swap and nothing else, which is tighter than the Megatron A/B rather than looser. The baseline manifest advertises its own Dockerfile. Left at the default a compute node would rebuild the swapped image under the baseline tag, giving an A/B that compares the CI RCCL against itself. Overlay ------- The swap for this family is discovery-based -- find, per-target RUNPATH preservation, and an unresolved-soname set gate -- ported from MAD#228, rather than aimed at the Primus SDK venv paths, which this base lacks. The libamd_smi soname bridge is kept and now runs over every directory the sweep finds, ordered after the copy and before the post-swap measurement so a .27-against-.26 skew is reconciled instead of failing the gate. Unlike MAD#228's dockerfile the Primus payload is gated on primus-cli, not examples/run_pretrain.sh: AMD-AGI/Primus#999 retired that file and run.sh drives primus-cli. v26.7.0 still carries both so MAD's gate passes today, but it would hard-fail on jax-maxtext-v26.8. PRIMUS_REF is pinned here. --skip-overlay-build is refused for this family: its bind-mount targets are Primus-image paths absent from this base, and `docker run -v` creates a missing source as an empty directory rather than refusing, so the run would look healthy while measuring the image's own RCCL. Scheduling ---------- Its own job rather than a third matrix entry: the weekly gate has to be a job-level `if:`, where the matrix context is unavailable, and it is fixed at 4 nodes while the matrix entries take madengine_nodes. WORK_DIR already carries inputs.workload, so they do not collide. Not gated on the matrix job -- a 70B failure should not hide a 405B result. time_limit is 05:00:00 for the A/B pair, against MAD#228's 12:00:00 for a single run: a 4-node exclusive request with a 12 h limit backfills badly, the measured run is ~17 min of training (MAD#228 reports 17.40 s/step) plus prep, and 300 min stays inside the step's existing 330 so the script still reports. The existing 360/330 budget needed no change. Risk to the existing workloads ------------------------------ MAD_REF moves 07ecef61 -> 7acfac7c, which is exactly one commit, ROCm/MAD#279, repointing tools/fetch_primus.sh at a ref that still exists. It touches benchmark/jax_maxtext/README.md, docker/primus_maxdiffusion.ubuntu.amd.Dockerfile and tools/fetch_primus.sh -- nothing under scripts/primus_megatron-lm. I rendered both A/B phase manifests for llama-3.1-70b, llama-4-scout and gpt-oss-120b against this branch and its base: byte-identical, all six. Also verified: the 405B manifest matches MAD#228's validated template field for field; the baseline and candidate manifests differ only in image and dockerfile; both generated Dockerfiles parse under `bash -n`, with the bridge ordered after the copy and before the gate. Not yet run on hardware. The first run has to establish whether the TheRock kpack librccl loads on this base at all, and whether its libamd_smi soname needs the bridge here. JIRA ID AICOMNET-420
i-kosarev
added a commit
to ROCm/rocm-systems
that referenced
this pull request
Oct 14, 2026
…e gate Adds `llama-3.1-405b-maxtext` at 4 nodes (32 GPUs) as the first JAX workload in the gate, and the first whose base image is not rocm/primus. Stacked on Draft until #12899 and #13180 land. The enabling work is ROCm/MAD#228, which validated this workload at 4N. Model family ------------ Until now every workload was PyTorch/Megatron-LM out of MAD's scripts/primus_megatron-lm. A "family" key now dispatches the three places that differ: install_madengine() copies scripts/jax-maxtext and checks Primus out, generate_manifest() emits `--config_path <yaml>` plus a `distributed.primus` block instead of `--model_repo`, and the overlay builder uses a different recipe. Settings lifted from MAD#228's template are each a failure mode, not taste: XLA_GPU_AUTOTUNE_LEVEL=0 Primus defaults to level 4, which at 32 ranks presents as a dead hang -- no steps, no NCCL output -- rather than a slow start. The MAD-native XLA_AUTOTUNE_LEVEL is a silent no-op. per_device_batch_size=1 the shipped config OOMs at 241.25 GiB for jit_train_step at 4 nodes. --ulimit nofile=1048576 a 32-device clique exhausts the default and stalls silently at "Initialize clique". They ride the per-workload env_vars key #13180 introduced. Why this family needs its own A/B baseline ------------------------------------------ The Megatron baseline is the stock registry image. That cannot work here: rocm/jax-training:maxtext-* ships MaxText but no Primus, so scripts/jax-maxtext/run.sh would not find primus-cli and the baseline phase would die before step 0. So the baseline is built: base image plus the Primus payload, stock RCCL. The candidate is that same Dockerfile plus the librccl swap -- literally a prefix, so the payload layers come out of the cache and only the swap builds twice. The pair therefore differs by the RCCL swap and nothing else, which is tighter than the Megatron A/B rather than looser. The baseline manifest advertises its own Dockerfile. Left at the default a compute node would rebuild the swapped image under the baseline tag, giving an A/B that compares the CI RCCL against itself. Overlay ------- The swap for this family is discovery-based -- find, per-target RUNPATH preservation, and an unresolved-soname set gate -- ported from MAD#228, rather than aimed at the Primus SDK venv paths, which this base lacks. The libamd_smi soname bridge is kept and now runs over every directory the sweep finds, ordered after the copy and before the post-swap measurement so a .27-against-.26 skew is reconciled instead of failing the gate. Unlike MAD#228's dockerfile the Primus payload is gated on primus-cli, not examples/run_pretrain.sh: AMD-AGI/Primus#999 retired that file and run.sh drives primus-cli. v26.7.0 still carries both so MAD's gate passes today, but it would hard-fail on jax-maxtext-v26.8. PRIMUS_REF is pinned here. --skip-overlay-build is refused for this family: its bind-mount targets are Primus-image paths absent from this base, and `docker run -v` creates a missing source as an empty directory rather than refusing, so the run would look healthy while measuring the image's own RCCL. Scheduling ---------- Its own job rather than a third matrix entry: the weekly gate has to be a job-level `if:`, where the matrix context is unavailable, and it is fixed at 4 nodes while the matrix entries take madengine_nodes. WORK_DIR already carries inputs.workload, so they do not collide. Not gated on the matrix job -- a 70B failure should not hide a 405B result. time_limit is 05:00:00 for the A/B pair, against MAD#228's 12:00:00 for a single run: a 4-node exclusive request with a 12 h limit backfills badly, the measured run is ~17 min of training (MAD#228 reports 17.40 s/step) plus prep, and 300 min stays inside the step's existing 330 so the script still reports. The existing 360/330 budget needed no change. Risk to the existing workloads ------------------------------ MAD_REF moves 07ecef61 -> 7acfac7c, which is exactly one commit, ROCm/MAD#279, repointing tools/fetch_primus.sh at a ref that still exists. It touches benchmark/jax_maxtext/README.md, docker/primus_maxdiffusion.ubuntu.amd.Dockerfile and tools/fetch_primus.sh -- nothing under scripts/primus_megatron-lm. I rendered both A/B phase manifests for llama-3.1-70b, llama-4-scout and gpt-oss-120b against this branch and its base: byte-identical, all six. Also verified: the 405B manifest matches MAD#228's validated template field for field; the baseline and candidate manifests differ only in image and dockerfile; both generated Dockerfiles parse under `bash -n`, with the bridge ordered after the copy and before the gate. Not yet run on hardware. The first run has to establish whether the TheRock kpack librccl loads on this base at all, and whether its libamd_smi soname needs the bridge here. JIRA ID AICOMNET-420
i-kosarev
added a commit
to ROCm/rocm-systems
that referenced
this pull request
Oct 14, 2026
…e gate Adds `llama-3.1-405b-maxtext` at 4 nodes (32 GPUs) as the first JAX workload in the gate, and the first whose base image is not rocm/primus. Stacked on Draft until #12899 and #13180 land. The enabling work is ROCm/MAD#228, which validated this workload at 4N. Model family ------------ Until now every workload was PyTorch/Megatron-LM out of MAD's scripts/primus_megatron-lm. A "family" key now dispatches the three places that differ: install_madengine() copies scripts/jax-maxtext and checks Primus out, generate_manifest() emits `--config_path <yaml>` plus a `distributed.primus` block instead of `--model_repo`, and the overlay builder uses a different recipe. Settings lifted from MAD#228's template are each a failure mode, not taste: XLA_GPU_AUTOTUNE_LEVEL=0 Primus defaults to level 4, which at 32 ranks presents as a dead hang -- no steps, no NCCL output -- rather than a slow start. The MAD-native XLA_AUTOTUNE_LEVEL is a silent no-op. per_device_batch_size=1 the shipped config OOMs at 241.25 GiB for jit_train_step at 4 nodes. --ulimit nofile=1048576 a 32-device clique exhausts the default and stalls silently at "Initialize clique". They ride the per-workload env_vars key #13180 introduced. Why this family needs its own A/B baseline ------------------------------------------ The Megatron baseline is the stock registry image. That cannot work here: rocm/jax-training:maxtext-* ships MaxText but no Primus, so scripts/jax-maxtext/run.sh would not find primus-cli and the baseline phase would die before step 0. So the baseline is built: base image plus the Primus payload, stock RCCL. The candidate is that same Dockerfile plus the librccl swap -- literally a prefix, so the payload layers come out of the cache and only the swap builds twice. The pair therefore differs by the RCCL swap and nothing else, which is tighter than the Megatron A/B rather than looser. The baseline manifest advertises its own Dockerfile. Left at the default a compute node would rebuild the swapped image under the baseline tag, giving an A/B that compares the CI RCCL against itself. Overlay ------- The swap for this family is discovery-based -- find, per-target RUNPATH preservation, and an unresolved-soname set gate -- ported from MAD#228, rather than aimed at the Primus SDK venv paths, which this base lacks. The libamd_smi soname bridge is kept and now runs over every directory the sweep finds, ordered after the copy and before the post-swap measurement so a .27-against-.26 skew is reconciled instead of failing the gate. Unlike MAD#228's dockerfile the Primus payload is gated on primus-cli, not examples/run_pretrain.sh: AMD-AGI/Primus#999 retired that file and run.sh drives primus-cli. v26.7.0 still carries both so MAD's gate passes today, but it would hard-fail on jax-maxtext-v26.8. PRIMUS_REF is pinned here. --skip-overlay-build is refused for this family: its bind-mount targets are Primus-image paths absent from this base, and `docker run -v` creates a missing source as an empty directory rather than refusing, so the run would look healthy while measuring the image's own RCCL. Scheduling ---------- Its own job rather than a third matrix entry: the weekly gate has to be a job-level `if:`, where the matrix context is unavailable, and it is fixed at 4 nodes while the matrix entries take madengine_nodes. WORK_DIR already carries inputs.workload, so they do not collide. Not gated on the matrix job -- a 70B failure should not hide a 405B result. time_limit is 05:00:00 for the A/B pair, against MAD#228's 12:00:00 for a single run: a 4-node exclusive request with a 12 h limit backfills badly, the measured run is ~17 min of training (MAD#228 reports 17.40 s/step) plus prep, and 300 min stays inside the step's existing 330 so the script still reports. The existing 360/330 budget needed no change. Risk to the existing workloads ------------------------------ MAD_REF moves 07ecef61 -> 7acfac7c, which is exactly one commit, ROCm/MAD#279, repointing tools/fetch_primus.sh at a ref that still exists. It touches benchmark/jax_maxtext/README.md, docker/primus_maxdiffusion.ubuntu.amd.Dockerfile and tools/fetch_primus.sh -- nothing under scripts/primus_megatron-lm. I rendered both A/B phase manifests for llama-3.1-70b, llama-4-scout and gpt-oss-120b against this branch and its base: byte-identical, all six. Also verified: the 405B manifest matches MAD#228's validated template field for field; the baseline and candidate manifests differ only in image and dockerfile; both generated Dockerfiles parse under `bash -n`, with the bridge ordered after the copy and before the gate. Not yet run on hardware. The first run has to establish whether the TheRock kpack librccl loads on this base at all, and whether its libamd_smi soname needs the bridge here. JIRA ID AICOMNET-420
i-kosarev
added a commit
to ROCm/rocm-systems
that referenced
this pull request
Oct 15, 2026
…e gate Adds `llama-3.1-405b-maxtext` at 4 nodes (32 GPUs) as the first JAX workload in the gate, and the first whose base image is not rocm/primus. Stacked on Draft until #12899 and #13180 land. The enabling work is ROCm/MAD#228, which validated this workload at 4N. Model family ------------ Until now every workload was PyTorch/Megatron-LM out of MAD's scripts/primus_megatron-lm. A "family" key now dispatches the three places that differ: install_madengine() copies scripts/jax-maxtext and checks Primus out, generate_manifest() emits `--config_path <yaml>` plus a `distributed.primus` block instead of `--model_repo`, and the overlay builder uses a different recipe. Settings lifted from MAD#228's template are each a failure mode, not taste: XLA_GPU_AUTOTUNE_LEVEL=0 Primus defaults to level 4, which at 32 ranks presents as a dead hang -- no steps, no NCCL output -- rather than a slow start. The MAD-native XLA_AUTOTUNE_LEVEL is a silent no-op. per_device_batch_size=1 the shipped config OOMs at 241.25 GiB for jit_train_step at 4 nodes. --ulimit nofile=1048576 a 32-device clique exhausts the default and stalls silently at "Initialize clique". They ride the per-workload env_vars key #13180 introduced. Why this family needs its own A/B baseline ------------------------------------------ The Megatron baseline is the stock registry image. That cannot work here: rocm/jax-training:maxtext-* ships MaxText but no Primus, so scripts/jax-maxtext/run.sh would not find primus-cli and the baseline phase would die before step 0. So the baseline is built: base image plus the Primus payload, stock RCCL. The candidate is that same Dockerfile plus the librccl swap -- literally a prefix, so the payload layers come out of the cache and only the swap builds twice. The pair therefore differs by the RCCL swap and nothing else, which is tighter than the Megatron A/B rather than looser. The baseline manifest advertises its own Dockerfile. Left at the default a compute node would rebuild the swapped image under the baseline tag, giving an A/B that compares the CI RCCL against itself. Overlay ------- The swap for this family is discovery-based -- find, per-target RUNPATH preservation, and an unresolved-soname set gate -- ported from MAD#228, rather than aimed at the Primus SDK venv paths, which this base lacks. The libamd_smi soname bridge is kept and now runs over every directory the sweep finds, ordered after the copy and before the post-swap measurement so a .27-against-.26 skew is reconciled instead of failing the gate. Unlike MAD#228's dockerfile the Primus payload is gated on primus-cli, not examples/run_pretrain.sh: AMD-AGI/Primus#999 retired that file and run.sh drives primus-cli. v26.7.0 still carries both so MAD's gate passes today, but it would hard-fail on jax-maxtext-v26.8. PRIMUS_REF is pinned here. --skip-overlay-build is refused for this family: its bind-mount targets are Primus-image paths absent from this base, and `docker run -v` creates a missing source as an empty directory rather than refusing, so the run would look healthy while measuring the image's own RCCL. Scheduling ---------- Its own job rather than a third matrix entry: the weekly gate has to be a job-level `if:`, where the matrix context is unavailable, and it is fixed at 4 nodes while the matrix entries take madengine_nodes. WORK_DIR already carries inputs.workload, so they do not collide. Not gated on the matrix job -- a 70B failure should not hide a 405B result. time_limit is 05:00:00 for the A/B pair, against MAD#228's 12:00:00 for a single run: a 4-node exclusive request with a 12 h limit backfills badly, the measured run is ~17 min of training (MAD#228 reports 17.40 s/step) plus prep, and 300 min stays inside the step's existing 330 so the script still reports. The existing 360/330 budget needed no change. Risk to the existing workloads ------------------------------ MAD_REF moves 07ecef61 -> 7acfac7c, which is exactly one commit, ROCm/MAD#279, repointing tools/fetch_primus.sh at a ref that still exists. It touches benchmark/jax_maxtext/README.md, docker/primus_maxdiffusion.ubuntu.amd.Dockerfile and tools/fetch_primus.sh -- nothing under scripts/primus_megatron-lm. I rendered both A/B phase manifests for llama-3.1-70b, llama-4-scout and gpt-oss-120b against this branch and its base: byte-identical, all six. Also verified: the 405B manifest matches MAD#228's validated template field for field; the baseline and candidate manifests differ only in image and dockerfile; both generated Dockerfiles parse under `bash -n`, with the bridge ordered after the copy and before the gate. Not yet run on hardware. The first run has to establish whether the TheRock kpack librccl loads on this base at all, and whether its libamd_smi soname needs the bridge here. JIRA ID AICOMNET-420
jianhan-amd
added a commit
that referenced
this pull request
Oct 15, 2026
main retired examples/run_pretrain.sh in favour of primus-cli (#999), so wan_train.sh, which exec'd it, no longer ran. No other model keeps a launcher at the repository root either: Flux's Megatron configs are started with primus-cli directly and carry any recipe-specific environment in a top-level env: block, which TrainRuntime applies after the launcher's defaults. The WAN example configs now do the same, with each variable only where it has an effect: - MIOPEN_FIND_MODE=5 in all three. base_env.sh still pins Fast mode, which sends the Conv3d patch-embed backward to a naive im2col path. - PRIMUS_TE_GEMM_WORKSPACE_MIB=128 in te_spec only. The local specs never reach a TransformerEngine GEMM, so the workspace patch has nothing to resize. - PRIMUS_TURBO_GEMM_BACKEND=fp4:AITER in local_spec_mxfp4 only, since only the FP4 dispatcher reads it. PYTORCH_CUDA_ALLOC_CONF is left out on purpose. A YAML env value wins over anything exported on the CLI, and expandable_segments:True deadlocks in the KFD VA-map ioctl on some ROCm builds, so it has to stay a per-environment choice. The WAN_DEBUG block goes too; those are plain NCCL and torch variables to export when needed. The usage comments now show the primus-cli command, and the trainer comment that cited examples/run_pretrain.sh as a source of the atomic-fp32 default now cites only base_env.sh, which still sets it. Co-authored-by: Cursor <cursoragent@cursor.com>
i-kosarev
added a commit
to ROCm/rocm-systems
that referenced
this pull request
Oct 19, 2026
…e gate Adds `llama-3.1-405b-maxtext` at 4 nodes (32 GPUs) as the first JAX workload in the gate, and the first whose base image is not rocm/primus. Stacked on Draft until #12899 and #13180 land. The enabling work is ROCm/MAD#228, which validated this workload at 4N. Model family ------------ Until now every workload was PyTorch/Megatron-LM out of MAD's scripts/primus_megatron-lm. A "family" key now dispatches the three places that differ: install_madengine() copies scripts/jax-maxtext and checks Primus out, generate_manifest() emits `--config_path <yaml>` plus a `distributed.primus` block instead of `--model_repo`, and the overlay builder uses a different recipe. Settings lifted from MAD#228's template are each a failure mode, not taste: XLA_GPU_AUTOTUNE_LEVEL=0 Primus defaults to level 4, which at 32 ranks presents as a dead hang -- no steps, no NCCL output -- rather than a slow start. The MAD-native XLA_AUTOTUNE_LEVEL is a silent no-op. per_device_batch_size=1 the shipped config OOMs at 241.25 GiB for jit_train_step at 4 nodes. --ulimit nofile=1048576 a 32-device clique exhausts the default and stalls silently at "Initialize clique". They ride the per-workload env_vars key #13180 introduced. Why this family needs its own A/B baseline ------------------------------------------ The Megatron baseline is the stock registry image. That cannot work here: rocm/jax-training:maxtext-* ships MaxText but no Primus, so scripts/jax-maxtext/run.sh would not find primus-cli and the baseline phase would die before step 0. So the baseline is built: base image plus the Primus payload, stock RCCL. The candidate is that same Dockerfile plus the librccl swap -- literally a prefix, so the payload layers come out of the cache and only the swap builds twice. The pair therefore differs by the RCCL swap and nothing else, which is tighter than the Megatron A/B rather than looser. The baseline manifest advertises its own Dockerfile. Left at the default a compute node would rebuild the swapped image under the baseline tag, giving an A/B that compares the CI RCCL against itself. Overlay ------- The swap for this family is discovery-based -- find, per-target RUNPATH preservation, and an unresolved-soname set gate -- ported from MAD#228, rather than aimed at the Primus SDK venv paths, which this base lacks. The libamd_smi soname bridge is kept and now runs over every directory the sweep finds, ordered after the copy and before the post-swap measurement so a .27-against-.26 skew is reconciled instead of failing the gate. Unlike MAD#228's dockerfile the Primus payload is gated on primus-cli, not examples/run_pretrain.sh: AMD-AGI/Primus#999 retired that file and run.sh drives primus-cli. v26.7.0 still carries both so MAD's gate passes today, but it would hard-fail on jax-maxtext-v26.8. PRIMUS_REF is pinned here. --skip-overlay-build is refused for this family: its bind-mount targets are Primus-image paths absent from this base, and `docker run -v` creates a missing source as an empty directory rather than refusing, so the run would look healthy while measuring the image's own RCCL. Scheduling ---------- Its own job rather than a third matrix entry: the weekly gate has to be a job-level `if:`, where the matrix context is unavailable, and it is fixed at 4 nodes while the matrix entries take madengine_nodes. WORK_DIR already carries inputs.workload, so they do not collide. Not gated on the matrix job -- a 70B failure should not hide a 405B result. time_limit is 05:00:00 for the A/B pair, against MAD#228's 12:00:00 for a single run: a 4-node exclusive request with a 12 h limit backfills badly, the measured run is ~17 min of training (MAD#228 reports 17.40 s/step) plus prep, and 300 min stays inside the step's existing 330 so the script still reports. The existing 360/330 budget needed no change. Risk to the existing workloads ------------------------------ MAD_REF moves 07ecef61 -> 7acfac7c, which is exactly one commit, ROCm/MAD#279, repointing tools/fetch_primus.sh at a ref that still exists. It touches benchmark/jax_maxtext/README.md, docker/primus_maxdiffusion.ubuntu.amd.Dockerfile and tools/fetch_primus.sh -- nothing under scripts/primus_megatron-lm. I rendered both A/B phase manifests for llama-3.1-70b, llama-4-scout and gpt-oss-120b against this branch and its base: byte-identical, all six. Also verified: the 405B manifest matches MAD#228's validated template field for field; the baseline and candidate manifests differ only in image and dockerfile; both generated Dockerfiles parse under `bash -n`, with the bridge ordered after the copy and before the gate. Not yet run on hardware. The first run has to establish whether the TheRock kpack librccl loads on this base at all, and whether its libamd_smi soname needs the bridge here. JIRA ID AICOMNET-420
i-kosarev
added a commit
to ROCm/rocm-systems
that referenced
this pull request
Oct 19, 2026
…e gate Adds `llama-3.1-405b-maxtext` at 4 nodes (32 GPUs) as the first JAX workload in the gate, and the first whose base image is not rocm/primus. Stacked on Draft until #12899 and #13180 land. The enabling work is ROCm/MAD#228, which validated this workload at 4N. Model family ------------ Until now every workload was PyTorch/Megatron-LM out of MAD's scripts/primus_megatron-lm. A "family" key now dispatches the three places that differ: install_madengine() copies scripts/jax-maxtext and checks Primus out, generate_manifest() emits `--config_path <yaml>` plus a `distributed.primus` block instead of `--model_repo`, and the overlay builder uses a different recipe. Settings lifted from MAD#228's template are each a failure mode, not taste: XLA_GPU_AUTOTUNE_LEVEL=0 Primus defaults to level 4, which at 32 ranks presents as a dead hang -- no steps, no NCCL output -- rather than a slow start. The MAD-native XLA_AUTOTUNE_LEVEL is a silent no-op. per_device_batch_size=1 the shipped config OOMs at 241.25 GiB for jit_train_step at 4 nodes. --ulimit nofile=1048576 a 32-device clique exhausts the default and stalls silently at "Initialize clique". They ride the per-workload env_vars key #13180 introduced. Why this family needs its own A/B baseline ------------------------------------------ The Megatron baseline is the stock registry image. That cannot work here: rocm/jax-training:maxtext-* ships MaxText but no Primus, so scripts/jax-maxtext/run.sh would not find primus-cli and the baseline phase would die before step 0. So the baseline is built: base image plus the Primus payload, stock RCCL. The candidate is that same Dockerfile plus the librccl swap -- literally a prefix, so the payload layers come out of the cache and only the swap builds twice. The pair therefore differs by the RCCL swap and nothing else, which is tighter than the Megatron A/B rather than looser. The baseline manifest advertises its own Dockerfile. Left at the default a compute node would rebuild the swapped image under the baseline tag, giving an A/B that compares the CI RCCL against itself. Overlay ------- The swap for this family is discovery-based -- find, per-target RUNPATH preservation, and an unresolved-soname set gate -- ported from MAD#228, rather than aimed at the Primus SDK venv paths, which this base lacks. The libamd_smi soname bridge is kept and now runs over every directory the sweep finds, ordered after the copy and before the post-swap measurement so a .27-against-.26 skew is reconciled instead of failing the gate. Unlike MAD#228's dockerfile the Primus payload is gated on primus-cli, not examples/run_pretrain.sh: AMD-AGI/Primus#999 retired that file and run.sh drives primus-cli. v26.7.0 still carries both so MAD's gate passes today, but it would hard-fail on jax-maxtext-v26.8. PRIMUS_REF is pinned here. --skip-overlay-build is refused for this family: its bind-mount targets are Primus-image paths absent from this base, and `docker run -v` creates a missing source as an empty directory rather than refusing, so the run would look healthy while measuring the image's own RCCL. Scheduling ---------- Its own job rather than a third matrix entry: the weekly gate has to be a job-level `if:`, where the matrix context is unavailable, and it is fixed at 4 nodes while the matrix entries take madengine_nodes. WORK_DIR already carries inputs.workload, so they do not collide. Not gated on the matrix job -- a 70B failure should not hide a 405B result. time_limit is 05:00:00 for the A/B pair, against MAD#228's 12:00:00 for a single run: a 4-node exclusive request with a 12 h limit backfills badly, the measured run is ~17 min of training (MAD#228 reports 17.40 s/step) plus prep, and 300 min stays inside the step's existing 330 so the script still reports. The existing 360/330 budget needed no change. Risk to the existing workloads ------------------------------ MAD_REF moves 07ecef61 -> 7acfac7c, which is exactly one commit, ROCm/MAD#279, repointing tools/fetch_primus.sh at a ref that still exists. It touches benchmark/jax_maxtext/README.md, docker/primus_maxdiffusion.ubuntu.amd.Dockerfile and tools/fetch_primus.sh -- nothing under scripts/primus_megatron-lm. I rendered both A/B phase manifests for llama-3.1-70b, llama-4-scout and gpt-oss-120b against this branch and its base: byte-identical, all six. Also verified: the 405B manifest matches MAD#228's validated template field for field; the baseline and candidate manifests differ only in image and dockerfile; both generated Dockerfiles parse under `bash -n`, with the bridge ordered after the copy and before the gate. Not yet run on hardware. The first run has to establish whether the TheRock kpack librccl loads on this base at all, and whether its libamd_smi soname needs the bridge here. JIRA ID AICOMNET-420
i-kosarev
added a commit
to ROCm/rocm-systems
that referenced
this pull request
Oct 19, 2026
…e gate Adds `llama-3.1-405b-maxtext` at 4 nodes (32 GPUs) as the first JAX workload in the gate, and the first whose base image is not rocm/primus. Stacked on Draft until #12899 and #13180 land. The enabling work is ROCm/MAD#228, which validated this workload at 4N. Model family ------------ Until now every workload was PyTorch/Megatron-LM out of MAD's scripts/primus_megatron-lm. A "family" key now dispatches the three places that differ: install_madengine() copies scripts/jax-maxtext and checks Primus out, generate_manifest() emits `--config_path <yaml>` plus a `distributed.primus` block instead of `--model_repo`, and the overlay builder uses a different recipe. Settings lifted from MAD#228's template are each a failure mode, not taste: XLA_GPU_AUTOTUNE_LEVEL=0 Primus defaults to level 4, which at 32 ranks presents as a dead hang -- no steps, no NCCL output -- rather than a slow start. The MAD-native XLA_AUTOTUNE_LEVEL is a silent no-op. per_device_batch_size=1 the shipped config OOMs at 241.25 GiB for jit_train_step at 4 nodes. --ulimit nofile=1048576 a 32-device clique exhausts the default and stalls silently at "Initialize clique". They ride the per-workload env_vars key #13180 introduced. Why this family needs its own A/B baseline ------------------------------------------ The Megatron baseline is the stock registry image. That cannot work here: rocm/jax-training:maxtext-* ships MaxText but no Primus, so scripts/jax-maxtext/run.sh would not find primus-cli and the baseline phase would die before step 0. So the baseline is built: base image plus the Primus payload, stock RCCL. The candidate is that same Dockerfile plus the librccl swap -- literally a prefix, so the payload layers come out of the cache and only the swap builds twice. The pair therefore differs by the RCCL swap and nothing else, which is tighter than the Megatron A/B rather than looser. The baseline manifest advertises its own Dockerfile. Left at the default a compute node would rebuild the swapped image under the baseline tag, giving an A/B that compares the CI RCCL against itself. Overlay ------- The swap for this family is discovery-based -- find, per-target RUNPATH preservation, and an unresolved-soname set gate -- ported from MAD#228, rather than aimed at the Primus SDK venv paths, which this base lacks. The libamd_smi soname bridge is kept and now runs over every directory the sweep finds, ordered after the copy and before the post-swap measurement so a .27-against-.26 skew is reconciled instead of failing the gate. Unlike MAD#228's dockerfile the Primus payload is gated on primus-cli, not examples/run_pretrain.sh: AMD-AGI/Primus#999 retired that file and run.sh drives primus-cli. v26.7.0 still carries both so MAD's gate passes today, but it would hard-fail on jax-maxtext-v26.8. PRIMUS_REF is pinned here. --skip-overlay-build is refused for this family: its bind-mount targets are Primus-image paths absent from this base, and `docker run -v` creates a missing source as an empty directory rather than refusing, so the run would look healthy while measuring the image's own RCCL. Scheduling ---------- Its own job rather than a third matrix entry: the weekly gate has to be a job-level `if:`, where the matrix context is unavailable, and it is fixed at 4 nodes while the matrix entries take madengine_nodes. WORK_DIR already carries inputs.workload, so they do not collide. Not gated on the matrix job -- a 70B failure should not hide a 405B result. time_limit is 05:00:00 for the A/B pair, against MAD#228's 12:00:00 for a single run: a 4-node exclusive request with a 12 h limit backfills badly, the measured run is ~17 min of training (MAD#228 reports 17.40 s/step) plus prep, and 300 min stays inside the step's existing 330 so the script still reports. The existing 360/330 budget needed no change. Risk to the existing workloads ------------------------------ MAD_REF moves 07ecef61 -> 7acfac7c, which is exactly one commit, ROCm/MAD#279, repointing tools/fetch_primus.sh at a ref that still exists. It touches benchmark/jax_maxtext/README.md, docker/primus_maxdiffusion.ubuntu.amd.Dockerfile and tools/fetch_primus.sh -- nothing under scripts/primus_megatron-lm. I rendered both A/B phase manifests for llama-3.1-70b, llama-4-scout and gpt-oss-120b against this branch and its base: byte-identical, all six. Also verified: the 405B manifest matches MAD#228's validated template field for field; the baseline and candidate manifests differ only in image and dockerfile; both generated Dockerfiles parse under `bash -n`, with the bridge ordered after the copy and before the gate. Not yet run on hardware. The first run has to establish whether the TheRock kpack librccl loads on this base at all, and whether its libamd_smi soname needs the bridge here. JIRA ID AICOMNET-420
mkuznet1
pushed a commit
to ROCm/rocm-systems
that referenced
this pull request
Oct 19, 2026
…e gate Adds `llama-3.1-405b-maxtext` at 4 nodes (32 GPUs) as the first JAX workload in the gate, and the first whose base image is not rocm/primus. Stacked on Draft until #12899 and #13180 land. The enabling work is ROCm/MAD#228, which validated this workload at 4N. Model family ------------ Until now every workload was PyTorch/Megatron-LM out of MAD's scripts/primus_megatron-lm. A "family" key now dispatches the three places that differ: install_madengine() copies scripts/jax-maxtext and checks Primus out, generate_manifest() emits `--config_path <yaml>` plus a `distributed.primus` block instead of `--model_repo`, and the overlay builder uses a different recipe. Settings lifted from MAD#228's template are each a failure mode, not taste: XLA_GPU_AUTOTUNE_LEVEL=0 Primus defaults to level 4, which at 32 ranks presents as a dead hang -- no steps, no NCCL output -- rather than a slow start. The MAD-native XLA_AUTOTUNE_LEVEL is a silent no-op. per_device_batch_size=1 the shipped config OOMs at 241.25 GiB for jit_train_step at 4 nodes. --ulimit nofile=1048576 a 32-device clique exhausts the default and stalls silently at "Initialize clique". They ride the per-workload env_vars key #13180 introduced. Why this family needs its own A/B baseline ------------------------------------------ The Megatron baseline is the stock registry image. That cannot work here: rocm/jax-training:maxtext-* ships MaxText but no Primus, so scripts/jax-maxtext/run.sh would not find primus-cli and the baseline phase would die before step 0. So the baseline is built: base image plus the Primus payload, stock RCCL. The candidate is that same Dockerfile plus the librccl swap -- literally a prefix, so the payload layers come out of the cache and only the swap builds twice. The pair therefore differs by the RCCL swap and nothing else, which is tighter than the Megatron A/B rather than looser. The baseline manifest advertises its own Dockerfile. Left at the default a compute node would rebuild the swapped image under the baseline tag, giving an A/B that compares the CI RCCL against itself. Overlay ------- The swap for this family is discovery-based -- find, per-target RUNPATH preservation, and an unresolved-soname set gate -- ported from MAD#228, rather than aimed at the Primus SDK venv paths, which this base lacks. The libamd_smi soname bridge is kept and now runs over every directory the sweep finds, ordered after the copy and before the post-swap measurement so a .27-against-.26 skew is reconciled instead of failing the gate. Unlike MAD#228's dockerfile the Primus payload is gated on primus-cli, not examples/run_pretrain.sh: AMD-AGI/Primus#999 retired that file and run.sh drives primus-cli. v26.7.0 still carries both so MAD's gate passes today, but it would hard-fail on jax-maxtext-v26.8. PRIMUS_REF is pinned here. --skip-overlay-build is refused for this family: its bind-mount targets are Primus-image paths absent from this base, and `docker run -v` creates a missing source as an empty directory rather than refusing, so the run would look healthy while measuring the image's own RCCL. Scheduling ---------- Its own job rather than a third matrix entry: the weekly gate has to be a job-level `if:`, where the matrix context is unavailable, and it is fixed at 4 nodes while the matrix entries take madengine_nodes. WORK_DIR already carries inputs.workload, so they do not collide. Not gated on the matrix job -- a 70B failure should not hide a 405B result. time_limit is 05:00:00 for the A/B pair, against MAD#228's 12:00:00 for a single run: a 4-node exclusive request with a 12 h limit backfills badly, the measured run is ~17 min of training (MAD#228 reports 17.40 s/step) plus prep, and 300 min stays inside the step's existing 330 so the script still reports. The existing 360/330 budget needed no change. Risk to the existing workloads ------------------------------ MAD_REF moves 07ecef61 -> 7acfac7c, which is exactly one commit, ROCm/MAD#279, repointing tools/fetch_primus.sh at a ref that still exists. It touches benchmark/jax_maxtext/README.md, docker/primus_maxdiffusion.ubuntu.amd.Dockerfile and tools/fetch_primus.sh -- nothing under scripts/primus_megatron-lm. I rendered both A/B phase manifests for llama-3.1-70b, llama-4-scout and gpt-oss-120b against this branch and its base: byte-identical, all six. Also verified: the 405B manifest matches MAD#228's validated template field for field; the baseline and candidate manifests differ only in image and dockerfile; both generated Dockerfiles parse under `bash -n`, with the bridge ordered after the copy and before the gate. Not yet run on hardware. The first run has to establish whether the TheRock kpack librccl loads on this base at all, and whether its libamd_smi soname needs the bridge here. JIRA ID AICOMNET-420
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
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
examples/run_*launchers, including the Kubernetes submission client, and migrate packaged examples, benchmarks, hybrid tools, MaxDiffusion, and ODC to the root./primus-clientry point.runner/helpers/launch/slurm_pretrain.shfor callers that still use theEXP/NNODES/DATA_PATHcontract, with child exit-code propagation and Spur-compatiblesrunflag normalization.runner/: writable cache fallback, dynamic Slurm GPU count, BNXT tar mounting, MLflow/Databricks and UEP passthrough, andPRIMUS_SKIP_PIPhandling.--split_seedso Megatron's--seedremains the training RNG.site-packageswith an in-process ROCm patch that narrows unsafe Triton autotune configs while preserving cache hits.Behavior removal
examples/run_k8s_pretrain.shand its documentation are removed. This repository no longer ships a Kubernetes workload-submission client; Kubernetes deployments can still run./primus-cli directin each pod with consistent distributed environment variables.Review fixes
primus-clifailures throughteein the shared Slurm helper, ODC launcher, and auto-benchmark runner.HF_DATASETS_CACHE/HF_HOMEdefaults.mainto the new launcher path.