Primus v26.6 - #239
Open
coketaste wants to merge 12 commits into
Open
Primus v26.6#239coketaste wants to merge 12 commits into
coketaste wants to merge 12 commits into
Conversation
Bump the Primus integration from v26.5 to v26.6.
- docker/primus.ubuntu.amd.Dockerfile: BASE_DOCKER rocm/primus:v26.5 -> v26.6.
Note the v26.6 image is not on Docker Hub yet (latest published is v26.5.1),
so builds will fail with 'manifest unknown' until it is pushed.
- scripts/Primus: pin 30cf451 -> 2aa05ea, the head of release/v26.6. Record
branch = release/v26.6 in .gitmodules so the pin no longer drifts onto main.
Two fixes the bump makes user-visible:
- scripts/primus_train/run.sh: derive BACKEND from the config's
modules.{pre,post}_trainer.framework instead of the launcher directory.
prepare_experiment.py compares BACKEND against that field and aborts on
mismatch, and the directory name is not always the framework: v26.6 adds
maxdiffusion and nemo_automodel (which fell through to megatron), while
diffusion and moe_package were already wrong at v26.5. Path inference is
kept as a fallback, and MaxText/MaxDiffusion keep their exact casing because
run_pretrain.sh string-matches those literals.
- scripts/primus_train/get_models_json.py: pass recursive=True so ** spans
directories. Without it, configs nested one level deeper were skipped,
hiding all four nemo_automodel configs. Discovery goes 414 -> 438.
- benchmark/primus/README.md: correct the BASE_DOCKER example (still v26.4) and
the submodule-init snippet (git submodule update only runs at the worktree
toplevel); drop the eight tags that no longer resolve (Zebra-Llama removed
upstream, TorchTitan deepseek_v3_16b split into BF16/FP8); regenerate the
Megatron and TorchTitan tables from the pinned checkout and add MI325X; add a
summary table for the backends that are discovered but not tabulated.
Verified: BACKEND resolution matches the declared framework for all 444 configs
in the checkout; madengine discover reports 439 entries; all 355 README tags
resolve to discovered models; the Dockerfile builds and bakes the v26.6 tree.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The v26.6 table regeneration added a separate Precision column; the v26.5 README instead folded precision into a single Model column. Restore that two-column shape (`| Model | Tag |`) across all 6 tables: underscores in the model name become spaces and the precision (when present) is appended to the end of the Model cell. Tag values are untouched. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
v26.5 capitalized display names (Llama, DeepSeek, Mixtral, ...); apply the same first-letter capitalization mechanically across all 355 v26.6 rows so the Model column isn't all-lowercase. Tag values are untouched. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
coketaste
requested review from
Cemberk,
amathews-amd,
gargrahul,
itej89,
lcskrishna and
leconcio
as code owners
August 28, 2026 21:40
Contributor
There was a problem hiding this comment.
Pull request overview
This PR bumps the Primus integration from v26.5 to v26.6, updating the base image reference, the Primus submodule tracking metadata, and the MAD-side Primus training/discovery utilities so they work correctly with the newer Primus launcher/config layout.
Changes:
- Update Primus training launcher selection by deriving
BACKENDfrom the config’s declared framework (with a path-based fallback). - Fix Primus config discovery to actually recurse through
**config directories. - Bump the Primus base Docker image tag to v26.6 and refresh the benchmark README to match the pinned Primus release content.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| scripts/primus_train/run.sh | Derives BACKEND from YAML framework fields (with fallback) to avoid mismatches as Primus adds/changes launchers. |
| scripts/primus_train/get_models_json.py | Enables recursive globbing so nested configs are discovered. |
| docker/primus.ubuntu.amd.Dockerfile | Updates the default Primus base image tag to v26.6. |
| benchmark/primus/README.md | Updates submodule init instructions, base image example, and regenerates supported-model tables for v26.6. |
| .gitmodules | Records the intended submodule branch (release/v26.6) to avoid drift when using --remote. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
# Conflicts: # scripts/Primus # scripts/primus_train/run.sh
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
The updated benchmark/primus/README.md currently documents JAX backends (maxtext/maxdiffusion) as primus_train/* tags discovered by scripts/primus_train/get_models_json.py, which conflicts with the repository’s actual discovery layout.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (1)
benchmark/primus/README.md:129
- This implies all runnable configs are discovered by
scripts/primus_train/get_models_json.py, but JAX backends are discovered by separate scripts (scripts/jax-maxtext/get_models_json.pyandscripts/jax-maxdiffusion/get_models_json.py) and use different dockerfiles. Clarifying this avoids confusion when users don’t seemaxtext/maxdiffusionconfigs underprimus_train/*.
Tags are discovered from the Primus submodule at run time by
`scripts/primus_train/get_models_json.py`, so any config present in the pinned Primus
checkout is runnable even if it is not listed below.
- Files reviewed: 6/6 changed files
- Comments generated: 1
- Review effort level: Lite
Restrict core.* to core.[0-9]* in .dockerignore so it only matches actual core dumps, and align .gitignore's env dir pattern with it (.jax-*_env/ -> .*_env/) to also cover .primus_train_*_env/ dumps. Also drop stale perf_*.json negations that were never tracked and just left run output showing as untracked clutter.
Contributor
There was a problem hiding this comment.
🔵 Needs a closer look
The updated run.sh YAML probe can terminate the script under set -e (preventing the intended fallback), and the README currently describes backend discovery in a way that conflicts with the repository’s dedicated JAX discovery scripts.
Review details
Suppressed comments (5)
Previously missed (3) — in code that hasn't changed since the last review.
scripts/primus_train/run.sh:76
- Because this script runs with
set -e, thepython3 -c ...inside the command substitution will cause the whole wrapper to exit if PyYAML is missing or the YAML parse fails. That prevents the intended fallback-to-path inference from ever running; force a zero exit status for the probe so the fallback can take over.
benchmark/primus/README.md:65 git submodule updatecan be run from any directory within the MAD worktree (Git resolves the repo root automatically). The current wording implies it fails outside the repo root, which is misleading.
benchmark/primus/README.md:125- Some Primus configs can be nested under additional subdirectories (e.g.,
.../configs/<GPU_ARCH>/diffusion/...). Saying the YAML lives directly underexamples/<backend>/configs/<GPU_ARCH>/is inaccurate and can confuse users when they try to locate configs.
benchmark/primus/README.md:123
- This README section says
maxtextandmaxdiffusionare validprimus_train/<backend>...values, butscripts/primus_train/get_models_json.pyexplicitly skips those launchers to avoid duplicating the dedicated JAX discoveries (scripts/jax-maxtext/scripts/jax-maxdiffusion). This mismatch will lead users to tags that are never discovered.
- `<backend>` is the Primus launcher: `megatron`, `torchtitan`, `megatron_bridge`, `maxtext`,
`maxdiffusion`, `nemo_automodel`, `diffusion`, or `moe_package`
benchmark/primus/README.md:129
- This paragraph implies all Primus launcher tags are discovered via
scripts/primus_train/get_models_json.py, but that script intentionally excludes the JAX backends (MaxText/MaxDiffusion). Consider clarifying that JAX backends are discovered via their dedicated scripts so users know where those tags come from.
Tags are discovered from the Primus submodule at run time by
`scripts/primus_train/get_models_json.py`, so any config present in the pinned Primus
checkout is runnable even if it is not listed below.
- Files reviewed: 7/8 changed files
- Comments generated: 0 new
- Review effort level: Lite
maxtext/maxdiffusion are excluded from get_models_json.py's discovery and run via dedicated jax-maxtext/jax-maxdiffusion benchmarks instead, so listing them in the primus_train tag convention was misleading. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
scripts/primus_train/run.sh runs with set -e, and the new YAML-parsing command substitution can abort the script (silently, due to stderr redirection) instead of falling back when parsing dependencies/configs are unavailable.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
- Files reviewed: 7/8 changed files
- Comments generated: 1
- Review effort level: Lite
Comment on lines
+68
to
+76
| framework="$(cd "$PRIMUS_ROOT" && python3 -c ' | ||
| import sys, yaml | ||
| mods = (yaml.safe_load(open(sys.argv[1])) or {}).get("modules") or {} | ||
| for key in ("pre_trainer", "post_trainer"): | ||
| fw = (mods.get(key) or {}).get("framework") | ||
| if fw: | ||
| print(fw) | ||
| break | ||
| ' "$EXP" 2>/dev/null)" |
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
The updated Primus README contains a couple of now-inaccurate discovery/tagging statements that can mislead users when locating configs and selecting the correct tag suite.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (1)
benchmark/primus/README.md:597
- This paragraph says all of these backends are discovered/runnable via the same
primus_train/tags, but MaxText and MaxDiffusion are explicitly excluded from primus_train discovery and instead come from thejax-maxtext/jax-maxdiffusionsuites. Clarifying the tag prefix difference will prevent users from trying non-existentprimus_train/maxtext_*orprimus_train/maxdiffusion_*tags.
These launchers are discovered and runnable through the same `primus_train/` tags, but are not
tabulated here. Use `madengine discover --tags <backend>` for the full list, or browse the configs
in the Primus repository.
- Files reviewed: 7/8 changed files
- Comments generated: 2
- Review effort level: Lite
| - `<backend>` is the Primus launcher: `megatron`, `torchtitan`, `megatron_bridge`, `nemo_automodel`, | ||
| `diffusion`, or `moe_package` | ||
| - `<GPU_ARCH>` is the target accelerator (`MI300X`, `MI325X`, or `MI355X`) | ||
| - `<MODEL_CONFIG>` matches the YAML filename under `examples/<backend>/configs/<GPU_ARCH>/` |
| # which hides all of the nemo_automodel backend added in Primus v26.6. | ||
| for yaml_path in sorted(glob.glob(CONFIGS_GLOB, recursive=True)): | ||
| rel_path = os.path.relpath(yaml_path, PRIMUS_ROOT) | ||
| # Path shape: examples/<launcher>/configs/<arch>/<file>.yaml |
scripts/primus_train/run.sh launched everything through
examples/run_pretrain.sh, which never loads Primus' runner/helpers/envs/
layer (base_env.sh + <GPU_MODEL>.sh). That layer is what the documented
standalone commands get via primus-cli, so a MAD run and a standalone run
of the same config were not measuring the same configuration:
HSA_NO_SCRATCH_RECLAIM defaulted to 0 instead of 1, and
NVTE_CK_IS_V3_ATOMIC_FP32 was never set at all. Apply the equivalent
architecture-aware environment in the wrapper, keyed on
MAD_SYSTEM_GPU_ARCHITECTURE / MAD_SYSTEM_GPU_PRODUCT_NAME:
HSA_NO_SCRATCH_RECLAIM=1 all non-JAX backends
NVTE_CK_IS_V3_ATOMIC_FP32=1 gfx942 (MI300X/MI325X)
PRIMUS_TURBO_ATTN_V3_ATOMIC_FP32=1 gfx942 (MI300X/MI325X)
RCCL_WARP_SPEED_AUTO=0 MI355X
NVTE_USE_CAST_TRANSPOSE_TRITON=0 *MXFP4* configs
Every value is ${VAR:-...}-guarded so an explicit override still wins,
and the effective values are echoed as "[primus_train] ..." so a run log
records what it actually ran with.
Also read the suite (pretrain vs posttrain) alongside the framework from
the config and launch post-training through primus-cli. run_pretrain.sh
hardcodes `train pretrain` and resolves prepare to
examples/<framework>/prepare.py, which does not exist for
megatron_bridge -- Qwen3-32B SFT/LoRA died in prepare_experiment.py
before training started. primus-cli's posttrain hooks install the bridge
requirements and convert checkpoints first.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
get_models_json.py hardcoded training_precision="bf16" on every discovered config, so the performance CSV reported BF16 for all of them regardless of the FP8/MXFP8/MXFP4 variant being run. Derive it from the precision token in the config name instead, longest token first so MXFP8 is not matched as FP8, and report "" (madengine's unknown convention) for configs that carry no token rather than guessing. extract_primus_perf.py had the same bug in its MFU fallback: the peak was doubled on a substring test for "fp8", which scored MXFP4 against the BF16 dense peak and overstated its MFU by 4x. Replace it with a shared multiplier table (FP8/MXFP8 2x, FP4/MXFP4 4x). The table is duplicated rather than imported because this script runs inside the training container, where madengine -- and therefore get_models_json -- is not importable. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
train_args.substitute_sdpa_with_flash_attn was passed with a "+" prefix, which tells Hydra to append a new key and errors out when the key already exists. It does exist in the current AMDiffusionBenchmark config, so Stable-Diffusion-XL failed before training started. Drop the prefix to override the existing value. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
scripts/primus_train/run.sh can exit early under set -e if the YAML parsing command substitution fails, preventing the intended fallback logic from running.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
- Files reviewed: 9/10 changed files
- Comments generated: 2
- Review effort level: Lite
| if fw: | ||
| print(fw, suite) | ||
| break | ||
| ' "$EXP" 2>/dev/null)" |
Comment on lines
+30
to
34
| in the container env), and precision is inferred from the log filename | ||
| (e.g. "...-FP8-pretrain.txt" vs "...-MXFP4-pretrain.txt"), which scales the | ||
| peak by _PEAK_MULTIPLIER. | ||
| If either can't be determined, model_flops_utilization is left blank | ||
| rather than guessed. |
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
Bump the Primus integration from v26.5 to v26.6.
docker/primus.ubuntu.amd.Dockerfile:BASE_DOCKERrocm/primus:v26.5 -> v26.6.Note the v26.6 image is not on Docker Hub yet (latest published is v26.5.1), so
builds will fail with "manifest unknown" until it is pushed.
scripts/Primus: pin30cf451->2aa05ea, the head ofrelease/v26.6. Recordbranch = release/v26.6in.gitmodulesso the pin no longer drifts ontomain.Two fixes the bump makes user-visible:
scripts/primus_train/run.sh: deriveBACKENDfrom the config'smodules.{pre,post}_trainer.frameworkinstead of the launcher directory.prepare_experiment.pycomparesBACKENDagainst that field and aborts onmismatch, and the directory name is not always the framework — v26.6 adds
maxdiffusionandnemo_automodel(which fell through tomegatron), whilediffusionandmoe_packagewere already wrong at v26.5. Path inference iskept as a fallback, and MaxText/MaxDiffusion keep their exact casing since
run_pretrain.shstring-matches those literals.scripts/primus_train/get_models_json.py: passrecursive=Trueso**spansdirectories. Without it, configs nested one level deeper were silently
skipped, hiding all four
nemo_automodelconfigs. Discovery goes 414 -> 438.benchmark/primus/README.md:BASE_DOCKERexample (still showed v26.4) and the submodule-initsnippet (
git submodule updateonly runs at the worktree toplevel).TorchTitan
deepseek_v3_16bsplit into BF16/FP8).add MI325X; add a summary table for backends that are discovered but not
tabulated.
| Model | Tag |table shape (fold precisioninto the Model cell) and capitalize the Model column's leading letter to
match prior formatting.
Verification
checkout.
madengine discoverreports 439 entries; all 355 README tags resolve todiscovered models.
Caveat
The v26.6 Primus image isn't published to Docker Hub yet (only v26.5.1 is
live), so image builds referencing
BASE_DOCKER=rocm/primus:v26.6will failuntil AMD-AGI pushes it.