Conversation
* initial v26.7 MAD change * Update docs * Use Primus MaxDiffusion throughput metrics in MAD extraction. Stop recomputing fps/frames from config; copy Primus per-device samples, frames, tokens, and TFLOP/s into madengine results. * Launch JAX MAD jobs with primus-cli and pin Primus to jax-maxtext-v26.7. examples/run_pretrain.sh is gone after Primus#999, so run.sh now calls primus-cli direct inside the MAD image. fetch_primus.sh defaults to the v26.7 branch that matches the maxtext-v26.7 stack. * Filter JAX configs by GPU product, not only gfx ISA. MI300X and MI325X are both gfx942, so madengine skip_gpu_arch cannot tell them apart and would run both Primus directories. Discovery now uses the rocminfo marketing name (overridable with JAX_HOST_DEVICE). * fix the fetching of Primus repo --------- Co-authored-by: Fuyuan Jing <Fuyuan.Jing@amd.com>
* Primus v26.6
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>
* README: merge Model/Precision columns to match v26.5 table pattern
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>
* README: capitalize leading letter of Model column
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>
* Tighten core dump and env dir ignore patterns
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.
* Fix primus README: exclude JAX launchers from primus_train tag list
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>
* Apply Primus perf env and route posttrain through primus-cli
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>
* Report actual training precision for Primus configs
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>
* Fix SDXL duplicate Hydra override for flash-attn substitution
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>
* docs: design ROCM-30673 diffusion perf fix
Co-authored-by: Cursor <cursoragent@cursor.com>
* docs: plan ROCM-30673 diffusion perf fix
Co-authored-by: Cursor <cursoragent@cursor.com>
* test: cover Primus diffusion environment defaults
Co-authored-by: Cursor <cursoragent@cursor.com>
* fix: align diffusion perf environment with Primus
Co-authored-by: Cursor <cursoragent@cursor.com>
* chore: keep agent plan/spec docs out of git
These superpowers artifacts are local working notes and should not ship in the repository.
Co-authored-by: Cursor <cursoragent@cursor.com>
* fix: overwrite diffusion perf CSV instead of appending
Stale NaN rows in perf_$MODEL.csv were ingested as Failed runs even when training succeeded.
Co-authored-by: Cursor <cursoragent@cursor.com>
* fix: parse Megatron-Bridge TPS for MAD posttrain results
Bridge logs often omit printed tokens/s/GPU, so extract_primus_perf now
derives TPS from elapsed time and batch size, and posttrain extraction
failures are no longer ignored.
Co-authored-by: Cursor <cursoragent@cursor.com>
* fix: parse Megatron-Bridge MODEL_TFLOP when throughput is omitted
SFT/LoRA logs print GPU utilization as MODEL_TFLOP/s/GPU instead of
throughput per GPU; use that only when no existing TFLOPS format matches.
Co-authored-by: Cursor <cursoragent@cursor.com>
* fix: align megatron pretrain NCCL PXN with primus-cli (ROCM-31034)
MAD left NCCL_PXN_DISABLE unset so run_pretrain.sh enabled PXN, unlike primus-cli/base_env.sh, which can drop GDN/megatron throughput ~10-15% vs QA.
Co-authored-by: Cursor <cursoragent@cursor.com>
* fix: force NVTE_USE_CAST_TRANSPOSE_TRITON=0 for MI355X MXFP4
Image/base_env default is 1, so ${VAR:-0} left MAD llama3.1_8B-MXFP4-pretrain on the slow Triton path vs the documented Primus CLI recipe.
Co-authored-by: Cursor <cursoragent@cursor.com>
---------
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
i-kosarev
requested review from
Cemberk,
amathews-amd,
coketaste,
gargrahul,
itej89,
lcskrishna and
leconcio
as code owners
September 21, 2026 23:50
Contributor
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Critical Primus launch and version-alignment blockers, along with additional discovery, override, fallback, and output issues, remain unresolved.
Get a fresh assessment by requesting another Copilot review.
Review effort: Lite
Findings: 2
Open (3)
What changed in this PR
This PR updates Primus-based PyTorch and JAX workflows, model discovery, performance extraction, containers, and documentation.
Changes:
- Migrates JAX launchers and metrics handling to
primus-cli. - Adds SKU-aware discovery, environment handling, and regression tests.
- Updates reporting tools, Docker images, submodule metadata, and documentation.
| File | Summary |
|---|---|
tools/fetch_primus.sh |
Primus checkout synchronization and validation |
scripts/pytorch_train/tests/test_perf_env.sh |
Performance environment regression tests |
scripts/pytorch_train/tests/test_overwrite_perf_csv.py |
CSV overwrite regression test |
scripts/pytorch_train/pytorch_benchmark_report.sh |
Performance reporting environment defaults |
scripts/pytorch_train/pytorch_benchmark_report.py |
Performance CSV output handling |
scripts/primus_train/tests/test_pretrain_env.sh |
Primus environment parity tests |
scripts/primus_train/tests/test_extract_primus_perf.py |
Primus metric extraction tests |
scripts/primus_train/run.sh |
Primus launch, detection, and environment handling |
scripts/primus_train/get_models_json.py |
Dynamic model and precision discovery |
scripts/primus_train/extract_primus_perf.py |
Primus performance metric extraction |
scripts/jax-maxtext/run.sh |
MaxText primus-cli launcher |
scripts/jax-maxtext/get_models_json.py |
MaxText SKU-aware discovery |
scripts/jax-maxdiffusion/run.sh |
MaxDiffusion primus-cli launcher |
scripts/jax-maxdiffusion/get_models_json.py |
MaxDiffusion SKU-aware discovery |
scripts/jax-maxdiffusion/extract_maxdiffusion_perf.py |
MaxDiffusion throughput extraction |
scripts/jax_host_device.py |
GPU SKU detection |
README.md |
Updated training workflow overview |
docker/primus.ubuntu.amd.Dockerfile |
Primus base image update |
docker/primus_maxtext.ubuntu.amd.Dockerfile |
MaxText image and CLI validation |
docker/primus_maxdiffusion.ubuntu.amd.Dockerfile |
MaxDiffusion image and CLI validation |
benchmark/primus/README.md |
Primus usage and model documentation |
benchmark/jax_maxtext/README.md |
JAX benchmark documentation |
.gitmodules |
Primus branch configuration |
.gitignore |
Generated artifact exclusions |
.dockerignore |
Docker build-context exclusions |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| bash "$PRIMUS_ROOT/runner/primus-cli" direct --log_file "$TRAIN_LOG" -- \ | ||
| train posttrain --config "$EXP" "${forward_args[@]}" | ||
| else | ||
| bash "$PRIMUS_ROOT/examples/run_pretrain.sh" "${forward_args[@]}" --job.dump_folder "$RUN_DIR/outputs" |
|
|
||
| PRIMUS_URL="${PRIMUS_URL:-https://github.com/AMD-AGI/Primus}" | ||
| PRIMUS_REF="${PRIMUS_REF:-main}" | ||
| PRIMUS_REF="${PRIMUS_REF:-jax-maxtext-v26.7}" |
| [submodule "scripts/Primus"] | ||
| path = scripts/Primus | ||
| url = https://github.com/AMD-AGI/Primus | ||
| branch = release/v26.6 |
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.
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist