Update Primus Docker base image from v26.1 to v26.2 (#642)

Update all references to the Primus base image across documentation,
configuration files, CI/CD workflows, benchmark helpers, and example
scripts to use the latest v26.2 release.

Keep existing JAX/MaxText image references unchanged.

Author: WangLingxun

.github/workflows/ci.yaml

@@ -14,7 +14,7 @@ env:
   PRIMUS_TURBO_AITER_COMMIT: e83f9903c07001a0ec29e85d223f6e6cdbe00859
   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)
-  BASE_IMAGE: docker.io/rocm/primus:v26.1
+  BASE_IMAGE: docker.io/rocm/primus:v26.2
   MAXTEXT_BASE_IMAGE: docker.io/rocm/jax-training:maxtext-v26.1
 
 jobs:
@@ -115,24 +115,26 @@ jobs:
           docker push docker.io/tasimage/primus:${{env.IMAGE_TAG}}
           docker login -u rocmshared -p ${{ secrets.ROCM_DOCKER_HUB_TOKEN }}
 
-          echo "> Build Docker Image with tag: ${{ env.IMAGE_TAG }}-ainic"
-          start_time=$(date +%s)
-          mkdir -p $GITHUB_WORKSPACE/.github/workflows/docker/ainic
-          cp /apps/tas/0_public/primus_docker_ci/ainic/ainic_bundle_1.117.5-a-56.tar.gz $GITHUB_WORKSPACE/.github/workflows/docker/ainic/ || { echo "Error: Failed to copy ainic bundle"; exit 1; }
-          docker build -f $GITHUB_WORKSPACE/.github/workflows/docker/Dockerfile.ainic \
-               --network=host \
-               -t tasimage/primus:${{env.IMAGE_TAG}}-ainic \
-               --build-arg BASE_IMAGE=docker.io/tasimage/primus:${{env.IMAGE_TAG}} \
-               --build-arg AINIC_BUNDLE_PATH=ainic \
-               $GITHUB_WORKSPACE/.github/workflows/docker
-          end_time=$(date +%s)
-          elapsed=$((end_time - start_time))
-          echo "⏱️ [build primus docker-ainic] Total elapsed time: ${elapsed} seconds"
+          # # Primus v26.2 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)
+          # mkdir -p $GITHUB_WORKSPACE/.github/workflows/docker/ainic
+          # cp /apps/tas/0_public/primus_docker_ci/ainic/ainic_bundle_1.117.5-a-56.tar.gz $GITHUB_WORKSPACE/.github/workflows/docker/ainic/ || { echo "Error: Failed to copy ainic bundle"; exit 1; }
+          # docker build -f $GITHUB_WORKSPACE/.github/workflows/docker/Dockerfile.ainic \
+          #      --network=host \
+          #      -t tasimage/primus:${{env.IMAGE_TAG}}-ainic \
+          #      --build-arg BASE_IMAGE=docker.io/tasimage/primus:${{env.IMAGE_TAG}} \
+          #      --build-arg AINIC_BUNDLE_PATH=ainic \
+          #      $GITHUB_WORKSPACE/.github/workflows/docker
+          # end_time=$(date +%s)
+          # elapsed=$((end_time - start_time))
+          # echo "⏱️ [build primus docker-ainic] Total elapsed time: ${elapsed} seconds"
 
-          docker tag tasimage/primus:${{env.IMAGE_TAG}}-ainic docker.io/tasimage/primus:${{env.IMAGE_TAG}}-ainic
-          docker login -u tasimage -p ${{ secrets.PRIMUS_DOCKER_HUB_TOKEN }}
-          docker push docker.io/tasimage/primus:${{env.IMAGE_TAG}}-ainic
-          docker login -u rocmshared -p ${{ secrets.ROCM_DOCKER_HUB_TOKEN }}
+          # docker tag tasimage/primus:${{env.IMAGE_TAG}}-ainic docker.io/tasimage/primus:${{env.IMAGE_TAG}}-ainic
+          # docker login -u tasimage -p ${{ secrets.PRIMUS_DOCKER_HUB_TOKEN }}
+          # docker push docker.io/tasimage/primus:${{env.IMAGE_TAG}}-ainic
+          # docker login -u rocmshared -p ${{ secrets.ROCM_DOCKER_HUB_TOKEN }}
 
           echo "> Build Docker Image with tag: ${{ env.IMAGE_TAG }}-v25.09-ainic"
           start_time=$(date +%s)
@@ -206,7 +208,7 @@ jobs:
       # PRIMUS_WORKDIR: /wekafs/primus-data/primus_safe_ci/torch
     needs: [code-lint]
     # runs-on: [primus-lm-cicd-torch-j8knc]
-    runs-on: [primus-lm-cicd-torch-tas8n-a16-40]
+    runs-on: [primus-lm-cicd-v26.2-tas8n-a16-40]
     steps:
       - run: echo "🎉 Begin Primus-Turbo Checkout."
       - name: Set commit hash to env

.github/workflows/docker/Dockerfile

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=docker.io/rocm/primus:v26.1
+ARG BASE_IMAGE=docker.io/rocm/primus:v26.2
 FROM ${BASE_IMAGE}
 
 ARG PRIMUS_TURBO_COMMIT
@@ -22,12 +22,14 @@ RUN rm -rf /var/lib/apt/lists/*
 # ---------------------------------------------------------------------------
 ENV ROCSHMEM_HOME=/opt/rocshmem
 ENV UCX_HOME=/opt/ucx
-ENV MPI_HOME=/opt/ompi
+# ENV MPI_HOME=/opt/ompi
+# Use the system OpenMPI prefix from the v26.2 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 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"

README.md

@@ -58,7 +58,7 @@ Primus leverages AMD’s ROCm Docker images to provide a consistent, ready-to-ru
 1. **Pull the latest Docker image**
 
     ```bash
-    docker pull docker.io/rocm/primus:v26.1
+    docker pull docker.io/rocm/primus:v26.2
     ```
 
 2. **Clone the repository**
@@ -74,7 +74,7 @@ Primus leverages AMD’s ROCm Docker images to provide a consistent, ready-to-ru
     # Run training in container
     # NOTE: If your config downloads weights/tokenizer from Hugging Face Hub,
     #       you typically need to pass HF_TOKEN into the container.
-    ./primus-cli container --image rocm/primus:v26.1 \
+    ./primus-cli container --image rocm/primus:v26.2 \
       --env HF_TOKEN="hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx" \
       -- train pretrain --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.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.1 NNODES=2 sbatch -N2 -w smci355-ccs-aus-n04-[25,29] -p Compute-DCPT ./run_slurm.sh
+#   DOCKER_IMAGE=rocm/primus:v26.2 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)

benchmark/kernel/rccl/submit_pairs.sh

@@ -1,7 +1,7 @@
 #!/usr/bin/env bash
 set -euo pipefail
 
-DOCKER_IMAGE="rocm/primus:v26.1"
+DOCKER_IMAGE="rocm/primus:v26.2"
 #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

docs/cli/PRIMUS-CLI-GUIDE.md

@@ -97,7 +97,7 @@ Primus CLI supports three execution modes, each suitable for different scenarios
 **Common Options**:
 | Option | Description | Example |
 |--------|-------------|---------|
-| `--image IMAGE` | Specify container image | `--image rocm/primus:v26.1` |
+| `--image IMAGE` | Specify container image | `--image rocm/primus:v26.2` |
 | `--volume PATH[:PATH]` | Mount directory | `--volume /data:/data` |
 | `--cpus N` | Limit CPU cores | `--cpus 16` |
 | `--memory SIZE` | Limit memory size | `--memory 128G` |
@@ -174,7 +174,7 @@ Primus CLI supports three execution modes, each suitable for different scenarios
 ./primus-cli slurm srun -N 4 -- preflight --report-file-name preflight-report-4N
 
 # if you are using AINIC in your cluster, use the appropriate configuration file
-# for preflight test, set docker image to rocm/primus:v26.1 in the configuration file
+# for preflight test, set docker image to rocm/primus:v26.2 in the configuration file
 ./primus-cli --config runner/use_ainic.yaml slurm srun -N 2 -- preflight --report-file-name preflight-report-2N
 ```
 
@@ -215,7 +215,7 @@ slurm:
 
 # Container configuration
 container:
-  image: "rocm/primus:v26.1"
+  image: "rocm/primus:v26.2"
   options:
     cpus: "32"
     memory: "256G"
@@ -645,7 +645,7 @@ Step 4: primus-cli-container.sh (on each node)
   ├─ Load container.* config (image, devices, mounts, etc.)
   ├─ Parse container params: --image rocm/megatron-lm:v25.8_py310
   ├─ Merge config and CLI params
-  │   Config: image=rocm/primus:v26.1
+  │   Config: image=rocm/primus:v26.2
   │   CLI: --image rocm/megatron-lm:v25.8_py310
   │   Result: image=rocm/megatron-lm:v25.8_py310
   ├─ Build container options

docs/cli/README.md

@@ -48,7 +48,7 @@ primus-cli direct -- benchmark gemm -M 4096 -N 4096 -K 4096
 | Mode | Use Case | Command Example |
 |------|----------|-----------------|
 | **Direct** | Local development, quick validation | `primus-cli direct -- train pretrain` |
-| **Container** | Environment isolation, dependency management | `primus-cli container --image rocm/primus:v26.1 -- train pretrain` |
+| **Container** | Environment isolation, dependency management | `primus-cli container --image rocm/primus:v26.2 -- train pretrain` |
 | **Slurm** | Multi-node distributed training | `primus-cli slurm srun -N 8 -- train pretrain` |
 
 ## 📖 Learn More

docs/quickstart.md

@@ -21,7 +21,7 @@ rocm-smi && docker --version
 
 ```bash
 # Pull Docker image
-docker pull docker.io/rocm/primus:v26.1
+docker pull docker.io/rocm/primus:v26.2
 
 # Clone repository
 git clone --recurse-submodules https://github.com/AMD-AIG-AIMA/Primus.git
@@ -32,7 +32,7 @@ cd Primus
 
 ```bash
 # Run a quick benchmark in container
-./primus-cli container --image rocm/primus:v26.1 \
+./primus-cli container --image rocm/primus:v26.2 \
   -- benchmark gemm -M 4096 -N 4096 -K 4096
 ```
 
@@ -50,7 +50,7 @@ Use the Docker image you just pulled:
 
 ```bash
 # Run training in container (recommended for getting started)
-./primus-cli container --image rocm/primus:v26.1 \
+./primus-cli container --image rocm/primus:v26.2 \
   -- train pretrain --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml
 ```
 
@@ -62,7 +62,7 @@ Use the Docker image you just pulled:
   --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml
 
 # Slurm mode (for multi-node cluster)
-./primus-cli slurm srun -N 8 -p gpu -- container --image rocm/primus:v26.1 \
+./primus-cli slurm srun -N 8 -p gpu -- container --image rocm/primus:v26.2 \
   -- train pretrain --config examples/megatron/configs/MI300X/llama2_7B-pretrain.yaml
 ```
 

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.1
+docker pull docker.io/rocm/primus:v26.2
 
 ```
 
@@ -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.1"
+export DOCKER_IMAGE="docker.io/rocm/primus:v26.2"
 export NNODES=8
 
 # Example for megatron llama3.1_8B
@@ -285,7 +285,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.1` |
+| `--image`      | Docker image to use                                  | `docker.io/rocm/primus:v26.2` |
 | `--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/run_k8s_pretrain.sh

@@ -15,7 +15,7 @@ GPU="8"
 EXP_PATH=""
 DATA_PATH=""
 BACKEND="megatron"
-IMAGE="docker.io/rocm/primus:v26.1"
+IMAGE="docker.io/rocm/primus:v26.2"
 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.1)
+    --image <docker_image>      Docker image to use (default: docker.io/rocm/primus:v26.2)
     --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)