Skip to content

JAX/MaxText multi-node RCCL overlay - #228

Merged
i-kosarev merged 10 commits into
ROCm:mad-rcclfrom
i-kosarev:ilkosare/jax-maxtext-multinode-rccl-fi
Sep 8, 2026
Merged

i-kosarev merged 10 commits into
ROCm:mad-rcclfrom
i-kosarev:ilkosare/jax-maxtext-multinode-rccl-fi

Conversation

@i-kosarev

@i-kosarev i-kosarev commented Aug 25, 2026 •

Copy link
Copy Markdown
Contributor

Enables the JAX/MaxText workload to run against any RCCL commit rather than only the
librccl the base image ships:

--build-arg RCCL_COMMIT=<sha>          # any commit of ROCm/rocm-systems
--build-arg RCCL_CMAKE_OPTIONS=<flags> # and any build options

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 MaxText
harness (the per-model env_scripts/, models.json, and the three jax-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_step metric, the skill notes and manifest templates, plus the fixes
found while validating and in review.

Worked example

One pinned RCCL commit, -DFAULT_INJECTION=ON vs =OFF, only librccl.so differing.
MI355X/gfx950, Broadcom Thor2 RoCEv2, base rocm/jax-training:maxtext-v26.6.

Model Nodes Inter-node traffic per step FI=ON FI=OFF Delta
8B FP8 1 none 2.9240 s 2.9200 s +0.14 %
8B FP8 2 grad all-reduce, QPS=1 3.0335 s 3.0320 s +0.05 %
8B FP8 2 grad all-reduce, QPS=8 3.0590 s 3.0300 s +0.96 %
8B FP8 4 grad all-reduce, QPS=8 3.2205 s 3.0835 s +4.44 %
405B FP8 4 FSDP across nodes, QPS=8 61.02 s 17.40 s +250.8 %

The 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 confirmed
NET/IB/*/GDRDMA with 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 32
    ranks that presents as a dead hang: zero steps, zero NCCL INFO, one thread pinned at 100 % in
    libhsa-runtime64.so.1 for 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-native XLA_AUTOTUNE_LEVEL is a silent no-op on this path, so
    carrying the old key across does nothing. Both templates set it in context.docker_env_vars and
    deployment_config.env_vars, since only the former reaches the container.
  • 405B at 4 nodes needs per_device_batch_size=1. The shipped config OOMs at
    241.25 GiB for jit_train_step. Setting dcn_fsdp_parallelism=-1 / ici_fsdp_parallelism=8
    does 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.json
    carries it, so the row is reproducible from a shipped card rather than from overrides typed at
    the prompt.

Traps worth knowing beyond this PR

  • The standard SLURM path exports NCCL_IB_DISABLE=1 itself. It is injected into the generated
    sbatch, absent from the manifest, and nothing fails — RCCL reports via NET/Socket and the job
    just runs slower. Worse, an A/B taken over sockets is insensitive. Set it to 0 in both env
    blocks; the templates do.
  • --ulimit nofile=1048576 is required at 4+ nodes. A 32-device clique exhausts the default and
    stalls silently at Initialize clique rather than reporting an error.
  • A librccl swap must preserve each destination's own RUNPATH — including an empty one. The copy
    in the ROCm pip wheels resolves libamd_smi.so.26 through a RUNPATH reaching a sibling wheel,
    while the candidate is built for /opt/rccl. Dropping the candidate in verbatim leaves a librccl
    that cannot resolve its own dependencies, and import torch dies before step 1. Where the
    destination carried no RUNPATH the candidate's $ORIGIN/../lib has to be removed rather than
    left behind, or that location quietly begins resolving its dependencies somewhere else.
  • Gate that swap on the set of unresolved sonames, not on how many there are. Comparing counts
    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_devel copy already reports six unresolved deps: it is only ever dlopen'd into a
    namespace 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, .rodata marker 2 vs 0 with a positive control, bnxt_re provider present. 1-, 2- and 4-node
runs go through madengine run end to end. Every run manifest passes the skill's
validate_manifest.sh with zero FAILs; the two templates report the same profile as the existing
ones 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 with shared_workspace set it then fails to
create run_directory as container-root against NFS root_squash. That is a madengine issue, not
one this PR touches.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_multi self-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_step metric 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.

Comment thread scripts/jax-maxtext/run.sh Resolved
Copilot AI review requested due to automatic review settings August 25, 2026 10:51

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment thread scripts/jax-maxtext/run.sh Resolved
Copilot AI review requested due to automatic review settings August 25, 2026 10:56

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

  • || true at the end of the docker run | tee pipeline forces the overall command to succeed and (in practice) resets PIPESTATUS, so rc=${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"

Copilot AI review requested due to automatic review settings August 25, 2026 13:58

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.sh failures (e.g., HF auth / tokenizer download) are not checked here, and run.sh does not run with set -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"

Copilot AI review requested due to automatic review settings August 25, 2026 15:01

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copilot AI review requested due to automatic review settings August 26, 2026 16:03
@i-kosarev
i-kosarev force-pushed the ilkosare/jax-maxtext-multinode-rccl-fi branch from 154eb05 to b871803 Compare August 26, 2026 16:03

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_HOST defaults 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

Copilot AI review requested due to automatic review settings August 26, 2026 17:04
@i-kosarev
i-kosarev force-pushed the ilkosare/jax-maxtext-multinode-rccl-fi branch from b871803 to 736b44c Compare August 26, 2026 17:04

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_DIR defaults 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"

Copilot AI review requested due to automatic review settings August 26, 2026 17:43
@i-kosarev
i-kosarev force-pushed the ilkosare/jax-maxtext-multinode-rccl-fi branch from 736b44c to 8819eab Compare August 26, 2026 17:43

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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-node can under-request GPUs or leave them unavailable to the job. Default this to the normal directive and let Thor2 users opt out, as documented in references/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,

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 11 out of 11 changed files in this pull request and generated 2 comments.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 11 out of 11 changed files in this pull request and generated 1 comment.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 11 out of 11 changed files in this pull request and generated 1 comment.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 11 out of 11 changed files in this pull request and generated no new comments.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_arg omits RCCL_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 by example_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>"

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_TARGETS below explicitly allows gfx942, but this cache key always says gfx950. 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. Add RCCL_CMAKE_OPTIONS here 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>"

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 11 out of 11 changed files in this pull request and generated 4 comments.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 11 out of 11 changed files in this pull request and generated 2 comments.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_PORT is overwritten with 29500 whenever MASTER_PORT is 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}"

Comment thread scripts/jax-maxtext/run.sh Resolved
Comment thread scripts/jax-maxtext/models.json Resolved

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 11 out of 11 changed files in this pull request and generated 3 comments.

Comment thread scripts/jax-maxtext/run.sh Resolved

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 11 out of 11 changed files in this pull request and generated 1 comment.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 11 out of 11 changed files in this pull request and generated 1 comment.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 11 out of 11 changed files in this pull request and generated 1 comment.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 11 out of 11 changed files in this pull request and generated no new comments.

i-kosarev and others added 5 commits September 3, 2026 10:34
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>

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 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

Comment thread scripts/jax-maxtext/run.sh Resolved
i-kosarev and others added 3 commits September 6, 2026 23:30
…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>

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 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

i-kosarev and others added 2 commits September 6, 2026 23:43
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>

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔵 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

@mkuznet1 mkuznet1 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants