Repository navigation
JAX/MaxText multi-node RCCL overlay - #228
Conversation
There was a problem hiding this comment.
Pull request overview
Enables multi-node execution for JAX/MaxText workloads and adds an RCCL FAULT_INJECTION A/B harness (ROCM-27881) that can pin and validate two RCCL builds differing only by -DFAULT_INJECTION=ON/OFF, including cluster/run tooling to reproduce and attribute the performance delta.
Changes:
- Add a pinned RCCL overlay Dockerfile that builds FI=ON and FI=OFF variants and gates that they differ.
- Add a SLURM
slurm_multiself-orchestrating runner to launch one container per node for multi-node MaxText runs. - Update MaxText run/report scripts to support multi-node topology and add a median steady-state
seconds_per_stepmetric for A/B comparisons.
Reviewed changes
Copilot reviewed 28 out of 28 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| scripts/jax-maxtext/run.sh | Adds RCCL variant selection/attestation hooks; wires JAX coordinator env for multi-node; adds 405B model mapping; adjusts performance marker behavior. |
| scripts/jax-maxtext/run_multinode_rccl.slurm | New SLURM multi-node orchestrator: per-node docker daemon prep, image load, watchdog, and container launch with required env/mounts for A/B. |
| scripts/jax-maxtext/rccl_verify_mapped.sh | New post-init verification script intended to prove which librccl is actually mapped by a running process. |
| scripts/jax-maxtext/rccl_variant_select.sh | New selector that installs the requested RCCL variant and logs on-disk verification markers + version banner behavior. |
| scripts/jax-maxtext/models.json | Adds multi-node-oriented model entries pointing at the RCCL overlay dockerfile prefix. |
| scripts/jax-maxtext/jax-maxtext_benchmark_setup.sh | Adds tokenizer download behavior for the new 405B model mapping. |
| scripts/jax-maxtext/jax-maxtext_benchmark_report.sh | Makes topology multi-node-aware; fails loudly on missing env/config; adds MAXTEXT_EXTRA_ARGS; adds 405B dispatch. |
| scripts/jax-maxtext/jax-maxtext_benchmark_report.py | Drops unused pandas; adds median steady-state step-time parsing/metric and stricter model handling. |
| scripts/jax-maxtext/install_rdma_core.sh | New rdma-core-from-source installer with a bnxt_re provider gate (Thor2-specific requirement). |
| scripts/jax-maxtext/env_scripts/qwen3_30b_a3b_env.sh | Stops clobbering NCCL_DEBUG; adds append-only XLA_EXTRA_FLAGS hook. |
| scripts/jax-maxtext/env_scripts/qwen3_14b_env.sh | Stops clobbering NCCL_DEBUG; adds append-only XLA_EXTRA_FLAGS hook. |
| scripts/jax-maxtext/env_scripts/mixtral_8x7b_env.sh | Stops clobbering NCCL_DEBUG; adds append-only XLA_EXTRA_FLAGS hook. |
| scripts/jax-maxtext/env_scripts/llama3.3_70b_env.sh | Stops clobbering NCCL_DEBUG; adds append-only XLA_EXTRA_FLAGS hook. |
| scripts/jax-maxtext/env_scripts/llama3_8b_env.sh | Stops clobbering NCCL_DEBUG; adds append-only XLA_EXTRA_FLAGS hook. |
| scripts/jax-maxtext/env_scripts/llama3_70b_env.sh | Stops clobbering NCCL_DEBUG; adds append-only XLA_EXTRA_FLAGS hook. |
| scripts/jax-maxtext/env_scripts/llama2_7b_env.sh | Stops clobbering NCCL_DEBUG; adds append-only XLA_EXTRA_FLAGS hook. |
| scripts/jax-maxtext/env_scripts/llama2_70b_env.sh | Stops clobbering NCCL_DEBUG; adds append-only XLA_EXTRA_FLAGS hook. |
| scripts/jax-maxtext/env_scripts/gfx950_qwen3_30b_a3b_env.sh | Stops clobbering NCCL_DEBUG; adds append-only XLA_EXTRA_FLAGS hook. |
| scripts/jax-maxtext/env_scripts/gfx950_qwen3_14b_env.sh | Stops clobbering NCCL_DEBUG; adds append-only XLA_EXTRA_FLAGS hook. |
| scripts/jax-maxtext/env_scripts/gfx950_mixtral_8x7b_env.sh | Stops clobbering NCCL_DEBUG; adds append-only XLA_EXTRA_FLAGS hook. |
| scripts/jax-maxtext/env_scripts/gfx950_llama3.3_70b_env.sh | Stops clobbering NCCL_DEBUG; adds append-only XLA_EXTRA_FLAGS hook. |
| scripts/jax-maxtext/env_scripts/gfx950_llama3.1_405b_env.sh | New gfx950 env script for the 405B multi-node config with XLA_AUTOTUNE_LEVEL/XLA_EXTRA_FLAGS hooks. |
| scripts/jax-maxtext/env_scripts/gfx950_llama3_8b_env.sh | Makes XLA_AUTOTUNE_LEVEL overridable; stops clobbering NCCL_DEBUG; adds XLA_EXTRA_FLAGS hook. |
| scripts/jax-maxtext/env_scripts/gfx950_llama3_70b_env.sh | Stops clobbering NCCL_DEBUG; adds append-only XLA_EXTRA_FLAGS hook. |
| scripts/jax-maxtext/env_scripts/gfx950_llama2_7b_env.sh | Stops clobbering NCCL_DEBUG; adds append-only XLA_EXTRA_FLAGS hook. |
| scripts/jax-maxtext/env_scripts/gfx950_llama2_70b_env.sh | Stops clobbering NCCL_DEBUG; adds append-only XLA_EXTRA_FLAGS hook. |
| docker/jax_maxtext_rccl_overlay.ubuntu.amd.Dockerfile | New overlay Dockerfile: pinned RCCL commit, dual FI builds with a “must differ” build gate, plus rdma-core 63.0 install option. |
| benchmark/jax_maxtext/RCCL_FAULT_INJECTION_AB.md | New catalog/recipe page documenting the A/B method, attribution gates, and cluster prerequisites. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 28 out of 28 changed files in this pull request and generated 2 comments.
Suppressed comments (1)
Previously missed (1) — in code that hasn't changed since the last review.
scripts/jax-maxtext/rccl_variant_select.sh:141
- This script notes that exporting NCCL_DEBUG here won’t reach the training process because the selector runs as a child of run.sh, but then it still exports NCCL_DEBUG/NCCL_DEBUG_SUBSYS in attestation mode. That’s misleading (and, in practice, ineffective) and makes it harder to reason about where the attestation env is actually coming from.
if [[ "${RCCL_FI_ATTEST:-0}" == "1" ]]; then
export NCCL_DEBUG=INFO
export NCCL_DEBUG_SUBSYS="${NCCL_DEBUG_SUBSYS:-INIT,NET}"
echo "ATTESTATION MODE: NCCL_DEBUG=INFO NCCL_DEBUG_SUBSYS=$NCCL_DEBUG_SUBSYS (timings from this run are NOT comparable)"
fi
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 28 out of 28 changed files in this pull request and generated 1 comment.
Suppressed comments (1)
scripts/jax-maxtext/run_multinode_rccl.slurm:242
|| trueat the end of thedocker run | teepipeline forces the overall command to succeed and (in practice) resetsPIPESTATUS, sorc=${PIPESTATUS[0]}becomes 0 even when the container failed. That can mask failures and let the SLURM step look successful.
Remove the trailing || true so the job correctly fails and PIPESTATUS[0] remains meaningful.
bash -lc "bash ./run.sh --model_repo $MODEL_REPO --quantization '$QUANTIZATION'" 2>&1 | tee -a "$RLOG" || true
rc=${PIPESTATUS[0]}
kill "$WATCHDOG_PID" 2>/dev/null || true
echo "[$HOST rank=$RANK] container exited rc=$rc"
exit "$rc"
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 28 out of 28 changed files in this pull request and generated 2 comments.
Suppressed comments (1)
Previously missed (1) — in code that hasn't changed since the last review.
scripts/jax-maxtext/run.sh:151
jax-maxtext_benchmark_setup.shfailures (e.g., HF auth / tokenizer download) are not checked here, andrun.shdoes not run withset -e, so the script can continue into the report step and potentially emit misleading output after a setup failure. Propagate the setup step exit status similarly to the report step.
bash ./jax-maxtext_benchmark_setup.sh -m "$model"
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 28 out of 28 changed files in this pull request and generated no new comments.
Suppressed comments (3)
Previously missed (3) — in code that hasn't changed since the last review.
scripts/jax-maxtext/run_multinode_rccl.slurm:69
- BUILDS_DIR defaults to a user-specific path (/home/ilkosare/...), which will fail on most clusters and makes the slurm_multi harness non-portable. Since the script already documents MAD_DOCKER_BUILDS as a required manifest-provided contract, fail fast when it’s missing instead of silently falling back to a personal directory.
BUILDS_DIR="${MAD_DOCKER_BUILDS:-/home/ilkosare/MADstorage/docker_builds}"
IMAGE_TAR="${BUILDS_DIR}/$(echo "$DOCKER_IMAGE_NAME" | tr ':/' '__').tar"
scripts/jax-maxtext/run_multinode_rccl.slurm:263
- JAX_CACHE_HOST also defaults to a user-specific path (/home/ilkosare/...). This makes the run harness non-portable and can fail in environments where that path doesn’t exist. Prefer a generic default (e.g. under the submit dir) so the script works out of the box, while still allowing the manifest/env to override it.
export JAX_CACHE_HOST="${JAX_CACHE_HOST:-/home/ilkosare/MADstorage/ilk-jaxmt/cache/jax}"
mkdir -p "$JAX_CACHE_HOST"; chmod -R o+rwX "$JAX_CACHE_HOST" 2>/dev/null || true
scripts/jax-maxtext/rccl_variant_select.sh:143
- The comments here state that nothing is exported from this selector because it runs as a child of run.sh, but the attestation block still exports NCCL_DEBUG/NCCL_DEBUG_SUBSYS. Those exports cannot affect the training process (child env doesn’t propagate), so the block is misleading and contradicts the stated design. Either drop these exports or rewrite the comment to clarify they only affect the selector process itself.
# ----------------------------------------------------- verification 2 of 4 ---
# The RCCL version banner. NOTE: this script runs as a CHILD of run.sh, so exporting
# NCCL_DEBUG here would die with it - the env scripts set `NCCL_DEBUG="${NCCL_DEBUG:-VERSION}"`
# instead, which is what actually puts the banner in the log. Nothing is exported here on
# purpose; the caller is responsible for asserting the banner is present afterwards.
# Verification 3 of 4: the mapped-library check, which can only run once the
# training process exists - see rccl_verify_mapped.sh, invoked after JAX initialises.
#
# Verification 4 of 4: the shim announces itself at runtime with
# NET/IB ops-fault: installed shims on context 0x...
# (INFO, NCCL_NET subsystem) once per ibv context. That is the strongest possible
# evidence - it proves the shim actually installed, not merely that it was compiled
# in. It requires NCCL_DEBUG=INFO + NCCL_DEBUG_SUBSYS covering NET, which is verbose
# enough to perturb step times, so it belongs in a SEPARATE short attestation run
# per variant, NOT in the timed A/B runs. Measuring and attesting in the same run
# would mean reporting numbers taken under logging overhead.
if [[ "${RCCL_FI_ATTEST:-0}" == "1" ]]; then
export NCCL_DEBUG=INFO
export NCCL_DEBUG_SUBSYS="${NCCL_DEBUG_SUBSYS:-INIT,NET}"
echo "ATTESTATION MODE: NCCL_DEBUG=INFO NCCL_DEBUG_SUBSYS=$NCCL_DEBUG_SUBSYS (timings from this run are NOT comparable)"
fi
154eb05 to
b871803
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 12 out of 12 changed files in this pull request and generated 3 comments.
Suppressed comments (1)
scripts/jax-maxtext/run_multinode_rccl.slurm:252
JAX_CACHE_HOSTdefaults to a hard-coded user-specific path (/home/ilkosare/...). Since this directory is mounted into every container, a non-portable default risks either failing or writing caches into an unexpected location on other clusters. Consider defaulting relative to the submit/output directory (or requiring it explicitly) to keep the script self-contained.
export JAX_CACHE_HOST="${JAX_CACHE_HOST:-/home/ilkosare/MADstorage/ilk-jaxmt/cache/jax}"
mkdir -p "$JAX_CACHE_HOST"; chmod -R o+rwX "$JAX_CACHE_HOST" 2>/dev/null || true
b871803 to
736b44c
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 12 out of 12 changed files in this pull request and generated 3 comments.
Suppressed comments (1)
scripts/jax-maxtext/run_multinode_rccl.slurm:67
BUILDS_DIRdefaults to a user-specific absolute path (/home/ilkosare/...). That makes the script non-portable and can lead to confusing failures if MAD_DOCKER_BUILDS isn’t set. Prefer requiring MAD_DOCKER_BUILDS explicitly (it’s already part of the stated env contract) and derive IMAGE_TAR from it.
esac
done
BUILDS_DIR="${MAD_DOCKER_BUILDS:-/home/ilkosare/MADstorage/docker_builds}"
IMAGE_TAR="${BUILDS_DIR}/$(echo "$DOCKER_IMAGE_NAME" | tr ':/' '__').tar"
736b44c to
8819eab
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 11 out of 11 changed files in this pull request and generated 1 comment.
Suppressed comments (1)
.claude/skills/mad-slurm-multinode/assets/manifests/jax_maxtext_llama-3.1-8b.template.json:81
- This reusable template hard-codes a Thor2-only scheduler workaround. On CX7/AINIC or any SLURM partition that uses GPU GRES, suppressing
--gpus-per-nodecan under-request GPUs or leave them unavailable to the job. Default this to the normal directive and let Thor2 users opt out, as documented inreferences/cluster-types.md:77-83.
"_comment_skip_gpus": "This cluster's SLURM rejects the generated --gpus-per-node directive; without this the job is refused before launch.",
"skip_gpus_directive": true,
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 11 out of 11 changed files in this pull request and generated no new comments.
Suppressed comments (1)
.claude/skills/mad-slurm-multinode/assets/manifests/jax_maxtext_llama-3.1-8b.template.json:64
- The reusable template cannot select the advertised FAULT_INJECTION A/B arm. Although the image tag asks for
<FILL_BUILD_VARIANT>,docker_build_argomitsRCCL_CMAKE_OPTIONS, so a fully filled template always builds the commit's default and can label that build as ON or OFF incorrectly. Add the same explicit build argument used byexample_manifest_multinode_rccl.json.
"RDMA_CORE_VERSION": "<FILL_RDMA_CORE empty to keep the base image rdma-core; 63.0 is the Broadcom-Thor2 bnxt_re EFAULT fix, see references/cluster-types.md>"
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 11 out of 11 changed files in this pull request and generated no new comments.
Suppressed comments (2)
.claude/skills/mad-slurm-multinode/assets/manifests/jax_maxtext_llama-3.1-8b.template.json:5
BUILD_GPU_TARGETSbelow explicitly allowsgfx942, but this cache key always saysgfx950. A gfx942 build can consequently reuse an existing gfx950 image/tar with the same commit and variant instead of rebuilding for the requested architecture. Include the target placeholder in the tag so architecture-specific artifacts cannot collide.
"docker_image": "jax-maxtext-rccl:<FILL_RCCL_COMMIT>-<FILL_BUILD_VARIANT>-gfx950",
.claude/skills/mad-slurm-multinode/assets/manifests/jax_maxtext_llama-3.1-8b.template.json:65
- The template names a
<FILL_BUILD_VARIANT>in the image tag but never passes that variant to the Docker build. Filling this template for FI=ON and FI=OFF therefore produces the Dockerfile's same default build under two tags, so it cannot perform the advertised A/B comparison. AddRCCL_CMAKE_OPTIONShere as a required placeholder, as the worked example does.
"docker_build_arg": {
"BASE_DOCKER": "rocm/jax-training:maxtext-v26.5",
"RCCL_COMMIT": "<FILL_RCCL_COMMIT>",
"BUILD_GPU_TARGETS": "<FILL_GPU_TARGETS gfx942=MI300X/MI325X, gfx950=MI350X/MI355X>",
"RDMA_CORE_VERSION": "<FILL_RDMA_CORE empty to keep the base image rdma-core; 63.0 is the Broadcom-Thor2 bnxt_re EFAULT fix, see references/cluster-types.md>"
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 11 out of 11 changed files in this pull request and generated 2 comments.
Suppressed comments (1)
scripts/jax-maxtext/run.sh:93
- An explicitly supplied
JAX_COORDINATOR_PORTis overwritten with 29500 wheneverMASTER_PORTis absent. Existing standalone launchers use custom JAX ports, so routing them through this wrapper would make ranks connect to the wrong endpoint. Retain the native JAX value as the fallback.
export JAX_COORDINATOR_PORT="${MASTER_PORT:-29500}"
Adds docker/primus_maxtext_rccl_overlay.ubuntu.amd.Dockerfile: the Primus JAX/MaxText image with a candidate RCCL built from source, so a MaxText workload can be run against any commit of ROCm/rocm-systems and any RCCL build options: --build-arg RCCL_COMMIT=<sha> --build-arg RCCL_CMAKE_OPTIONS=<flags> Comparing two RCCL builds is then two images and the same model card twice - no A/B driver, no variant switch, nothing added to the model path. It is a drop-in for docker/primus_maxtext.ubuntu.amd.Dockerfile: same PRIMUS_ROOT, MAXTEXT_PATH and mad.launcher=primus label, so the same discovered model card and the same scripts/jax-maxtext/run.sh work against either image and the only difference between them is librccl. The Primus payload is duplicated rather than layered because primus_maxtext is not published under a stable tag, and FROM-ing a locally built one would make this build order-dependent; the two `test -f` gates fail loudly if it drifts. The name matters functionally, not just for grouping: madengine gives repo-root build context to any dockerfile whose path contains "primus" (DockerBuilder.get_context_path), which is what lets `COPY scripts/Primus/` resolve. Under ./docker context it would not. This is the JAX/MaxText counterpart of primus_megatron_train_rccl_overlay and follows its structure; what differs is the base image, the Primus payload (rocm/primus:* already ships Primus, rocm/jax-training:maxtext-* does not) and the librccl splice - there is no torch/lib here, so it is one find-based sweep over every librccl on disk. Beyond the shared undefined-symbol gate, the build fails when RCCL_CMAKE_OPTIONS names FAULT_INJECTION but did not take effect. cmake only warns about an unrecognised -D cache variable, so on a commit that does not define the option both arms build the same library while the provenance file claims they differ - a silent no-op A/B. The option is located as a whole token and read with cmake's own truth rules, so TRUE/YES/:BOOL=ON are not mistaken for a request for OFF, and a positive-control string guards the marker count against `strings` having read nothing at all.
extract_maxtext_perf.py reported only Tokens/s/device and TFLOP/s/device. For an A/B of the collective library the quantity of interest is step latency, and neither existing metric is a good proxy for it: both are means over a 10-sample trailing window, so one straggler step - a checkpoint flush, an eval, a host hiccup - moves them by percent while the effect under test can itself be a few percent. Measured on a synthetic log with a single 12.5 s step among 0.85 s steps, the mean TFLOP/s reads 380 against a true ~421. Adds a third row, seconds_per_step, as the MEDIAN of the same window. The two throughput metrics keep their existing mean so their published numbers stay comparable across releases; only the new metric is robust by construction. The regex is anchored on "completed step:" so an unrelated line that happens to carry "seconds:" - a checkpoint timer, a wall-time summary - cannot contribute a sample. The row is omitted rather than written blank when a log carries no step timings: an empty performance cell reads downstream as a measured value, and 0 s/step would look like an impossibly good result instead of a missing one. Checked against synthetic logs covering: a straggler inside the window, unrelated "seconds:" lines inside the window (ignored, median unchanged), step lines with no seconds field (2 rows, exit 0), and an empty log (exit 1, unchanged behaviour).
…multinode
The campaign write-up does not belong in the repo, so what is durable goes where the next
person will actually look: the skill that is read before a run.
references/gotchas.md gains a jax_maxtext section - the JAX coordinator contract and why a
loopback result must be rejected, one container per node and why NPROC_PER_NODE must not
stand in for the GPU count, a stale primus_perf_output.csv outliving the run that failed
to overwrite it, the harmless gated-repo tokenizer failure, that madengine scrapes the
FIRST performance match, and that step-time medians only compare at equal `steps`.
Four cross-cutting entries too, each of which cost a debugging session: that the SLURM
multi-node profile sets NCCL_IB_DISABLE=1 for you and an A/B taken over sockets is not
just slow but insensitive (+192 % collapses to +0.59 %); that built_models.env_vars is not
read on the docker path; that the fd limit and /dev/infiniband belong in
additional_docker_run_options, with the silent-stall and TCP-fallback symptoms they
produce; and that a verification which cannot fail is worse than none.
The XLA autotune entry is deliberately not a recommendation. Primus now owns XLA_FLAGS
for MaxText and defaults autotune_level=4 because >=4 is what keeps fp8 numerically
stable, while the old MAD-native measurement said levels 1-2 would not finish compiling
at 32 ranks. Those were different stacks and level 4 was never timed at 32 ranks, so the
entry records both facts and the tension rather than advising a value that could trade an
fp8 correctness guarantee for compile time.
assets/manifests/ gains a primus_maxtext template, since Step 5 picks from that directory
and without one the workload cannot be run through the skill at all. It uses
launcher: primus with distributed.primus.{config_path,backend}, which is madengine's
supported path for JAX/MaxText. Checked with the skill's own validate_manifest.sh: the
same 3 PASS / 2 WARN / 3 FAIL profile as the existing primus_llama-3.1-70b template, the
FAILs being unfilled placeholders.
run_pretrain.sh appends whatever it does not recognise to the training command, and --config_path is not a MaxText flag. Stripping it only when PRIMUS_CONFIG_PATH was unset left it in the forwarded args on exactly the path that sets it: madengine's `launcher: primus` exports PRIMUS_CONFIG_PATH, and every discovered model card carries `args: --config_path <rel>`. So the flag reached MaxText and it died with ValueError: 'config_path' not in 'emb_dim', ... before step 1 on every multi-node run. Strip it always, and let the launcher's value win when both are present: when madengine sets PRIMUS_CONFIG_PATH that is the authoritative config for the deployment, and the card's args are just how the same path reaches a local run. Co-Authored-By: Claude <noreply@anthropic.com>
The overlay installs the candidate librccl over every librccl on disk. Copying the bytes in verbatim also replaces the destination's RUNPATH: a librccl inside the ROCm pip wheels resolves its DT_NEEDED libamd_smi.so.26 through a RUNPATH reaching the sibling wheel ($ORIGIN/../../_rocm_sdk_core/lib), while the candidate is built for /opt/rccl and carries $ORIGIN/../lib. The result is a librccl that cannot resolve its own dependencies, and `import torch` dies in rocm_sdk.initialize_process with libamd_smi.so.26: cannot open shared object file before a single step runs. Copy the bytes, then restore the destination's own RUNPATH with patchelf. The ldd gate is what would have caught this at build time instead of four minutes into a two-node job. It compares each target before and after the swap and fails only on an INCREASE. An absolute "must resolve everything" test would fail a good build: on this base the pristine _rocm_sdk_devel copy already reports 6 unresolved deps, because it is only ever dlopen'd into a namespace where those libraries are already loaded. What must not happen is the swap making linkage worse than it found it. Co-Authored-By: Claude <noreply@anthropic.com>
There was a problem hiding this comment.
🟡 Changes recommended
Multi-node rendezvous setup, RCCL build arguments, metric gating, and the advertised 405B fix are incomplete.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
- Files reviewed: 5/5 changed files
- Comments generated: 5
- Review effort level: Balanced
…swap gate Counting `not found` lines before and after the librccl swap passes a build in which the swap removed one missing soname and introduced a different one — the count is unchanged and the image is broken in exactly the way the gate exists to catch. Record the set of unresolved sonames per target instead and fail when the post-swap set contains any name that was not already there. An absolute "must resolve everything" test is still not usable here: the pristine _rocm_sdk_devel copy already reports 6 unresolved deps on this base, because it is only ever dlopen'd into a namespace where those libraries are loaded. Checked the comparison against the three cases that matter: same count with a different soname now fails, an unchanged set passes, and a strictly smaller set passes. Co-Authored-By: Claude <noreply@anthropic.com>
…5B one The template's `docker_build_arg` was empty while its image tag advertised a pinned RCCL build. With `local_image: true` the skill's launch path builds a missing image from this manifest, so a first run silently took the Dockerfile defaults — RCCL_BRANCH=develop, no cmake options, the base image's rdma-core — and produced an image whose tag claimed something it was not. The build args that decide the comparison are now in the template: RCCL_COMMIT, RCCL_BRANCH, RCCL_CMAKE_OPTIONS, BUILD_GPU_TARGETS, RDMA_CORE_VERSION. Adds the 405B template, so the 405B row in the results table is reproducible from a shipped card rather than from overrides typed at the prompt. It carries `per_device_batch_size=1`: the shipped config OOMs at four nodes asking for 241.25 GiB in jit_train_step, and the sharding knobs (dcn_fsdp_parallelism, ici_fsdp_parallelism) do not move that number at all, so what does not fit is an activation, not parameter or optimizer state. Both templates report the same validate_manifest.sh profile as the existing ones (3 PASS, 2 WARN, 3 FAIL, all of them unfilled placeholders). Co-Authored-By: Claude <noreply@anthropic.com>
The metric is emitted for any log carrying step timings, single- and multi-node alike, and that is deliberate rather than an oversight: this parser writes the madengine `multiple_results` CSV and does not print a `performance:` line, so the throughput rows every card already reports are untouched and madengine's first-match scraping is not involved. The opt-in gate that existed in the retired MAD-native harness protected that log-line path, for cards which no longer exist here. Co-Authored-By: Claude <noreply@anthropic.com>
There was a problem hiding this comment.
🟡 Changes recommended
The 405B template retains a known hanging autotune default, and empty destination RUNPATHs are not preserved during library replacement.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (1)
Previously missed (1) — in code that hasn't changed since the last review.
.claude/skills/mad-slurm-multinode/references/gotchas.md:227
- This says level 4 was never timed at 32 ranks and cites only the retired harness, but the PR's current Primus validation reports that the default level 4 stalled for hours at both 24 and 32 ranks. As written, the durable guidance omits the observed failure and tells users to repeat it; update this section to distinguish the current Primus result and document level 0 for these RCCL A/B runs.
- Files reviewed: 6/6 changed files
- Comments generated: 2
- Review effort level: Balanced
Both templates shipped without the one setting the results they document depend on. At 24 and 32 ranks Primus's default autotune level 4 does not finish: the run shows zero steps, zero NCCL INFO and a thread pinned in libhsa-runtime64.so.1 for hours, which reads as a hang rather than as compilation. A reference manifest that cannot reach step 1 is worse than no reference manifest, because the reader blames the workload. Set in both env blocks, since only context.docker_env_vars reaches the container. The gotchas entry is updated to match: it still said level 4 had never been timed at 32 ranks, which was true when written and is not any more. It now records that level 4 does not finish at 24 or 32 ranks while level 0 reached RCCL init in about a minute and trained 50 steps, that this holds on the stock image as well as an overlay and on fp8 as well as bf16, and that sixteen ranks survive it — so the failure scales with rank count and says nothing about the library under test. Co-Authored-By: Claude <noreply@anthropic.com>
The swap restored the destination's RUNPATH only when it had one. Where it had none, the candidate's own $ORIGIN/../lib survived the copy, so that location silently began resolving its dependencies somewhere else than it did before. The soname gate cannot catch this: it compares which names are unresolved, not which file each name resolved to, so every name can still resolve and resolve differently. Remove the RUNPATH explicitly in the empty case, and fail the build if that cannot be done — an unreported failure here would leave exactly the state the gate is blind to. Co-Authored-By: Claude <noreply@anthropic.com>
There was a problem hiding this comment.
🔵 Needs a closer look
The broad library replacement and optional rdma-core surgery require final human validation despite no specific defect found.
Review details
- Files reviewed: 6/6 changed files
- Comments generated: 0 new
- Review effort level: Balanced
Enables the JAX/MaxText workload to run against any RCCL commit rather than only the
librccl the base image ships:
Comparing two builds is two images and the same model card twice — no A/B driver, no variant
switch, nothing added to the model path.
Rewritten onto the current
mad-rccl. The previous revision targeted the pre-#238 MaxTextharness (the per-model
env_scripts/,models.json, and the threejax-maxtext_benchmark_*scripts), all of which #238 removed. Rebasing was not possible — every file it touched had been
deleted underneath it — so the same capability is re-done on the Primus path: the overlay image,
the median
seconds_per_stepmetric, the skill notes and manifest templates, plus the fixesfound while validating and in review.
Worked example
One pinned RCCL commit,
-DFAULT_INJECTION=ONvs=OFF, onlylibrccl.sodiffering.MI355X/gfx950, Broadcom Thor2 RoCEv2, base
rocm/jax-training:maxtext-v26.6.QPS=1QPS=8QPS=8QPS=8The cost scales with inter-node collective volume. The single-node row is the control — the flag
also compiles a delay into the collective device kernels, so a slowdown under it is not
automatically a network result. Every run is 50 steps (step-time medians only compare at equal
steps), each A/B pair ran back-to-back on the same nodes, and every run confirmedNET/IB/*/GDRDMAwith zero socket fallback. Single repeats.These percentages are not comparable with the v26.5 numbers this PR previously carried: the
Primus configs have a ~3.5× larger baseline step time, so the same absolute overhead is a much
smaller fraction of it (v26.5 at 4 nodes: 1.064 − 0.888 = 0.176 s; here: 3.2205 − 3.0835 = 0.137 s).
Two settings the table depends on — both are in the shipped templates
XLA_GPU_AUTOTUNE_LEVEL=0. Primus defaults--xla_gpu_autotune_level=4, and at 24 and 32ranks that presents as a dead hang: zero steps, zero
NCCL INFO, one thread pinned at 100 % inlibhsa-runtime64.so.1for hours. Sixteen ranks survive it. It is compilation, not a defect —but it is indistinguishable from a hang without timing it. Note that Primus reads
XLA_GPU_AUTOTUNE_LEVEL; the MAD-nativeXLA_AUTOTUNE_LEVELis a silent no-op on this path, socarrying the old key across does nothing. Both templates set it in
context.docker_env_varsanddeployment_config.env_vars, since only the former reaches the container.per_device_batch_size=1. The shipped config OOMs at241.25 GiBforjit_train_step. Settingdcn_fsdp_parallelism=-1/ici_fsdp_parallelism=8does not change that figure by a byte (the overrides verifiably reach MaxText), so the 241.25 GiB
is an activation, not parameter or optimizer state.
primus_maxtext_llama3.1-405b.template.jsoncarries it, so the row is reproducible from a shipped card rather than from overrides typed at
the prompt.
Traps worth knowing beyond this PR
NCCL_IB_DISABLE=1itself. It is injected into the generatedsbatch, absent from the manifest, and nothing fails — RCCL reports
via NET/Socketand the jobjust runs slower. Worse, an A/B taken over sockets is insensitive. Set it to 0 in both env
blocks; the templates do.
--ulimit nofile=1048576is required at 4+ nodes. A 32-device clique exhausts the default andstalls silently at
Initialize cliquerather than reporting an error.in the ROCm pip wheels resolves
libamd_smi.so.26through a RUNPATH reaching a sibling wheel,while the candidate is built for
/opt/rccl. Dropping the candidate in verbatim leaves a librcclthat cannot resolve its own dependencies, and
import torchdies before step 1. Where thedestination carried no RUNPATH the candidate's
$ORIGIN/../libhas to be removed rather thanleft behind, or that location quietly begins resolving its dependencies somewhere else.
passes a swap that drops one missing soname and introduces a different one — the number is
unchanged and the image is broken in exactly the way the gate exists to catch. The overlay records
the set per target before the swap and fails on any name that was not already there. It stays a
relative check rather than "must resolve everything", because on this base the pristine
_rocm_sdk_develcopy already reports six unresolved deps: it is only everdlopen'd into anamespace where those libraries are loaded. Note the limit of any such check — it sees which
names are unresolved, never which file a name resolved to, which is why the RUNPATH is
restored rather than merely inspected.
The durable findings are in
.claude/skills/mad-slurm-multinode/references/gotchas.md.Testing
MI355X/gfx950, Broadcom Thor2 RoCEv2. Both images built from scratch and verified in-image: pinned
SHA,
.rodatamarker 2 vs 0 with a positive control,bnxt_reprovider present. 1-, 2- and 4-noderuns go through
madengine runend to end. Every run manifest passes the skill'svalidate_manifest.shwith zero FAILs; the two templates report the same profile as the existingones in
assets/manifests/(their FAILs are the unfilled<FILL_...>placeholders, by design).The 1-node control could not be run through madengine's single-node SLURM path on this cluster and
needed a node-local workspace: that path opens the manifest by basename while its local-workspace
branch copies it in as
build_manifest.json, and withshared_workspaceset it then fails tocreate
run_directoryas container-root against NFSroot_squash. That is a madengine issue, notone this PR touches.