ci/docker: bump Primus base image to v26.5

.github/workflows/ci.yaml

@@ -45,7 +45,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: 09500db9f0fe66fd176d1f080e2017b37e7e995d # pinned triton-lang/triton commit built from source in Dockerfile
-  BASE_IMAGE: docker.io/rocm/primus:v26.4
+  BASE_IMAGE: docker.io/rocm/primus:v26.5
   MAXTEXT_BASE_IMAGE: docker.io/rocm/jax-training:maxtext-v26.2
 
 jobs:
@@ -187,7 +187,7 @@ jobs:
           docker push docker.io/tasimage/primus:${{env.IMAGE_TAG}}
           # docker login -u rocmshared -p ${{ secrets.ROCM_DOCKER_HUB_TOKEN }}
 
-          # Primus v26.4 already includes AINIC under /workspace. Re-enable this
+          # Primus v26.5 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)

.github/workflows/docker/Dockerfile

@@ -1,10 +1,9 @@
-ARG BASE_IMAGE=docker.io/rocm/primus:v26.4
+ARG BASE_IMAGE=docker.io/rocm/primus:v26.5
 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
@@ -27,19 +26,11 @@ ENV GPU_ARCHS="gfx942;gfx950"
 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.
-#
-# 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
+# v26.5 exposes rocSHMEM through the pip-installed ROCm SDK under $ROCM_PATH
+# instead of exporting a standalone $ROCSHMEM_HOME. Primus-Turbo auto-detects
+# that library and builds its optional odc_rocshmem_* extensions. The v26.4
+# team-symbol shim is no longer needed: the v26.5 static library links and
+# imports cleanly.
 
 # ---------------------------------------------------------------------------
 # Install Primus-Turbo
@@ -64,22 +55,6 @@ 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
-
 # ---------------------------------------------------------------------------
 # Remove the CUDA-only CUTLASS DSL stack
 # ---------------------------------------------------------------------------
@@ -127,7 +102,7 @@ RUN if [ "$PRIMUS_TURBO_FRAMEWORK" != "JAX" ]; then \
     rm -rf /opt/uccl; \
     fi
 
-# Install origami (rocm-libraries@223648a). The v26.4 base image ships no
+# Install origami (rocm-libraries@223648a). The v26.5 base image ships no
 # origami at all. primus_turbo's MoE grouped-gemm kernel selection tolerates
 # its absence (falls back to a heuristic, no measurable perf loss on gfx942),
 # but Primus's performance-projection tool (primus projection performance)

README.md

@@ -111,9 +111,9 @@ primus-cli deps sync --dir ~/.cache/Primus/third_party
 
     ```bash
     # For Megatron-LM and TorchTitan backends
-    docker pull rocm/primus:v26.4
+    docker pull rocm/primus:v26.5
     # For MaxText backend
-    docker pull rocm/jax-training:maxtext-v26.4-jax0.9.1-te2.12.0
+    docker pull rocm/jax-training:maxtext-v26.5
     ```
 
 2. **Clone the repository**
@@ -122,7 +122,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.4
+    git checkout release/v26.5
     git submodule update --init --recursive
     ```
 
@@ -133,7 +133,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.4 \
+    ./primus-cli container --image rocm/primus:v26.5 \
       --env HF_TOKEN="hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx" \
       -- train pretrain --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml
     ```
@@ -151,7 +151,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.4.0" --no-deps --extra-index-url https://amd-agi.github.io/Primus/simple/
+    pip install "primus==26.5.0" --no-deps --extra-index-url https://amd-agi.github.io/Primus/simple/
 
     ```
 
@@ -162,7 +162,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.4 \
+    primus-cli container --image rocm/primus:v26.5 \
     --env HF_TOKEN="hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx" \
     --volume /path/to/your/data:/data  -- --log_file /data/run.log \
     -- train pretrain --config /data/your/config.yaml

benchmark/kernel/rccl/run_slurm.sh

@@ -10,7 +10,7 @@
 # Usage:
 #   DOCKER_IMAGE=<image> sbatch run_slurm.sh
 #   DOCKER_IMAGE=<image> NNODES=2 PARTITION=my-gpu sbatch run_slurm.sh
-#   DOCKER_IMAGE=rocm/primus:v26.4 NNODES=2 sbatch -N2 -w node[01-02] -p my-gpu ./run_slurm.sh
+#   DOCKER_IMAGE=rocm/primus:v26.5 NNODES=2 sbatch -N2 -w node[01-02] -p my-gpu ./run_slurm.sh
 #
 # Environment variables (all optional except DOCKER_IMAGE):
 #   DOCKER_IMAGE        Docker image to use (required)

examples/README.md

@@ -50,7 +50,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.4
+docker pull docker.io/rocm/primus:v26.5
 
 ```
 
@@ -127,7 +127,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.4"
+export DOCKER_IMAGE="docker.io/rocm/primus:v26.5"
 export NNODES=8
 
 # Example for megatron llama3.1_8B
@@ -293,7 +293,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.4` |
+| `--image`      | Docker image to use                                  | `docker.io/rocm/primus:v26.5` |
 | `--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     | —                                        |

