From c6d16f3138ff964662d48dadde8285f203a600b7 Mon Sep 17 00:00:00 2001 From: WangLingxun Date: Tue, 21 Jul 2026 05:56:55 +0000 Subject: [PATCH 1/7] ci/docs: move primus base image to v26.4 - ci.yaml: BASE_IMAGE + runs-on v26.3 -> v26.4; drop the v26.3-only 'Purge stale TorchTitan' step (v26.4 ships no torchtitan to shadow). Keep 'Install fixed origami' (v26.4 still ships no origami). - Dockerfile: default BASE_IMAGE + comments v26.3 -> v26.4. - docs/scripts: bump default image refs to v26.4, release branch to release/v26.4, and pip pin to primus==26.4.0. --- .github/workflows/ci.yaml | 45 ++++++--------------------- .github/workflows/docker/Dockerfile | 13 ++++---- README.md | 10 +++--- benchmark/kernel/rccl/run_slurm.sh | 2 +- benchmark/kernel/rccl/submit_pairs.sh | 2 +- docs/01-getting-started/quickstart.md | 2 +- examples/README.md | 6 ++-- examples/run_k8s_pretrain.sh | 4 +-- examples/run_local_pretrain.sh | 4 +-- examples/run_local_pretrain_cli.sh | 2 +- examples/run_slurm_pretrain_cli.sh | 2 +- runner/README.md | 2 +- runner/primus-cli-container.sh | 2 +- tools/docker/start_container.sh | 2 +- 14 files changed, 36 insertions(+), 62 deletions(-) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 002706592..ffe08f556 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -28,7 +28,7 @@ env: ROCSHMEM_COMMIT: 17ff985c026f9f97f85068647e863ab541dd5645 # Update version to 3.2.0 for 7.2.0 rocm release (#351) (#355) UCCL_COMMIT: 5afb4117893c58cc0c8557d9286336141a301053 # [EP]: fix fp8 error of internode_ll on amd gfx950 arch. (#710) TRITON_COMMIT: 88b227e23f0445f3f695bad05bbf1a363b4f50e0 - BASE_IMAGE: docker.io/rocm/primus:v26.3 + BASE_IMAGE: docker.io/rocm/primus:v26.4 MAXTEXT_BASE_IMAGE: docker.io/rocm/jax-training:maxtext-v26.2 jobs: @@ -133,7 +133,7 @@ jobs: docker push docker.io/tasimage/primus:${{env.IMAGE_TAG}} # docker login -u rocmshared -p ${{ secrets.ROCM_DOCKER_HUB_TOKEN }} - # Primus v26.3 already includes AINIC under /workspace. Re-enable this + # Primus v26.4 already includes AINIC under /workspace. Re-enable this # Dockerfile.ainic build only when we need to refresh the tasimage -ainic image. echo "> Build Docker Image with tag: ${{ env.IMAGE_TAG }}-ainic" start_time=$(date +%s) @@ -211,7 +211,7 @@ jobs: needs: [code-lint] if: ${{ github.event_name != 'pull_request' || !github.event.pull_request.draft }} # runs-on: [primus-lm-cicd-torch-j8knc] - runs-on: [primus-lm-cicd-v26.3-tas8n-a16-40] + runs-on: [primus-lm-cicd-v26.4-tas8n-a16-40] steps: - run: echo "🎉 Begin Primus-Turbo Checkout." - name: Clean stale Primus-Turbo checkout @@ -285,41 +285,14 @@ jobs: - name: Install Primus run: | pip install -r requirements.txt - # v26.3-only: drop this whole step on v26.4. Unit tests import torchtitan - # from the submodule source via conftest sys.path injection (v0.2.2 has - # PEP-420 namespace packages that need the source root on sys.path). This - # just removes any stale torchtitan on the persistent runner that would - # shadow it (old install or leftover editable finder / .pth / dist-info). - - name: Purge stale TorchTitan install (v26.3-only, drop on v26.4) - run: | - pip uninstall -y torchtitan || true - python - <<'PY' - import glob, os, shutil, site - dirs = set(site.getsitepackages() + [site.getusersitepackages()]) - patterns = ( - "__editable__*torchtitan*", - "__editable__.torchtitan*.pth", - "torchtitan*.pth", - "torchtitan*.egg-link", - "torchtitan*.dist-info", - "torchtitan*.egg-info", - ) - for d in dirs: - for pat in patterns: - for f in glob.glob(os.path.join(d, pat)): - try: - shutil.rmtree(f) if os.path.isdir(f) else os.remove(f) - print("removed", f) - except OSError as e: - print("skip", f, e) - PY - name: Install fixed origami run: | - # rocm/primus:v26.3 bundles origami 0.1.0, whose rank_configs() raises - # ValueError: vector::reserve during MoE grouped-gemm kernel selection and - # crashes training (turbo's _safe_rank_configs only catches RuntimeError). - # Install the fixed origami (rocm-libraries@223648a) over the bundled 0.1.0. - # TODO: drop this once a base image ships origami with the fix. + # The v26.4 base image ships no origami at all, so primus_turbo falls + # back to a heuristic kernel selection instead of origami's tuned + # rank_configs(). Install origami (rocm-libraries@223648a) so CI exercises + # the same grouped-gemm path as the published image (Dockerfile installs + # the same pin). Keep in sync with .github/workflows/docker/Dockerfile. + # TODO: drop once a base image ships origami with the fix. rm -rf /tmp/rocm-libraries git clone --filter=blob:none --no-checkout https://github.com/ROCm/rocm-libraries.git /tmp/rocm-libraries cd /tmp/rocm-libraries diff --git a/.github/workflows/docker/Dockerfile b/.github/workflows/docker/Dockerfile index d25e6c77b..464d9135c 100644 --- a/.github/workflows/docker/Dockerfile +++ b/.github/workflows/docker/Dockerfile @@ -1,4 +1,4 @@ -ARG BASE_IMAGE=docker.io/rocm/primus:v26.3 +ARG BASE_IMAGE=docker.io/rocm/primus:v26.4 FROM ${BASE_IMAGE} ARG PRIMUS_TURBO_COMMIT @@ -24,7 +24,7 @@ RUN rm -rf /var/lib/apt/lists/* ENV ROCSHMEM_HOME=/opt/rocshmem ENV UCX_HOME=/opt/ucx # ENV MPI_HOME=/opt/ompi -# Use the system OpenMPI prefix from the v26.3 base image. +# Use the system OpenMPI prefix from the v26.4 base image. ENV MPI_HOME=/usr/lib/x86_64-linux-gnu/openmpi ENV ROCM_HOME=/opt/rocm ENV PRIMUS_TURBO_FRAMEWORK=${PRIMUS_TURBO_FRAMEWORK} @@ -101,10 +101,11 @@ RUN if [ "$PRIMUS_TURBO_FRAMEWORK" != "JAX" ]; then \ fi # --------------------------------------------------------------------------- -# Install fixed origami (rocm-libraries@223648a) over the base image's bundled -# 0.1.0. The bundled origami's rank_configs() raises -# `ValueError: vector::reserve` during MoE grouped-gemm kernel selection and -# crashes training (turbo's _safe_rank_configs only catches RuntimeError). +# Install origami (rocm-libraries@223648a). The v26.4 base image ships no origami +# at all, so primus_turbo would fall back to heuristic kernel selection instead +# of origami's tuned rank_configs() for MoE grouped-gemm. (The earlier v26.3 base +# bundled a broken origami 0.1.0 whose rank_configs() raised +# `ValueError: vector::reserve`; this pin also supersedes that.) # Skipped for JAX. TODO: drop once a base image ships origami with the fix. RUN if [ "$PRIMUS_TURBO_FRAMEWORK" != "JAX" ]; then \ rm -rf /tmp/rocm-libraries && \ diff --git a/README.md b/README.md index 12b8810f9..504c7b09b 100644 --- a/README.md +++ b/README.md @@ -91,7 +91,7 @@ primus-cli deps sync --dir ~/.cache/Primus/third_party ```bash # For Megatron-LM and TorchTitan backends - docker pull rocm/primus:v26.3 + docker pull rocm/primus:v26.4 # For MaxText backend docker pull rocm/jax-training:maxtext-v26.4-jax0.9.1-te2.12.0 ``` @@ -102,7 +102,7 @@ primus-cli deps sync --dir ~/.cache/Primus/third_party git clone --recurse-submodules https://github.com/AMD-AGI/Primus.git cd Primus # checkout the branch for the specific release - git checkout release/v26.3 + git checkout release/v26.4 git submodule update --init --recursive ``` @@ -113,7 +113,7 @@ primus-cli deps sync --dir ~/.cache/Primus/third_party # NOTE: If your config downloads weights/tokenizer from Hugging Face Hub, # you typically need to pass HF_TOKEN into the container. # Run in the Primus repository root directory - ./primus-cli container --image rocm/primus:v26.3 \ + ./primus-cli container --image rocm/primus:v26.4 \ --env HF_TOKEN="hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx" \ -- train pretrain --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml ``` @@ -131,7 +131,7 @@ For more detailed usage instructions, see the [CLI User Guide](./docs/02-user-gu python -m venv primus-env source primus-env/bin/activate # Install Primus - pip install "primus==26.3.1" --no-deps --extra-index-url https://amd-agi.github.io/Primus/simple/ + pip install "primus==26.4.0" --no-deps --extra-index-url https://amd-agi.github.io/Primus/simple/ ``` @@ -142,7 +142,7 @@ For more detailed usage instructions, see the [CLI User Guide](./docs/02-user-gu 2. **Run training in container using pip-installed Primus** ```bash - primus-cli container --image rocm/primus:v26.3 \ + primus-cli container --image rocm/primus:v26.4 \ --env HF_TOKEN="hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx" \ --volume /path/to/your/data:/data -- --log_file /data/run.log \ -- train pretrain --config /data/your/config.yaml diff --git a/benchmark/kernel/rccl/run_slurm.sh b/benchmark/kernel/rccl/run_slurm.sh index 2a6710147..26ba72aa1 100644 --- a/benchmark/kernel/rccl/run_slurm.sh +++ b/benchmark/kernel/rccl/run_slurm.sh @@ -10,7 +10,7 @@ # Usage: # DOCKER_IMAGE= sbatch run_slurm.sh # DOCKER_IMAGE= NNODES=2 PARTITION=my-gpu sbatch run_slurm.sh -# DOCKER_IMAGE=rocm/primus:v26.3 NNODES=2 sbatch -N2 -w smci355-ccs-aus-n04-[25,29] -p Compute-DCPT ./run_slurm.sh +# DOCKER_IMAGE=rocm/primus:v26.4 NNODES=2 sbatch -N2 -w smci355-ccs-aus-n04-[25,29] -p Compute-DCPT ./run_slurm.sh # # Environment variables (all optional except DOCKER_IMAGE): # DOCKER_IMAGE Docker image to use (required) diff --git a/benchmark/kernel/rccl/submit_pairs.sh b/benchmark/kernel/rccl/submit_pairs.sh index ba64c93a1..3f5a83a12 100755 --- a/benchmark/kernel/rccl/submit_pairs.sh +++ b/benchmark/kernel/rccl/submit_pairs.sh @@ -1,7 +1,7 @@ #!/usr/bin/env bash set -euo pipefail -DOCKER_IMAGE="rocm/primus:v26.3" +DOCKER_IMAGE="rocm/primus:v26.4" #DOCKER_IMAGE="docker.gpuperf:5000/aai_2026_training/rocm/primus_megatron:v25.11_gpt_oss_sink" #DOCKER_IMAGE="docker.gpuperf:5000/gpuperf/primus:v26.1_sinkfa" NNODES=2 diff --git a/docs/01-getting-started/quickstart.md b/docs/01-getting-started/quickstart.md index b9075430c..bb9a4789a 100644 --- a/docs/01-getting-started/quickstart.md +++ b/docs/01-getting-started/quickstart.md @@ -65,7 +65,7 @@ Install Primus in a virtual environment: ```bash python -m venv primus-env source primus-env/bin/activate -pip install "primus==26.3.1" --no-deps --extra-index-url https://amd-agi.github.io/Primus/simple/ +pip install "primus==26.4.0" --no-deps --extra-index-url https://amd-agi.github.io/Primus/simple/ ``` > **Note:** This installs only the Primus CLI into your virtual environment (under `site-packages`), without other dependencies. Third-party submodules are downloaded on the first run of the container, and the complete training software stack is provided in the AMD-published Docker images. You can launch `primus-cli` from any directory. diff --git a/examples/README.md b/examples/README.md index 1a2aec904..eedff4c60 100644 --- a/examples/README.md +++ b/examples/README.md @@ -49,7 +49,7 @@ We recommend using the official [rocm/megatron-lm Docker image](https://hub.dock ```bash # Pull the latest Docker image -docker pull docker.io/rocm/primus:v26.3 +docker pull docker.io/rocm/primus:v26.4 ``` @@ -126,7 +126,7 @@ Multi-node training is launched via **SLURM**. Specify the number of nodes and the model config: ```bash -export DOCKER_IMAGE="docker.io/rocm/primus:v26.3" +export DOCKER_IMAGE="docker.io/rocm/primus:v26.4" export NNODES=8 # Example for megatron llama3.1_8B @@ -292,7 +292,7 @@ When using the `create` command to start a new training workload, the following | `--gpu` | Number of GPUs | 8 | | `--exp` | Path to experiment (training config) file (required) | — | | `--data_path` | Path to training data | — | -| `--image` | Docker image to use | `docker.io/rocm/primus:v26.3` | +| `--image` | Docker image to use | `docker.io/rocm/primus:v26.4` | | `--hf_token` | HuggingFace token | Read from env var `HF_TOKEN` | | `--workspace` | Workspace name | `primus-safe-pretrain` | | `--nodelist` | Comma-separated list of node hostnames to run on | — | diff --git a/examples/run_k8s_pretrain.sh b/examples/run_k8s_pretrain.sh index a57eda70d..ae07ab2bd 100644 --- a/examples/run_k8s_pretrain.sh +++ b/examples/run_k8s_pretrain.sh @@ -15,7 +15,7 @@ GPU="8" EXP_PATH="" DATA_PATH="" BACKEND="megatron" -IMAGE="docker.io/rocm/primus:v26.3" +IMAGE="docker.io/rocm/primus:v26.4" HF_TOKEN="${HF_TOKEN:-}" WORKSPACE="primus-safe-pretrain" NODELIST="" @@ -38,7 +38,7 @@ Options for create: --backend Training backend, e.g. megatron | torchtitan(default: megatron) --exp Path to EXP config (optional) --data_path Data path (optional) - --image Docker image to use (default: docker.io/rocm/primus:v26.3) + --image Docker image to use (default: docker.io/rocm/primus:v26.4) --hf_token HuggingFace token (default: from env HF_TOKEN) --workspace Workspace name (default: safe-cluster-dev) --nodelist Comma-separated list of node names to run on (optional) diff --git a/examples/run_local_pretrain.sh b/examples/run_local_pretrain.sh index f40292308..f0db4a280 100755 --- a/examples/run_local_pretrain.sh +++ b/examples/run_local_pretrain.sh @@ -16,7 +16,7 @@ Usage: bash run_local_pretrain.sh This script launches a Primus pretraining task inside a Docker/Podman container. Environment Variables: - DOCKER_IMAGE Docker image to use [Default: docker.io/rocm/primus:v26.3] + DOCKER_IMAGE Docker image to use [Default: docker.io/rocm/primus:v26.4] MASTER_ADDR Master node IP or hostname [Default: localhost] MASTER_PORT Master node port [Default: 1234] NNODES Total number of nodes [Default: 1] @@ -44,7 +44,7 @@ export EXP=${EXP:-"examples/megatron/exp_pretrain.yaml"} if [ "${BACKEND:-}" = "MaxText" ]; then DOCKER_IMAGE=${DOCKER_IMAGE:-"docker.io/rocm/jax-training:maxtext-v26.2"} else - DOCKER_IMAGE=${DOCKER_IMAGE:-"docker.io/rocm/primus:v26.3"} + DOCKER_IMAGE=${DOCKER_IMAGE:-"docker.io/rocm/primus:v26.4"} fi # Project root diff --git a/examples/run_local_pretrain_cli.sh b/examples/run_local_pretrain_cli.sh index 7e6305c9b..df930b077 100755 --- a/examples/run_local_pretrain_cli.sh +++ b/examples/run_local_pretrain_cli.sh @@ -15,7 +15,7 @@ EXP=${EXP:-"examples/megatron/exp_pretrain.yaml"} if [ "${BACKEND:-}" = "MaxText" ]; then DOCKER_IMAGE=${DOCKER_IMAGE:-"docker.io/rocm/jax-training:maxtext-v25.9"} else - DOCKER_IMAGE=${DOCKER_IMAGE:-"docker.io/rocm/primus:v26.3"} + DOCKER_IMAGE=${DOCKER_IMAGE:-"docker.io/rocm/primus:v26.4"} fi # ------------------ Cluster Env Defaults ------------------ diff --git a/examples/run_slurm_pretrain_cli.sh b/examples/run_slurm_pretrain_cli.sh index 6d10d0231..bdfa8850a 100755 --- a/examples/run_slurm_pretrain_cli.sh +++ b/examples/run_slurm_pretrain_cli.sh @@ -35,7 +35,7 @@ mkdir -p "$LOG_DIR" # NOTE: The --env entries below are passed into the container and will be visible # to the Primus training process (and system hooks) inside the container. bash "$PRIMUS_PATH/runner/primus-cli" slurm "${SLURM_ARGS[@]}" \ --- --image "${DOCKER_IMAGE:-rocm/primus:v26.3}" \ +-- --image "${DOCKER_IMAGE:-rocm/primus:v26.4}" \ -- \ --env "USING_AINIC=${USING_AINIC:-0}" \ --env "PATCH_TE_FLASH_ATTN=${PATCH_TE_FLASH_ATTN:-0}" \ diff --git a/runner/README.md b/runner/README.md index e93241c6a..2317d5663 100644 --- a/runner/README.md +++ b/runner/README.md @@ -162,7 +162,7 @@ This is the most common pattern when you need a fixed software stack. bash runner/primus-cli slurm \ -N 4 \ --nodelist "node[01-04]" \ --- --image "rocm/primus:v26.3" \ +-- --image "rocm/primus:v26.4" \ -- --env NCCL_DEBUG=INFO \ -- train pretrain --config examples/megatron/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml ``` diff --git a/runner/primus-cli-container.sh b/runner/primus-cli-container.sh index 97928f15e..8e1883429 100755 --- a/runner/primus-cli-container.sh +++ b/runner/primus-cli-container.sh @@ -45,7 +45,7 @@ Docker/Podman Options: --cap-add Add Linux capabilities (e.g., SYS_PTRACE) Container Configuration: - --image Docker image [default: rocm/primus:v26.3] + --image Docker image [default: rocm/primus:v26.4] --name Container name --user Run as specific user (e.g., 1000:1000) --network Network mode (e.g., host, bridge) diff --git a/tools/docker/start_container.sh b/tools/docker/start_container.sh index d29289122..2fe081557 100755 --- a/tools/docker/start_container.sh +++ b/tools/docker/start_container.sh @@ -6,7 +6,7 @@ ############################################################################### PRIMUS_PATH=$(realpath "$(dirname "$0")/../..") -DOCKER_IMAGE=${DOCKER_IMAGE:-"docker.io/rocm/primus:v26.3"} +DOCKER_IMAGE=${DOCKER_IMAGE:-"docker.io/rocm/primus:v26.4"} DATA_PATH=${DATA_PATH:-"${PRIMUS_PATH}/data"} SANITIZED_USER=$(echo "${USER:-unknown}" | tr -cd '[:alnum:]_-') if [ -z "$SANITIZED_USER" ]; then From 7d93fb7bc0604aa57592e7bdb82811813e892647 Mon Sep 17 00:00:00 2001 From: WangLingxun Date: Tue, 21 Jul 2026 09:24:34 +0000 Subject: [PATCH 2/7] ci/patch: drop origami install, self-detect aiter deepbind need - Dockerfile/ci.yaml: drop the origami install; base image ships none, so primus_turbo falls back to heuristic MoE grouped-gemm kernel selection. - aiter_deepbind_patches.py: skip the RTLD_DEEPBIND hook once transformer_engine >= 2.14 (fixed in rocm/primus:v26.4); still applies on older images. - Add unit tests for the new gating logic. --- .github/workflows/ci.yaml | 21 +--- .github/workflows/docker/Dockerfile | 23 +--- .../patches/turbo/aiter_deepbind_patches.py | 50 ++++++++- .../megatron/test_aiter_deepbind_patches.py | 105 ++++++++++++++++++ 4 files changed, 160 insertions(+), 39 deletions(-) create mode 100644 tests/unit_tests/backends/megatron/test_aiter_deepbind_patches.py diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index ffe08f556..f781e0874 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -285,22 +285,11 @@ jobs: - name: Install Primus run: | pip install -r requirements.txt - - name: Install fixed origami - run: | - # The v26.4 base image ships no origami at all, so primus_turbo falls - # back to a heuristic kernel selection instead of origami's tuned - # rank_configs(). Install origami (rocm-libraries@223648a) so CI exercises - # the same grouped-gemm path as the published image (Dockerfile installs - # the same pin). Keep in sync with .github/workflows/docker/Dockerfile. - # TODO: drop once a base image ships origami with the fix. - rm -rf /tmp/rocm-libraries - git clone --filter=blob:none --no-checkout https://github.com/ROCm/rocm-libraries.git /tmp/rocm-libraries - cd /tmp/rocm-libraries - git sparse-checkout init --cone - git sparse-checkout set shared/origami - git checkout 223648a26928ebed7f3dd0ccdc044c09f1dccf9b - pip uninstall -y origami || true - pip install ./shared/origami/python + # This runner's image ships no origami, so primus_turbo falls back to + # heuristic kernel selection instead of origami's tuned rank_configs() + # for MoE grouped-gemm (benchmarked: no measurable gain on gfx942). If + # needed, install it the same way docker/Dockerfile documents, e.g.: + # pip install "git+https://github.com/ROCm/rocm-libraries.git@223648a26928ebed7f3dd0ccdc044c09f1dccf9b#subdirectory=shared/origami/python" - name: Set UT_LOG_PATH run: | ts="$(date +%Y%m%d-%H%M%S)" diff --git a/.github/workflows/docker/Dockerfile b/.github/workflows/docker/Dockerfile index 464d9135c..d477e531a 100644 --- a/.github/workflows/docker/Dockerfile +++ b/.github/workflows/docker/Dockerfile @@ -100,24 +100,11 @@ RUN if [ "$PRIMUS_TURBO_FRAMEWORK" != "JAX" ]; then \ rm -rf /opt/uccl; \ fi -# --------------------------------------------------------------------------- -# Install origami (rocm-libraries@223648a). The v26.4 base image ships no origami -# at all, so primus_turbo would fall back to heuristic kernel selection instead -# of origami's tuned rank_configs() for MoE grouped-gemm. (The earlier v26.3 base -# bundled a broken origami 0.1.0 whose rank_configs() raised -# `ValueError: vector::reserve`; this pin also supersedes that.) -# Skipped for JAX. TODO: drop once a base image ships origami with the fix. -RUN if [ "$PRIMUS_TURBO_FRAMEWORK" != "JAX" ]; then \ - rm -rf /tmp/rocm-libraries && \ - git clone --filter=blob:none --no-checkout https://github.com/ROCm/rocm-libraries.git /tmp/rocm-libraries && \ - cd /tmp/rocm-libraries && \ - git sparse-checkout init --cone && \ - git sparse-checkout set shared/origami && \ - git checkout 223648a26928ebed7f3dd0ccdc044c09f1dccf9b && \ - (pip uninstall -y origami || true) && \ - pip install --no-cache-dir ./shared/origami/python && \ - rm -rf /tmp/rocm-libraries; \ - fi +# This base image ships no origami, so primus_turbo falls back to heuristic +# kernel selection instead of origami's tuned rank_configs() for MoE +# grouped-gemm (benchmarked: no measurable gain on gfx942). If needed, install +# it the same way CI used to, e.g.: +# pip install "git+https://github.com/ROCm/rocm-libraries.git@223648a26928ebed7f3dd0ccdc044c09f1dccf9b#subdirectory=shared/origami/python" # Set the default working directory WORKDIR /opt diff --git a/primus/backends/megatron/patches/turbo/aiter_deepbind_patches.py b/primus/backends/megatron/patches/turbo/aiter_deepbind_patches.py index cf0852cda..da1abce71 100644 --- a/primus/backends/megatron/patches/turbo/aiter_deepbind_patches.py +++ b/primus/backends/megatron/patches/turbo/aiter_deepbind_patches.py @@ -26,12 +26,17 @@ below, gated on Turbo + affected arch (TE path never touched); ``before_train`` runs well before the first backward that JIT-imports these modules. -TEMPORARY: delete this file (and its turbo/__init__.py entry) once the base image's -TE is built against the updated aiter. +Self-detection: skipped once ``transformer_engine.__version__`` >= 2.14, which renamed +the vendored library and stopped exporting the colliding ``aiter::mha_bwd`` symbol -- +confirmed fixed in rocm/primus:v26.4. Older images (e.g. v26.3) still get the patch. + +TEMPORARY: delete this file (and its turbo/__init__.py entry) once every base image +still in support ships a TE built against the updated aiter. """ import importlib import os +import re import sys from primus.backends.megatron.patches.turbo.utils import is_primus_turbo_can_patch @@ -71,6 +76,34 @@ def _is_affected_arch() -> bool: return False +# First TE version that dropped the stale vendored libmha_bwd.so (see module docstring). +_TE_FIX_VERSION = (2, 14, 0) +_TE_FIX_VERSION_STR = ".".join(map(str, _TE_FIX_VERSION)) + + +def _te_version_tuple(): + """Best-effort (major, minor, patch) from ``transformer_engine.__version__``, or + None if TE isn't importable / the version string can't be parsed.""" + try: + import transformer_engine + + match = re.match(r"(\d+)\.(\d+)\.(\d+)", getattr(transformer_engine, "__version__", "") or "") + return tuple(int(part) for part in match.groups()) if match else None + except Exception: # pragma: no cover - defensive, never block training + return None + + +def _te_already_fixed() -> bool: + """True if the installed TE no longer vendors the stale libmha_bwd.so. + + Fails safe towards False (unfixed) on an unknown/unparseable version: applying + the hook on an already-fixed TE is a harmless no-op, but skipping it on an + unfixed TE reintroduces the hd128 backward crash. + """ + version = _te_version_tuple() + return version is not None and version >= _TE_FIX_VERSION + + def _install_deepbind_import_hook() -> bool: """ Idempotently wrap ``importlib.import_module`` to OR in RTLD_DEEPBIND for the pinned @@ -115,7 +148,7 @@ def _import_module_with_deepbind(name, package=None): def _can_install_aiter_deepbind(ctx: PatchContext) -> bool: - """Install only on affected arches (gfx942/gfx950) with the Turbo attention path active.""" + """Install only on affected arches (gfx942/gfx950) with an unfixed TE and Turbo attention active.""" args = get_args(ctx) if not bool(getattr(args, "use_turbo_attention", False)): return False @@ -127,6 +160,13 @@ def _can_install_aiter_deepbind(ctx: PatchContext) -> bool: "aiter DEEPBIND isolation not needed." ) return False + if _te_already_fixed(): + log_rank_0( + f"{_LOG_PREFIX} transformer_engine >= {_TE_FIX_VERSION_STR} detected " + "(no vendored stale libmha_bwd.so exported as aiter::mha_bwd); aiter DEEPBIND " + "isolation not needed on this image." + ) + return False return True @@ -137,8 +177,8 @@ def _can_install_aiter_deepbind(ctx: PatchContext) -> bool: description=( "On gfx942/gfx950, install the aiter mha RTLD_DEEPBIND import hook so the Turbo " "attention backward binds the pinned aiter::mha_bwd instead of TE's stale vendored " - "libmha. Gated on Turbo attention (TE path untouched). Temporary workaround " - "(ROCm/aiter#1332)." + "libmha. Gated on Turbo attention and an unfixed TE (auto-skips once " + f"transformer_engine >= {_TE_FIX_VERSION_STR}). Temporary workaround (ROCm/aiter#1332)." ), condition=_can_install_aiter_deepbind, ) diff --git a/tests/unit_tests/backends/megatron/test_aiter_deepbind_patches.py b/tests/unit_tests/backends/megatron/test_aiter_deepbind_patches.py new file mode 100644 index 000000000..13f60b42e --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_aiter_deepbind_patches.py @@ -0,0 +1,105 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Unit tests for the aiter RTLD_DEEPBIND gating logic +(``primus/backends/megatron/patches/turbo/aiter_deepbind_patches.py``). + +These verify the patch self-detects whether the installed transformer_engine +already fixed the stale-vendored-libmha issue (ROCm/aiter#1332), without +needing a GPU. The actual RTLD_DEEPBIND import hook installation is exercised +separately in the diffusion GPU tests via ``tests.utils.install_aiter_deepbind_hook``. +""" + +import sys +import types + +import pytest + +import primus.backends.megatron.patches.turbo.aiter_deepbind_patches as patch_mod + + +def _fake_te_module(version: str) -> types.ModuleType: + module = types.ModuleType("transformer_engine") + module.__version__ = version + return module + + +@pytest.mark.parametrize( + "version, expected", + [ + ("2.14.0", (2, 14, 0)), + ("2.14.1", (2, 14, 1)), + ("1.11.0+abcdef0", (1, 11, 0)), + ("2.14.0.dev0+rocm", (2, 14, 0)), + ], +) +def test_te_version_tuple_parses_leading_semver(monkeypatch, version, expected): + monkeypatch.setitem(sys.modules, "transformer_engine", _fake_te_module(version)) + assert patch_mod._te_version_tuple() == expected + + +def test_te_version_tuple_none_when_unimportable(monkeypatch): + monkeypatch.setitem(sys.modules, "transformer_engine", None) # forces ImportError on import + assert patch_mod._te_version_tuple() is None + + +def test_te_version_tuple_none_when_unparseable(monkeypatch): + monkeypatch.setitem(sys.modules, "transformer_engine", _fake_te_module("not-a-version")) + assert patch_mod._te_version_tuple() is None + + +@pytest.mark.parametrize( + "version, expected_fixed", + [ + ("2.13.9", False), # pre-fix: still vendors the stale libmha_bwd.so + ("2.14.0", True), # fixed: renamed to te_libmha_*.so / qola::te:: + ("2.15.0", True), + ("1.11.0+abcdef0", False), # e.g. rocm/primus:v26.3's TE + ], +) +def test_te_already_fixed_thresholds_on_fix_version(monkeypatch, version, expected_fixed): + monkeypatch.setitem(sys.modules, "transformer_engine", _fake_te_module(version)) + assert patch_mod._te_already_fixed() is expected_fixed + + +def test_te_already_fixed_fails_safe_towards_unfixed_when_unknown(monkeypatch): + # An unparseable/missing version must NOT be treated as "fixed": skipping the + # hook on an actually-unfixed TE reintroduces the hd128 backward crash. + monkeypatch.setitem(sys.modules, "transformer_engine", None) + assert patch_mod._te_already_fixed() is False + + +def _condition_ctx(monkeypatch, *, use_turbo_attention=True, turbo_can_patch=True, affected_arch=True): + args = types.SimpleNamespace(use_turbo_attention=use_turbo_attention) + monkeypatch.setattr(patch_mod, "get_args", lambda ctx: args) + monkeypatch.setattr(patch_mod, "is_primus_turbo_can_patch", lambda ctx: turbo_can_patch) + monkeypatch.setattr(patch_mod, "_is_affected_arch", lambda: affected_arch) + monkeypatch.setattr(patch_mod, "log_rank_0", lambda *a, **k: None) + return object() # PatchContext stand-in; unused by the stubbed collaborators above + + +def test_can_install_skips_when_te_already_fixed(monkeypatch): + ctx = _condition_ctx(monkeypatch) + monkeypatch.setitem(sys.modules, "transformer_engine", _fake_te_module("2.14.0")) + assert patch_mod._can_install_aiter_deepbind(ctx) is False + + +def test_can_install_applies_when_te_unfixed_on_affected_arch(monkeypatch): + ctx = _condition_ctx(monkeypatch) + monkeypatch.setitem(sys.modules, "transformer_engine", _fake_te_module("1.11.0+abcdef0")) + assert patch_mod._can_install_aiter_deepbind(ctx) is True + + +def test_can_install_skips_when_arch_unaffected_regardless_of_te(monkeypatch): + ctx = _condition_ctx(monkeypatch, affected_arch=False) + monkeypatch.setitem(sys.modules, "transformer_engine", _fake_te_module("1.11.0+abcdef0")) + assert patch_mod._can_install_aiter_deepbind(ctx) is False + + +def test_can_install_skips_when_turbo_attention_disabled(monkeypatch): + ctx = _condition_ctx(monkeypatch, use_turbo_attention=False) + monkeypatch.setitem(sys.modules, "transformer_engine", _fake_te_module("1.11.0+abcdef0")) + assert patch_mod._can_install_aiter_deepbind(ctx) is False From 11e1885e883dc43a4438ca3602dfa9f20a12199f Mon Sep 17 00:00:00 2001 From: WangLingxun Date: Thu, 23 Jul 2026 10:47:41 +0000 Subject: [PATCH 3/7] ci: fix stale unit-test container reuse on same-tag repush Container reuse compared only the image tag string, not content. On main, IMAGE_TAG is always "latest", so after build-docker repushes under that tag, the check still saw an identical string and reused the stale container instead of pulling the new image. Fix: pull first, then compare docker image inspect .Id (target) against the container's docker inspect .Image (actual) instead of the tag name. --- .github/workflows/ci.yaml | 23 ++++++------ .github/workflows/docker/Dockerfile | 6 +-- .../patches/turbo/aiter_deepbind_patches.py | 37 ++++++------------- 3 files changed, 25 insertions(+), 41 deletions(-) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 98a90135b..ff174b24e 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -304,16 +304,18 @@ jobs: : > "$RUNNER_TEMP/runtime.tsv" # reset the CI runtime log for this job IMG="docker.io/tasimage/primus:${{ env.IMAGE_TAG }}" echo "Target image: ${IMG}" - # Reuse the container only if it is already running on the same image; - # otherwise recreate it from the freshly pushed image. - cur_img="$(docker inspect -f '{{.Config.Image}}' "${UT_CONTAINER}" 2>/dev/null || true)" + # Pull first, then compare image id (not tag name): IMAGE_TAG is + # constant ("latest") on main, so a tag-only check can't detect a + # same-tag repush and would keep reusing a stale container. + docker pull "${IMG}" + target_id="$(docker image inspect -f '{{.Id}}' "${IMG}" 2>/dev/null || true)" + cur_id="$(docker inspect -f '{{.Image}}' "${UT_CONTAINER}" 2>/dev/null || true)" running="$(docker inspect -f '{{.State.Running}}' "${UT_CONTAINER}" 2>/dev/null || true)" - if [ "${running}" = "true" ] && [ "${cur_img}" = "${IMG}" ]; then - echo "Reusing container ${UT_CONTAINER} (already on ${IMG})." + if [ "${running}" = "true" ] && [ -n "${target_id}" ] && [ "${cur_id}" = "${target_id}" ]; then + echo "Reusing container ${UT_CONTAINER} (already on ${IMG}, id=${target_id})." else - echo "Recreating container ${UT_CONTAINER} (running=${running}, image=${cur_img})." + echo "Recreating container ${UT_CONTAINER} (running=${running}, cur_id=${cur_id}, target_id=${target_id})." docker rm -f "${UT_CONTAINER}" >/dev/null 2>&1 || true - docker pull "${IMG}" docker run -d --name "${UT_CONTAINER}" \ --network host --ipc host --privileged \ --group-add video \ @@ -333,11 +335,8 @@ jobs: set -e pip install -r requirements.txt IN - # This runner's image ships no origami, so primus_turbo falls back to - # heuristic kernel selection instead of origami's tuned rank_configs() - # for MoE grouped-gemm (benchmarked: no measurable gain on gfx942). If - # needed, install it the same way docker/Dockerfile documents, e.g.: - # pip install "git+https://github.com/ROCm/rocm-libraries.git@223648a26928ebed7f3dd0ccdc044c09f1dccf9b#subdirectory=shared/origami/python" + # No origami in this image (see docker/Dockerfile); primus_turbo falls + # back to heuristic kernel selection for MoE grouped-gemm. - name: Set UT_LOG_PATH run: | ts="$(date +%Y%m%d-%H%M%S)" diff --git a/.github/workflows/docker/Dockerfile b/.github/workflows/docker/Dockerfile index d477e531a..cd302f0e0 100644 --- a/.github/workflows/docker/Dockerfile +++ b/.github/workflows/docker/Dockerfile @@ -100,10 +100,8 @@ RUN if [ "$PRIMUS_TURBO_FRAMEWORK" != "JAX" ]; then \ rm -rf /opt/uccl; \ fi -# This base image ships no origami, so primus_turbo falls back to heuristic -# kernel selection instead of origami's tuned rank_configs() for MoE -# grouped-gemm (benchmarked: no measurable gain on gfx942). If needed, install -# it the same way CI used to, e.g.: +# No origami in this image; primus_turbo falls back to heuristic kernel +# selection for MoE grouped-gemm. To add it back: # pip install "git+https://github.com/ROCm/rocm-libraries.git@223648a26928ebed7f3dd0ccdc044c09f1dccf9b#subdirectory=shared/origami/python" # Set the default working directory diff --git a/primus/backends/megatron/patches/turbo/aiter_deepbind_patches.py b/primus/backends/megatron/patches/turbo/aiter_deepbind_patches.py index da1abce71..fa3f7a21e 100644 --- a/primus/backends/megatron/patches/turbo/aiter_deepbind_patches.py +++ b/primus/backends/megatron/patches/turbo/aiter_deepbind_patches.py @@ -7,31 +7,18 @@ """ hd128 backward crash fix via scoped RTLD_DEEPBIND (ROCm/aiter#1332). -Affects gfx942 (MI300X/MI325X) and gfx950 (MI350X/MI355X). - -rocm/primus:v26.3's transformer_engine bundles a STALE vendored ``libmha_bwd.so`` -(old aiter f299f579) that, via ``libck_fused_attn.so``'s DT_NEEDED, interposes the -global symbol ``aiter::mha_bwd`` over the freshly pinned aiter (b5e03ed1) -> Turbo's -hd128 backward mis-selects the swa variant and launches with an invalid grid config -(gridDim.x=0 on gfx942 / "invalid configuration argument" on gfx950) -> 8-rank crash. -The ``mha_bwd_args`` ABI differs between the two aiter revisions, so a -*global* fix (RTLD_GLOBAL preload / .so overwrite) repairs Turbo but breaks the TE -path (TE's libck calls the stale libmha by the old ABI) -- isolation must be scoped. - -Fix: wrap ``importlib.import_module`` (aiter loads its mha extensions through it, -honouring ``sys.setdlopenflags()``) and OR in ``RTLD_DEEPBIND`` only for the pinned -aiter mha modules (``module_fmha_v3_bwd`` / ``mha_bwd_bf16_*``), so they bind their -own fresh ``aiter::mha_bwd`` while TE keeps the stale one. DEEPBIND is scoped to these -few extensions, leaving torch/c10 untouched. Installed by the ``before_train`` patch -below, gated on Turbo + affected arch (TE path never touched); ``before_train`` runs -well before the first backward that JIT-imports these modules. - -Self-detection: skipped once ``transformer_engine.__version__`` >= 2.14, which renamed -the vendored library and stopped exporting the colliding ``aiter::mha_bwd`` symbol -- -confirmed fixed in rocm/primus:v26.4. Older images (e.g. v26.3) still get the patch. - -TEMPORARY: delete this file (and its turbo/__init__.py entry) once every base image -still in support ships a TE built against the updated aiter. +On gfx942/gfx950, rocm/primus:v26.3's transformer_engine bundles a stale vendored +``libmha_bwd.so`` that shadows the global symbol ``aiter::mha_bwd`` over the pinned +aiter build, crashing Turbo's hd128 backward. A global fix would break TE (it needs +the stale symbol), so isolation is scoped: wrap ``importlib.import_module`` and OR in +``RTLD_DEEPBIND`` only for the pinned aiter mha modules, installed via the +``before_train`` patch below. + +Self-detection: skipped once ``transformer_engine.__version__`` >= 2.14 (confirmed +fixed in rocm/primus:v26.4). Older images still get the patch. + +TEMPORARY: delete this file (and its turbo/__init__.py entry) once every supported +base image ships a TE built against the updated aiter. """ import importlib From 1a48a5bbd955876e41cb3c63eb517bf6a5f2c445 Mon Sep 17 00:00:00 2001 From: WangLingxun Date: Fri, 24 Jul 2026 07:52:00 +0000 Subject: [PATCH 4/7] ci/docker: drop stale rocSHMEM build + v26.3 env overrides for v26.4 v26.4 ships a prebuilt rocSHMEM plus correct ROCM_HOME/UCX_HOME/MPI_HOME for its pip-installed ROCm SDK layout. Rebuilding rocSHMEM from source and re-pinning those vars to their old v26.3 paths made the rebuild CMake step fail (ROCM_HOME pointed at a now-nonexistent /opt/rocm) and is redundant anyway, so drop both. The restore instructions left in the Dockerfile intentionally read the vars from the base image instead of re-declaring them, so they stay correct across future base image bumps. --- .github/workflows/ci.yaml | 8 ++++-- .github/workflows/docker/Dockerfile | 39 +++++++++++------------------ 2 files changed, 21 insertions(+), 26 deletions(-) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index ff174b24e..fce52adcd 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -25,7 +25,7 @@ permissions: env: PRIMUS_TURBO_COMMIT: edc8d2ccb0be4888e80ee7c6e765fd3956026a32 # feat: add quantized tensor support for blockwise grouped gemm (#422) PRIMUS_TURBO_AITER_COMMIT: 0f3c58e6edb6754940bcf9fd5f09ccb6f389f52e # AITER v0.1.14.post1 (tag commit) — required by Primus-Turbo main aiter_utils.py - ROCSHMEM_COMMIT: 17ff985c026f9f97f85068647e863ab541dd5645 # Update version to 3.2.0 for 7.2.0 rocm release (#351) (#355) + #ROCSHMEM_COMMIT: 17ff985c026f9f97f85068647e863ab541dd5645 # Update version to 3.2.0 for 7.2.0 rocm release (#351) (#355) UCCL_COMMIT: 5afb4117893c58cc0c8557d9286336141a301053 # [EP]: fix fp8 error of internode_ll on amd gfx950 arch. (#710) TRITON_COMMIT: 88b227e23f0445f3f695bad05bbf1a363b4f50e0 BASE_IMAGE: docker.io/rocm/primus:v26.4 @@ -123,16 +123,20 @@ jobs: echo "> Build Docker Image with tag: ${{ env.IMAGE_TAG }}" echo "> Build dependencies" start_time=$(date +%s) + # Optional build-args, off by default (kept as an array so a commented-out + # element here can't break the docker build \-continuation below). + EXTRA_BUILD_ARGS=() + # EXTRA_BUILD_ARGS+=(--build-arg ROCSHMEM_COMMIT=${ROCSHMEM_COMMIT}) docker build -f $CWS/.github/workflows/docker/Dockerfile \ --network=host \ -t tasimage/primus:${{env.IMAGE_TAG}} \ --build-arg BASE_IMAGE=${BASE_IMAGE} \ --build-arg PRIMUS_TURBO_COMMIT=${PRIMUS_TURBO_COMMIT} \ --build-arg PRIMUS_TURBO_AITER_COMMIT=${PRIMUS_TURBO_AITER_COMMIT} \ - --build-arg ROCSHMEM_COMMIT=${ROCSHMEM_COMMIT} \ --build-arg PRIMUS_TURBO_FRAMEWORK=PYTORCH \ --build-arg UCCL_COMMIT=${UCCL_COMMIT} \ --build-arg TRITON_COMMIT=${TRITON_COMMIT} \ + "${EXTRA_BUILD_ARGS[@]}" \ $CWS/.github/workflows/docker end_time=$(date +%s) elapsed=$((end_time - start_time)) diff --git a/.github/workflows/docker/Dockerfile b/.github/workflows/docker/Dockerfile index cd302f0e0..324cc4cef 100644 --- a/.github/workflows/docker/Dockerfile +++ b/.github/workflows/docker/Dockerfile @@ -4,7 +4,7 @@ FROM ${BASE_IMAGE} ARG PRIMUS_TURBO_COMMIT ARG PRIMUS_TURBO_AITER_COMMIT ARG PRIMUS_TURBO_FRAMEWORK -ARG ROCSHMEM_COMMIT +# ARG ROCSHMEM_COMMIT # only needed if the rocSHMEM rebuild below is restored ARG UCCL_COMMIT ARG TRITON_COMMIT # Non-interactive APT @@ -21,34 +21,25 @@ RUN rm -rf /var/lib/apt/lists/* # --------------------------------------------------------------------------- # Enviroment variables # --------------------------------------------------------------------------- -ENV ROCSHMEM_HOME=/opt/rocshmem -ENV UCX_HOME=/opt/ucx -# ENV MPI_HOME=/opt/ompi -# Use the system OpenMPI prefix from the v26.4 base image. -ENV MPI_HOME=/usr/lib/x86_64-linux-gnu/openmpi -ENV ROCM_HOME=/opt/rocm ENV PRIMUS_TURBO_FRAMEWORK=${PRIMUS_TURBO_FRAMEWORK} -# ENV PATH="/opt/ompi/bin:/opt/ompi/sbin:${PATH}" -# ENV LD_LIBRARY_PATH="/opt/ompi/lib:${LD_LIBRARY_PATH}" ENV GPU_ARCHS="gfx942;gfx950" ENV PYTORCH_ROCM_ARCH="gfx942;gfx950" ENV HCC_AMDGPU_TARGET="gfx942,gfx950" -# --------------------------------------------------------------------------- -# Install rocSHMEM -# --------------------------------------------------------------------------- -RUN mkdir -p /opt && cd /opt && \ - git clone https://github.com/ROCm/rocSHMEM.git && \ - cd rocSHMEM && \ - git checkout ${ROCSHMEM_COMMIT} && \ - mkdir build && \ - cd build && \ - MPI_ROOT=${MPI_HOME} UCX_ROOT=${UCX_HOME} INSTALL_PREFIX=${ROCSHMEM_HOME} ../scripts/build_configs/gda \ - -DGDA_IONIC=ON \ - -DGDA_MLX5=ON \ - -DGDA_BNXT=ON \ - -DUSE_IPC=ON -RUN rm -rf /opt/rocSHMEM + +# rocSHMEM: used by Primus ODC via Primus-Turbo's odc_rocshmem_* extensions. +# This base image already ships a prebuilt install (commit 17ff985c, version +# 3.2.0) at $ROCSHMEM_HOME=/opt/rocshmem, so building it again is redundant. +# +# If a future base image drops it, restore by uncommenting below (and the +# `ARG ROCSHMEM_COMMIT` above). Note: cmake reads the ROCm version from +# -DROCM_PATH=, not $ROCM_PATH; omitting it falls back to the nonexistent +# /opt/rocm and fails. +# RUN mkdir -p /opt && cd /opt && git clone https://github.com/ROCm/rocSHMEM.git && \ +# cd rocSHMEM && git checkout ${ROCSHMEM_COMMIT} && mkdir build && cd build && \ +# MPI_ROOT=${MPI_HOME} UCX_ROOT=${UCX_HOME} INSTALL_PREFIX=${ROCSHMEM_HOME} \ +# ../scripts/build_configs/gda -DROCM_PATH=${ROCM_PATH} \ +# -DGDA_IONIC=ON -DGDA_MLX5=ON -DGDA_BNXT=ON -DUSE_IPC=ON # --------------------------------------------------------------------------- # Install Primus-Turbo From bd04e39d06c7b75c5038bd6f009054178903116a Mon Sep 17 00:00:00 2001 From: WangLingxun Date: Fri, 24 Jul 2026 10:50:34 +0000 Subject: [PATCH 5/7] ci/docker: fix Dockerfile.ainic for v26.4's pip-installed ROCm SDK Same root cause as the earlier Dockerfile fix, in a file that hadn't been touched yet: ROCM_PATH/MPI_PATH were hardcoded to v26.3's /opt/rocm and system-OpenMPI paths, which no longer exist under v26.4's pip-installed rocm-sdk-devel. rccl/amd-anp also hardcode /opt/rocm internally, and amdclang++ can't find its sibling clang++ (moved under llvm/bin/), so symlink both in rather than patching those repos. Also add pipefail so `cmd | tee log` steps fail on `cmd`'s exit code instead of tee's, which was silently swallowing the above failures until a much later, unrelated- looking step. Verified locally end to end (rccl + amd-anp build install). --- .github/workflows/docker/Dockerfile.ainic | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/.github/workflows/docker/Dockerfile.ainic b/.github/workflows/docker/Dockerfile.ainic index ca103df2d..5308ae2b9 100644 --- a/.github/workflows/docker/Dockerfile.ainic +++ b/.github/workflows/docker/Dockerfile.ainic @@ -7,12 +7,16 @@ ARG AINIC_BUNDLE_PATH # Non-interactive APT ENV DEBIAN_FRONTEND=noninteractive +# `cmd | tee log` below reports only tee's exit code by default, silently +# swallowing real failures; pipefail makes the RUN fail with `cmd` instead. +SHELL ["/bin/bash", "-o", "pipefail", "-c"] # --------------------------------------------------------------------------- # Install build dependencies # --------------------------------------------------------------------------- +# initramfs-tools: required by the AINIC install.sh below. RUN apt-get update && \ - apt-get install jq dpkg-dev kmod xz-utils \ + apt-get install jq dpkg-dev kmod xz-utils initramfs-tools \ libfmt-dev libboost-all-dev \ libibverbs-dev ibverbs-utils infiniband-diags -y @@ -20,8 +24,18 @@ RUN apt-get update && \ # Enviroment variables # --------------------------------------------------------------------------- ENV WORKDIR=/workspace -ENV ROCM_PATH=/opt/rocm -ENV MPI_PATH=/usr/lib/x86_64-linux-gnu/openmpi +# ROCM_PATH/MPI_HOME already come from the base image; mirror MPI_HOME here +# instead of hardcoding the old v26.3 OpenMPI path. +ENV MPI_PATH=${MPI_HOME} + +# v26.4 moved ROCm to a pip-installed rocm-sdk-devel, so /opt/rocm no longer +# exists, and rocm-sdk-devel/bin/amdclang++ can't find its sibling clang++ +# (now under llvm/bin/). rccl/amd-anp still hardcode both, so symlink them in +# rather than patching those repos. +RUN ln -sfn "${ROCM_PATH}" /opt/rocm && \ + for f in "${ROCM_PATH}"/llvm/bin/clang*; do \ + ln -sf "$f" "${ROCM_PATH}/bin/$(basename "$f")"; \ + done # =============================== Build AINIC Driver =============================== # WARNING: Please ensure the following environment variables are correctly set: From 439030504557ccaf33441444cb6bd96506f1f9c9 Mon Sep 17 00:00:00 2001 From: WangLingxun Date: Fri, 24 Jul 2026 15:00:32 +0000 Subject: [PATCH 6/7] ci/docker: fix torch/triton metadata mismatch breaking pip installs Our custom triton build (pinned commit) leaves torch's dist-info Requires-Dist pointing at the base image's stock triton, so later pip installs (e.g. unit-test CI) see it as unsatisfiable and silently replace torch with a plain PyPI/CUDA build. Patch the metadata to the actual triton version right after installing it, and restore the TRITON_COMMIT build-arg that had been dead since #694. --- .github/workflows/ci.yaml | 2 +- .github/workflows/docker/Dockerfile | 9 ++++++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index fce52adcd..504d85d38 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -27,7 +27,7 @@ env: PRIMUS_TURBO_AITER_COMMIT: 0f3c58e6edb6754940bcf9fd5f09ccb6f389f52e # AITER v0.1.14.post1 (tag commit) — required by Primus-Turbo main aiter_utils.py #ROCSHMEM_COMMIT: 17ff985c026f9f97f85068647e863ab541dd5645 # Update version to 3.2.0 for 7.2.0 rocm release (#351) (#355) UCCL_COMMIT: 5afb4117893c58cc0c8557d9286336141a301053 # [EP]: fix fp8 error of internode_ll on amd gfx950 arch. (#710) - TRITON_COMMIT: 88b227e23f0445f3f695bad05bbf1a363b4f50e0 + TRITON_COMMIT: 09500db9f0fe66fd176d1f080e2017b37e7e995d # pinned triton-lang/triton commit built from source in Dockerfile BASE_IMAGE: docker.io/rocm/primus:v26.4 MAXTEXT_BASE_IMAGE: docker.io/rocm/jax-training:maxtext-v26.2 diff --git a/.github/workflows/docker/Dockerfile b/.github/workflows/docker/Dockerfile index 324cc4cef..66ac6ad27 100644 --- a/.github/workflows/docker/Dockerfile +++ b/.github/workflows/docker/Dockerfile @@ -70,12 +70,19 @@ RUN rm -rf /opt/Primus-Turbo RUN cd /opt && \ git clone https://github.com/triton-lang/triton.git && \ cd triton && \ - git checkout 09500db9 && \ + git checkout ${TRITON_COMMIT} && \ pip3 install -r python/requirements.txt && \ MAX_JOBS=96 pip3 install --no-build-isolation --force-reinstall --no-deps . RUN rm -rf /opt/triton +# torch's dist-info still pins the base image's stock triton build; without this, +# any later "pip install" (e.g. CI unit tests) sees torch's triton requirement as +# unsatisfiable and silently upgrades torch from PyPI, dragging in a CUDA stack. +RUN TORCH_META=$(python3 -c "import importlib.metadata as m; print(m.distribution('torch')._path / 'METADATA')") && \ + TRITON_VER=$(python3 -c "import importlib.metadata as m; print(m.version('triton'))") && \ + sed -i -E "s/^Requires-Dist: triton==.*/Requires-Dist: triton==${TRITON_VER}/" "${TORCH_META}" + # --------------------------------------------------------------------------- # Install UCCL-EP (skip for JAX framework) # --------------------------------------------------------------------------- From de445264cacdfca7cf029e79a4f33182f936dccb Mon Sep 17 00:00:00 2001 From: WangLingxun Date: Sat, 25 Jul 2026 09:45:16 +0000 Subject: [PATCH 7/7] ci/docker: fix primus_turbo import ImportError from missing rocshmem symbol Primus-Turbo links librocshmem.a via -l: (not --whole-archive) under -fgpu-rdc/--hip-link, so rocSHMEM's team.cpp.o - which defines the plain rocshmem::ROCSHMEM_TEAM_WORLD symbol - never gets pulled into libprimus_turbo_kernels.so. Being a shared object, the missing symbol doesn't fail the build; it only surfaces as an ImportError at runtime. Fix by re-exporting that symbol from a small shim .so and wiring it in via patchelf. --- .github/workflows/docker/Dockerfile | 31 +++++++++++++++++++---------- 1 file changed, 21 insertions(+), 10 deletions(-) diff --git a/.github/workflows/docker/Dockerfile b/.github/workflows/docker/Dockerfile index 66ac6ad27..58fa4e1e1 100644 --- a/.github/workflows/docker/Dockerfile +++ b/.github/workflows/docker/Dockerfile @@ -4,7 +4,6 @@ FROM ${BASE_IMAGE} ARG PRIMUS_TURBO_COMMIT ARG PRIMUS_TURBO_AITER_COMMIT ARG PRIMUS_TURBO_FRAMEWORK -# ARG ROCSHMEM_COMMIT # only needed if the rocSHMEM rebuild below is restored ARG UCCL_COMMIT ARG TRITON_COMMIT # Non-interactive APT @@ -28,18 +27,14 @@ ENV PYTORCH_ROCM_ARCH="gfx942;gfx950" ENV HCC_AMDGPU_TARGET="gfx942,gfx950" # rocSHMEM: used by Primus ODC via Primus-Turbo's odc_rocshmem_* extensions. -# This base image already ships a prebuilt install (commit 17ff985c, version -# 3.2.0) at $ROCSHMEM_HOME=/opt/rocshmem, so building it again is redundant. +# This base image already ships a prebuilt install (version 3.2.0) at +# $ROCSHMEM_HOME=/opt/rocshmem, so building it again is redundant (confirmed: +# rebuilding from source does not change the link issue below either). # -# If a future base image drops it, restore by uncommenting below (and the -# `ARG ROCSHMEM_COMMIT` above). Note: cmake reads the ROCm version from +# If a future base image drops it, restore the from-source build (see git +# history prior to the v26.4 bump) — note cmake reads the ROCm version from # -DROCM_PATH=, not $ROCM_PATH; omitting it falls back to the nonexistent # /opt/rocm and fails. -# RUN mkdir -p /opt && cd /opt && git clone https://github.com/ROCm/rocSHMEM.git && \ -# cd rocSHMEM && git checkout ${ROCSHMEM_COMMIT} && mkdir build && cd build && \ -# MPI_ROOT=${MPI_HOME} UCX_ROOT=${UCX_HOME} INSTALL_PREFIX=${ROCSHMEM_HOME} \ -# ../scripts/build_configs/gda -DROCM_PATH=${ROCM_PATH} \ -# -DGDA_IONIC=ON -DGDA_MLX5=ON -DGDA_BNXT=ON -DUSE_IPC=ON # --------------------------------------------------------------------------- # Install Primus-Turbo @@ -64,6 +59,22 @@ RUN cd /opt && \ RUN rm -rf /opt/Primus-Turbo +# Primus-Turbo links librocshmem.a with -l: (not --whole-archive) under +# -fgpu-rdc/--hip-link, so rocSHMEM's team.cpp.o (which defines the plain, +# non-namespaced rocshmem::ROCSHMEM_TEAM_WORLD symbol) never gets pulled into +# libprimus_turbo_kernels.so; since it's a shared object the missing symbol +# only surfaces as an ImportError at dlopen/import time, not at build time. +# Work around it by re-exporting that symbol from a tiny shim .so, placed on +# libprimus_turbo_kernels.so's existing rpath and wired in via patchelf. +RUN PT_LIB=$(python3 -c "import primus_turbo, os; print(os.path.join(os.path.dirname(primus_turbo.__file__), 'lib', 'libprimus_turbo_kernels.so'))") && \ + mkdir -p /tmp/rocshmem_team_shim && cd /tmp/rocshmem_team_shim && \ + ar x ${ROCSHMEM_HOME}/lib/librocshmem.a team.cpp.o && \ + g++ -shared -fPIC -o ${ROCSHMEM_HOME}/lib/librocshmem_team_shim.so team.cpp.o \ + -L${ROCSHMEM_HOME}/lib -l:librocshmem.a -L${MPI_HOME}/lib -l:libmpi.so \ + -Wl,-rpath,${ROCSHMEM_HOME}/lib -Wl,-rpath,${MPI_HOME}/lib && \ + patchelf --add-needed librocshmem_team_shim.so "${PT_LIB}" && \ + cd / && rm -rf /tmp/rocshmem_team_shim + # --------------------------------------------------------------------------- # Install Triton # ---------------------------------------------------------------------------