Repository navigation
fix(jax): point fetch_primus.sh at a ref that still exists - #279
Merged
Merged
Conversation
The default `PRIMUS_REF=jax-maxtext-v26.7` names a branch upstream has since
deleted, so the script fails at clone:
fatal: Remote branch jax-maxtext-v26.7 not found in upstream origin
Nothing downstream says so in those terms: with no checkout at scripts/Primus,
`get_models_json.py` reports zero JAX models and names this script, which reads
as "no MaxText models exist" rather than "its default ref is gone".
Switch the default to the `v26.7.0` tag, which carries the same release and
cannot move or disappear, and update the two places that named the branch.
Verified end to end: `tools/fetch_primus.sh` checks out 2631e68d and reports
47 maxtext and 6 maxdiffusion configs, including both llama3.1_405B yamls.
i-kosarev
requested review from
Cemberk,
amathews-amd,
coketaste,
gargrahul,
itej89,
lcskrishna and
leconcio
as code owners
October 9, 2026 10:22
Contributor
There was a problem hiding this comment.
🟢 Approval recommended
The verified tag exists at the reported commit, and all affected references are consistent.
0 open findings
What changed in this PR
Updates the Primus default ref to a valid release tag, restoring clean-tree JAX model discovery.
Changes:
- Defaults Primus fetching to
v26.7.0. - Updates related documentation and Docker build guidance.
| File | Description |
|---|---|
tools/fetch_primus.sh |
Uses the existing v26.7.0 tag. |
docker/primus_maxdiffusion.ubuntu.amd.Dockerfile |
Updates the missing-backend guidance. |
benchmark/jax_maxtext/README.md |
Documents the new default tag. |
🧠 Review effort: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
i-kosarev
requested review from
mkuznet1
and removed request for
amathews-amd,
coketaste,
gargrahul,
itej89,
lcskrishna and
leconcio
October 9, 2026 10:42
This was referenced Oct 9, 2026
Merged
i-kosarev
added a commit
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.
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
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.
Motivation
tools/fetch_primus.shdefaults toPRIMUS_REF=jax-maxtext-v26.7. That branch has been deleted upstream, so the script now fails at clone:Nothing downstream reports it in those terms. With no checkout at
scripts/Primus,get_models_json.pyfinds zero JAX models and points at this script — which reads as "no MaxText models exist" rather than "its default ref is gone". Anyone enabling a MaxText workload from a clean tree hits this first.Technical Details
v26.7.0tag. It carries the same release as the branch did, matchesdocker/primus_maxtext'srocm/jax-training:maxtext-v26.7base, and being a tag it cannot move or be deleted the same way.benchmark/jax_maxtext/README.mdand the error string indocker/primus_maxdiffusion.ubuntu.amd.Dockerfile.PRIMUS_REFremains the override, so anyone who wantsjax-maxtext-v26.8(which exists, alongside arocm/jax-training:maxtext-v26.8image published 2026-10-06) can still ask for it. Moving the whole default pairing to v26.8 felt like a bigger call than fixing a broken default, so this PR does not make it — say the word and I will.Test Plan
Run the script from a clean tree and confirm it checks out and discovery has configs to find.
Test Result
The old default still fails the same way, which is what this PR is about.
bash -n tools/fetch_primus.shis clean.Submission Checklist