examples/mlperf/llama2_70b/run_with_docker.sh

@@ -13,7 +13,7 @@
 # Quick start:
 #   export DATADIR=/data/mlperf_llama2
 #   export LOGDIR=/data/mlperf_llama2/results
-#   export CONT=rocm/primus:v26.4
+#   export CONT=rocm/primus:v26.5
 #   export DGXSYSTEM=MI355X_1x8x1
 #   bash examples/mlperf/llama2_70b/run_with_docker.sh
 #
@@ -30,7 +30,7 @@ PRIMUS_HOST="${PRIMUS_HOST:-$(cd "${SCRIPT_DIR}/../../.." && pwd)}"
 
 # Defaults (override via env or INTERACTIVE prompts)
 : "${DGXSYSTEM:=MI355X_1x8x1}"
-: "${CONT:=rocm/primus:v26.4}"
+: "${CONT:=rocm/primus:v26.5}"
 : "${DATADIR:=${HOME}/data/mlperf_llama2}"
 : "${LOGDIR:=${DATADIR}/results}"
 : "${NEXP:=1}"

examples/models/kimi-k3/run_kimi_k3_8L_official_pretrain_mi355x.sh

@@ -74,7 +74,7 @@
 ######################### Training Docker and Variables #########################
 # fla (flash-linear-attention) is not in the base image; the megatron pretrain
 # hook installs it (runner/helpers/hooks/train/pretrain/.../requirements-megatron.txt).
-export DOCKER_IMAGE=${DOCKER_IMAGE:-"docker.io/rocm/primus:v26.4"}
+export DOCKER_IMAGE=${DOCKER_IMAGE:-"docker.io/rocm/primus:v26.5"}
 export CLEAN_DOCKER_CONTAINER=${CLEAN_DOCKER_CONTAINER:-1}
 export SKIP_TRAIN=${SKIP_TRAIN:-0}
 

examples/models/kimi-k3/run_kimi_k3_curve_pretrain_mi355x.sh

@@ -111,7 +111,7 @@
 ######################### Training Docker and Variables #########################
 # fla (flash-linear-attention) is installed by the megatron pretrain hook, not
 # the base image.
-export DOCKER_IMAGE=${DOCKER_IMAGE:-"docker.io/rocm/primus:v26.4"}
+export DOCKER_IMAGE=${DOCKER_IMAGE:-"docker.io/rocm/primus:v26.5"}
 export CLEAN_DOCKER_CONTAINER=${CLEAN_DOCKER_CONTAINER:-1}
 export SKIP_TRAIN=${SKIP_TRAIN:-0}
 

examples/run_k8s_pretrain.sh

@@ -15,7 +15,7 @@ GPU="8"
 EXP_PATH=""
 DATA_PATH=""
 BACKEND="megatron"
-IMAGE="docker.io/rocm/primus:v26.4"
+IMAGE="docker.io/rocm/primus:v26.5"
 HF_TOKEN="${HF_TOKEN:-}"
 WORKSPACE="primus-safe-pretrain"
 NODELIST=""
@@ -38,7 +38,7 @@ Options for create:
     --backend <name>            Training backend, e.g. megatron | torchtitan(default: megatron)
     --exp <exp_path>            Path to EXP config (optional)
     --data_path <data_path>     Data path (optional)
-    --image <docker_image>      Docker image to use (default: docker.io/rocm/primus:v26.4)
+    --image <docker_image>      Docker image to use (default: docker.io/rocm/primus:v26.5)
     --hf_token <token>          HuggingFace token (default: from env HF_TOKEN)
     --workspace <workspace>     Workspace name (default: safe-cluster-dev)
     --nodelist <node1,node2>    Comma-separated list of node names to run on (optional)

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.4]
+    DOCKER_IMAGE   Docker image to use [Default: docker.io/rocm/primus:v26.5]
     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.4"}
+    DOCKER_IMAGE=${DOCKER_IMAGE:-"docker.io/rocm/primus:v26.5"}
 fi
 
 # Project root

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.4"}
+    DOCKER_IMAGE=${DOCKER_IMAGE:-"docker.io/rocm/primus:v26.5"}
 fi
 
 # ------------------ Cluster Env Defaults ------------------

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.4}" \
+-- --image "${DOCKER_IMAGE:-rocm/primus:v26.5}" \
 -- \
   --env "USING_AINIC=${USING_AINIC:-0}" \
   --env "PATCH_TE_FLASH_ATTN=${PATCH_TE_FLASH_ATTN:-0}" \

tools/docker/start_container.sh

@@ -6,7 +6,7 @@
 ###############################################################################
 
 PRIMUS_PATH=$(realpath "$(dirname "$0")/../..")
-DOCKER_IMAGE=${DOCKER_IMAGE:-"docker.io/rocm/primus:v26.4"}
+DOCKER_IMAGE=${DOCKER_IMAGE:-"docker.io/rocm/primus:v26.5"}
 DATA_PATH=${DATA_PATH:-"${PRIMUS_PATH}/data"}
 SANITIZED_USER=$(echo "${USER:-unknown}" | tr -cd '[:alnum:]_-')
 if [ -z "$SANITIZED_USER" ]; then