Bare-metal installation (JAX / MaxText): build the Primus JAX training stack from source (no Docker)

Bare-metal installation (JAX / MaxText): build the Primus JAX training stack from source (no Docker)#

This guide explains how to build the Primus JAX / MaxText training software stack directly on a host machine, without using the AMD published JAX training Docker image. It is intended for users who, for policy or operational reasons, cannot run containers and need to reproduce the same environment on bare metal.

It is derived from the official JAX training Dockerfile and installs the same components and versions. Wherever possible, everything is installed inside a Python virtual environment and without sudo. The only steps that require root are a small set of OS-level system libraries (installed with apt), and a couple of optional networking packages used for multi-node training.

Looking for the PyTorch/Megatron/TorchTitan stack instead? See bare-metal-installation.md. This document is the JAX MaxText counterpart. It is leaner than the PyTorch stack — no Flash Attention / aiter / Primus-Turbo / FBGEMM / rocSHMEM builds — but v26.5 still compiles TensorFlow (CPU) and RCCL from source (and TransformerEngine from source on hosts with glibc < 2.38), so it is not build-free.

Python 3.12+ required. MaxText requires Python ≥ 3.12 (the reference image is built on Ubuntu 24.04). Unlike the PyTorch recipe, Python 3.10 is not sufficient here.

⚠️ Host OS: Ubuntu 24.04 / glibc ≥ 2.38 strongly recommended. The prebuilt transformer_engine_rocm_jax wheel (and other JAX training wheels) are built against the Dockerfile’s ubuntu:24.04 base — they need glibc ≥ 2.38 and libstdc++ with GLIBCXX_3.4.32 (GCC 13/14). On an older host such as Ubuntu 22.04 (glibc 2.35) the TransformerEngine shared library fails to load with version 'GLIBC_2.38' not found (and this failure is silently swallowed by the launcher — training just exits right after JAX initializes the GPUs). libstdc++ can be side-loaded via LD_LIBRARY_PATH, but glibc cannot. For a manual install on Ubuntu 22.04 you must either (a) run on a 24.04 host, or (b) build TransformerEngine from source against the host toolchain — see Section 3.7. Check your host with ldd --version.

The automated setup.sh does (b) for you: its te stage detects the host glibc and, when it is < 2.38, transparently builds TransformerEngine from source instead of installing the prebuilt wheel — so bash setup.sh works on both Ubuntu 22.04 and 24.04. Only the prebuilt TE wheel needs glibc ≥ 2.38; the ROCm JAX/PJRT wheels load fine on glibc 2.35.

Quick path: automated install scripts#

If you just want the environment built for you, use the helper scripts in tools/installation-jax/. They automate everything in Section 3 (Python venv) of this guide — venv creation, the ROCm release tarball, MaxText and its dependencies, TensorFlow (built from source), JAX + ROCm PJRT/plugin, TransformerEngine, RCCL (built from source), and Primus itself — and provide a single env.sh to activate the environment before each job. Read the rest of this document if you want to understand or customize what they do, or if you need the multi-node networking stack (Section 4), which the scripts do not build.

There are two files:

  • env.sh — defines the install location and exports every environment variable the build and runtime need (ROCm paths, NVTE_* flags, XLA_FLAGS, MAXTEXT_PATH, cache locations, etc.). Source it both during the build and every time you use the environment.

  • setup.sh — runs the install in re-runnable stages. It sources env.sh automatically.

Choose where it installs (important)#

Everything lives under PRIMUS_JAX_BASE (venv, kept checkouts), with transient build sources on SRC_DIR (defaults to local /tmp for fast I/O). PRIMUS_JAX_BASE is required — there is no default, so set it to a directory you can write to that has tens of GB free (env.sh errors out if it is unset):

export PRIMUS_JAX_BASE=/path/to/big/disk/primus-jax-env   # venv + checkouts (persistent)
export SRC_DIR=/tmp/primus-jax-build                      # transient sources (optional override)

The scripts auto-detect your GPU architecture (env.sh reads rocminfo, or falls back to the kernel KFD sysfs gfx_target_version when no ROCm is installed yet), and install the matching device wheels — gfx942 (MI300X/MI325X), gfx950 (MI350X/MI355X), or both. To force a target, export it before running:

export PYTORCH_ROCM_ARCH="gfx942;gfx950"

PYTORCH_ROCM_ARCH is the variable name the ROCm SDK and TransformerEngine read to select gfx targets — it applies to the JAX build too, despite the name.

If your default python3 is older than 3.12 (e.g. Ubuntu 22.04 ships 3.10), you do not need sudo or a PPA. The scripts use uv to provide Python 3.12:

  • If uv is already installed, env.sh auto-detects a uv-managed >= 3.12 interpreter, and setup.sh runs uv python install 3.12 automatically when none is present yet. Just run bash setup.sh.

  • If uv is not installed, install it once (no root) and re-run:

    python3 -m pip install --user uv      # or: curl -LsSf https://astral.sh/uv/install.sh | sh
bash setup.sh                          # setup.sh will fetch Python 3.12 via uv

To force a specific interpreter instead, export it before running:

export PRIMUS_PYTHON=/path/to/python3.12

Build the environment#

cd tools/installation-jax

bash setup.sh            # run all default stages, in order
bash setup.sh --list     # list available stages
bash setup.sh te         # re-run a single stage (e.g. reinstall TransformerEngine)
bash setup.sh venv rocm jax  # run a subset of stages

Stages are idempotent and re-runnable, so if a step fails you can fix the cause and re-run just that stage. On failure the script stops immediately and prints which stage failed.

Default stages (v26.5):

venv → rocm → maxtext → tf_source → jax → te → primus → jaxreqs → rccl → manifest

This mirrors the v26.5 image, which builds TensorFlow (2.21 CPU) and RCCL from source (tf_source, rccl). Those are heavy: the TF bazel build alone is ~30–60 min. Lighter/alternative stages:

  • tf_cpu_fix — pip tensorflow-cpu instead of the tf_source bazel build.

  • te_source — force the from-source TransformerEngine build regardless of glibc. You normally don’t need to pass this: the default te stage auto-falls-back to a from-source build on glibc < 2.38 hosts (e.g. Ubuntu 22.04), where the prebuilt wheel won’t load (see the host-OS note above and Section 3.7). Heavy (~30–60 min).

# Ubuntu 22.04 (glibc < 2.38): `bash setup.sh` already builds TE from source
# automatically. If you also want to skip the heavy tf_source bazel build, swap
# in the lighter tf_cpu_fix:
bash setup.sh venv rocm maxtext tf_cpu_fix jax te primus jaxreqs rccl manifest

Use the environment for a training job#

# Use the SAME PRIMUS_JAX_BASE you built with
export PRIMUS_JAX_BASE=/path/to/big/disk/primus-jax-env
source tools/installation-jax/env.sh    # activates the venv + sets ROCm/NVTE/XLA vars

python -c "import jax; print('devices:', jax.devices())"

# Primus is checked out under $WORKSPACE_DIR
cd "$WORKSPACE_DIR/Primus"
./primus-cli direct -- train pretrain \
  --config examples/maxtext/configs/MI300X/llama2_7B-pretrain.yaml

Pick the config directory that matches your GPU: examples/maxtext/configs/MI300X/ for gfx942 (MI300X/MI325X) and examples/maxtext/configs/MI355X/ for gfx950 (MI350X/MI355X), e.g. --config examples/maxtext/configs/MI355X/llama2_7B-pretrain.yaml.

What the scripts do NOT do#

  • System (apt) packages (Section 2): skipped — they need root. A C++ compiler (g++/make), git, and the build basics must already be present, along with the small extra set MaxText expects (numactl, curl, lsb-release, …).

  • Multi-node networking (Section 4: UCX, OpenMPI, AINIC): not built. Single-node training works without them; follow Section 4 manually if you need distributed-over-RDMA.

  • gcsfuse: only needed to mount Google Cloud Storage buckets for data. Not required for synthetic-data or local-data runs.

The scripts also adapt a few host-specific details beyond the Dockerfile (the Python part of MaxText’s setup.sh is run directly to skip its apt/interactive steps, and the ROCm/MaxText/Primus checkouts are collapsed onto a single MAXTEXT_PATH). See tools/installation-jax/README.md for the full rationale.

0. The key idea: ROCm comes from a tarball, not from a system install#

This build does not require a system-wide ROCm installation. In v26.5 ROCm is delivered as a release tarball (AMD “TheRock” multi-arch dist) that is extracted into a user-writable directory ($ROCM_DIR, default $PRIMUS_JAX_BASE/rocm) — no /opt/rocm, no root:

  • The tarball (repo.amd.com/rocm/tarball-multi-arch/therock-dist-linux-multiarch-7.14.0.tar.gz) provides the full ROCm toolchain (HIP, hipBLASLt, compilers, headers, libraries). ROCM_PATH points at the extraction dir.

  • jax + jaxlib (upstream) plus the ROCm jax_rocm7_pjrt and jax_rocm7_plugin wheels (from repo.amd.com/rocm/whl-multi-arch/) provide GPU-accelerated JAX built against that ROCm.

The pinned release tarball is stable and avoids any pip version-skew that could break hipBLASLt GEMMs. RCCL is then rebuilt from source (see Section 3.11) and dropped into this ROCm tree.

This means almost the entire stack can be installed without root into a venv. The only host-level requirements from the system administrator are:

  • The AMD GPU kernel driver (amdgpu / ROCm KMD) must already be installed and loaded (/dev/kfd and /dev/dri must exist, and the user must be in the video and render groups).

  • A small set of build/runtime system libraries (see Section 2).

1. Required software stack for JAX MaxText training#

The complete environment is composed of the following layers:

Layer

Component

Source

Needs root?

Kernel / hardware

AMD GPU driver (amdgpu KMD), GPU device access

OS / admin

Yes (one-time, by admin)

OS libraries

Build toolchain + runtime libs (g++, git, numactl, RDMA, …)

apt

Yes (one-time)

ROCm user-space

TheRock ROCm dist (HIP, hipBLASLt, compilers, libs)

release tarball ($ROCM_DIR)

No (user dir)

Deep learning framework

JAX (jax, jaxlib) + ROCm jax_rocm7_pjrt / jax_rocm7_plugin

pip (upstream + repo.amd.com)

No (venv)

Accelerated kernels

TransformerEngine (JAX) — prebuilt ROCm wheel (or from source)

pip (staging index) / build

No (venv)

Training framework

MaxText (ROCm fork) + its Python deps

git + pip (uv)

No (venv)

Collectives fix

tensorflow-cpu 2.21 (built from source; no bundled NCCL/LLVM)

build from source (bazel)

No (venv)

Collectives lib

RCCL (rebuilt from source into the ROCm tree)

build from source

No (user dir)

Multi-node comms

UCX, OpenMPI, AMD AINIC (libionic)

build from source / apt

Mostly no (AINIC needs root)

Primus

Primus + third_party/maxtext submodule

git + pip

No (venv)

1.1 Version requirements (pins and host prerequisites)#

These are the exact versions the v26.5 reference Dockerfile pins. The install scripts use the same pins; change one and you may have to change the others.

Component

Pinned version / source

Notes

ROCm (TheRock dist tarball)

therock-dist-linux-multiarch-7.14.0.tar.gz

Multi-arch (gfx942 + gfx950). Extracted into $ROCM_DIR.

JAX / jaxlib

0.10.0

Upstream PyPI.

ROCm PJRT / plugin

jax_rocm7_pjrt / jax_rocm7_plugin 0.10.0+rocm7.14.0

From repo.amd.com/rocm/whl-multi-arch/.

TransformerEngine (JAX)

transformer_engine_rocm_jax 2.15.0.dev0+rocm7.15.0a20260707.72d01a0

Prebuilt wheel needs glibc ≥ 2.38; else build from source (te_source).

TensorFlow (CPU, from source)

ROCm tensorflow-upstream branch upstream-v2.21.0

Built with bazelisk v1.25.0. Needs host clang-18/lld-18.

RCCL (from source)

rocm-systems @ 9e5e4084a4b8e1e86551b0eb054725c62354a926

Installed into $ROCM_PATH/lib. Needs host clang-18/lld-18.

MaxText (ROCm fork)

release/v26.5

2-value initialize()/run() API; Primus main supports it (fix #912).

Primus

main

Includes the MaxText initialize()/run() compatibility shim.

scipy

1.16

amdsmi

7.0.2

pip, after the ROCm tarball.

Build front-end

cmake 3.31.6, ninja 1.11.1.3, wheel 0.45.1, packaging 25.0, setuptools 69.5.1

Plus uv (used by MaxText’s dep install).

TE deps

pybind11 3.0.4, importlib-metadata 8.7.1, pydantic 2.13.4, flax 0.12.2

Host prerequisites (independent of the pins above):

Requirement

Minimum / recommended

Why

Python

≥ 3.12 — 3.12 recommended/pinned

MaxText requires ≥ 3.12; the prebuilt TE/JAX wheels are cp312, so on a 3.13 venv you must build TE from source. uv provides 3.12 with no root.

glibc

≥ 2.38 for the prebuilt TE wheel

Ubuntu 24.04 = glibc 2.38+. On Ubuntu 22.04 (glibc 2.35) the prebuilt TE wheel won’t load — the te stage auto-falls-back to a from-source build. Check with ldd --version. glibc cannot be side-loaded via LD_LIBRARY_PATH.

libstdc++ (GCC)

GLIBCXX_3.4.32 (GCC 13/14) for prebuilt TE

Can be side-loaded via LD_LIBRARY_PATH if glibc itself is new enough.

C/C++ toolchain

g++ with C++17; clang-18/lld-18 for the TF/RCCL source builds

Source builds (tf_source, rccl, te_source) need LLVM 18.

GPU arch

AMD Instinct gfx942 (MI300/MI325) or gfx950 (MI350/MI355)

The pinned ROCm tarball + JAX wheels target these. Other archs (e.g. gfx90a/MI250) need matching wheels not pinned here.

GPU driver (KMD)

amdgpu / ROCm kernel driver loaded

/dev/kfd + /dev/dri present; user in video/render groups.

Validated: the scripts’ default flow (with te_source substituted for the prebuilt te) has been run end-to-end for single-node MaxText pretraining on gfx942 / Ubuntu 22.04 (glibc 2.35) / Python 3.12. The prebuilt-te flow is the path for Ubuntu 24.04 hosts.

For a distributed (multi-node) JAX MaxText job specifically, beyond JAX and ROCm you additionally need:

  • RCCL (AMD’s collective library) — rebuilt from source into the ROCm tree (see Section 3.11). This is what MaxText’s collectives run over.

  • AMD AINIC / RDMA stack (libibverbs, rdma-core, libionic) — for high-performance networking on AMD Pensando NICs.

  • Correct GPU/NIC device permissions and (often) hugepages / ulimit -l unlimited.

  • UCX + OpenMPI (Section 4) — optional for MaxText; carried over from the reference image for MPI-launched / other JAX workloads. MaxText itself does not use them (see the note in Section 4).

Unlike the PyTorch stack, the JAX MaxText image does not build rocSHMEM, and MaxText does not launch via mpirun. JAX forms its process group through the JAX distributed coordinator (JAX_COORDINATOR_IP, which Primus sets from MASTER_ADDR) and runs collectives over RCCL.

2. System packages (require sudo / administrator, one-time)#

These are OS-level libraries needed to build the rest of the stack and to run MaxText / RDMA networking. They must be installed by someone with root, but this is a one-time action; everything afterward is done unprivileged in a venv.

If you genuinely cannot get root at all, these packages must already be present on the host. The remainder of the guide (Sections 3+) then runs entirely without root.

2.1 Build toolchain and core libraries#

sudo apt update
sudo apt install -y \
    gfortran git git-lfs ninja-build g++ pkg-config xxd patchelf \
    automake libtool flex ccache \
    python3-venv python3-dev python3-pip python-is-python3 \
    libegl1-mesa-dev liblzma-dev libdw1 libdrm-dev \
    wget unzip zip

Source builds (v26.5): LLVM 18 toolchain. The tf_source (TensorFlow 2.21 bazel) and rccl source builds want a host clang-18/lld-18. The reference image adds the LLVM apt repo and installs them:

echo 'deb http://apt.llvm.org/jammy/ llvm-toolchain-jammy-18 main' | sudo tee /etc/apt/sources.list.d/llvm.list
wget -O - https://apt.llvm.org/llvm-snapshot.gpg.key | sudo apt-key add -
sudo apt update && sudo apt install -y clang-18 lld-18 llvm-18-dev llvm-18-tools

(Use jammy for Ubuntu 22.04, noble for 24.04.) You can skip this if you use the lighter tf_cpu_fix stage instead of tf_source.

Python 3.12+: MaxText needs Python ≥ 3.12. Ubuntu 24.04 ships 3.12 by default. On Ubuntu 22.04 (which ships 3.10) apt install python3.12 fails because jammy has no such package — do not rely on it. Two options:

  • Recommended, no sudo — uv (this is what the automated scripts use):

    python3 -m pip install --user uv    # or: curl -LsSf https://astral.sh/uv/install.sh | sh
uv python install 3.12              # downloads a standalone CPython 3.12 (no root)
export PRIMUS_PYTHON="$(uv python find '>=3.12')"

    The manual venv step below then uses "$PRIMUS_PYTHON" -m venv ....

  • Alternative — deadsnakes PPA (needs sudo, and the PPA must be reachable from your network):

    sudo add-apt-repository ppa:deadsnakes/ppa
sudo apt update
sudo apt install -y python3.12 python3.12-venv python3.12-dev
export PRIMUS_PYTHON=python3.12

2.2 Extra packages MaxText’s setup expects#

MaxText’s own setup.sh installs these via apt; on bare metal, pre-install them once:

sudo apt install -y \
    numactl lsb-release gnupg curl net-tools iproute2 procps lsof ethtool

Optional — only if you read training data from Google Cloud Storage:

# gcsfuse (mount GCS buckets)
export GCSFUSE_REPO=gcsfuse-$(lsb_release -c -s)
echo "deb https://packages.cloud.google.com/apt $GCSFUSE_REPO main" | \
    sudo tee /etc/apt/sources.list.d/gcsfuse.list
curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -
sudo apt update && sudo apt install -y gcsfuse

2.3 RDMA / networking libraries (needed for multi-node training)#

sudo apt install -y \
    rdma-core libibverbs-dev ibverbs-utils infiniband-diags \
    ethtool kmod dpkg-dev jq xz-utils \
    libevent-dev libhwloc-dev libmunge-dev \
    software-properties-common

2.4 AMD AINIC library (optional, for AMD Pensando NICs)#

This pulls a vendor .deb from the AMD radeon repository. Skip it if you are not using AMD AINIC networking.

# Pin to the version used by the reference image
AINIC_BUNDLE_VERSION="1.117.5-a-77"

sudo add-apt-repository -y \
  "deb https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_BUNDLE_VERSION} noble main"
sudo apt update --allow-insecure-repositories
sudo apt install -y --allow-unauthenticated libionic-dev

3. Build the Python environment (no sudo from here on)#

Everything below runs as a regular user inside a virtual environment.

3.1 Create and activate the virtual environment#

# Pick a stable location, e.g. ~/primus-jax-env. MaxText needs Python >= 3.12.
# On Ubuntu 22.04, get a 3.12 interpreter via uv first (see the Python 3.12+
# note in Section 2.1): export PRIMUS_PYTHON="$(uv python find '>=3.12')"
"${PRIMUS_PYTHON:-python3.12}" -m venv ~/primus-jax-env
source ~/primus-jax-env/bin/activate

# Build/runtime knobs (match the Dockerfile)
export MAX_JOBS=128                              # lower this if you have fewer CPU cores / less RAM
export PYTORCH_ROCM_ARCH="gfx942;gfx950"         # MI300/MI325 = gfx942, MI350/MI355 = gfx950
export ROCM_AMDGPU_TARGETS="gfx942,gfx950"

3.2 Bootstrap build tooling#

pip install --upgrade pip
pip uninstall -y wheel
pip install \
    cmake==3.31.6 \
    ninja==1.11.1.3 \
    wheel==0.45.1 \
    packaging==25.0 \
    setuptools==69.5.1 \
    uv

3.3 Workaround environment variables#

# Avoids HSA_STATUS_ERROR_OUT_OF_RESOURCES on some configurations
export HSA_ENABLE_SCRATCH_ASYNC_RECLAIM=0
export HSA_NO_SCRATCH_RECLAIM=1

# Fix the ROCm profiler hang issue
export ROCPROFILER_QUEUE_INTERPOSITION=0
export DEBUG_HIP_DYNAMIC_QUEUES=0

3.4 Install ROCm from the TheRock release tarball#

This step replaces a system ROCm install. v26.5 uses a release tarball (not pip wheels) extracted into a user-writable dir — no /opt/rocm, no root.

# Extract into a user-writable location (the automated env.sh uses
# $PRIMUS_JAX_BASE/rocm). ROCM_PATH will point here.
export ROCM_DIR=~/primus-jax-env/rocm
mkdir -p "$ROCM_DIR"
wget -O /tmp/therock-dist.tar.gz \
    https://repo.amd.com/rocm/tarball-multi-arch/therock-dist-linux-multiarch-7.14.0.tar.gz
tar -xzf /tmp/therock-dist.tar.gz -C "$ROCM_DIR"

# amdsmi (installed via pip in the reference image)
pip install amdsmi==7.0.2

The tarball is multi-arch (gfx942 + gfx950), so there is no per-arch package to pick, and there is no pip version-skew to worry about — a pinned tarball cannot drift out of sync and break hipBLASLt GEMMs.

3.5 Export ROCm paths#

Point the rest of the build/runtime at the extracted ROCm. These must be set every time you use the environment — Section 5 shows how to make them persistent (the automated env.sh does all of this for you).

export ROCM_PATH=$ROCM_DIR
export ROCM_HOME=$ROCM_PATH
export HIP_PLATFORM=amd
export HIP_PATH=$ROCM_PATH
export HIP_CLANG_PATH=$ROCM_PATH/llvm/bin
export HIP_INCLUDE_PATH=$ROCM_PATH/include
export HIP_LIB_PATH=$ROCM_PATH/lib
export HIP_DEVICE_LIB_PATH=$ROCM_PATH/lib/llvm/amdgcn/bitcode
# The reference Dockerfile puts $ROCM_PATH/lib on PATH too; mirror that.
export PATH="$ROCM_PATH/lib:$ROCM_PATH/bin:$HIP_CLANG_PATH:$PATH"
export LD_LIBRARY_PATH="$ROCM_PATH/lib:$ROCM_PATH/lib/rocm_sysdeps/lib:$ROCM_PATH/lib64:$ROCM_PATH/llvm/lib"
export LIBRARY_PATH="$ROCM_PATH/lib:$ROCM_PATH/lib64"
export CPATH=$HIP_INCLUDE_PATH
export PKG_CONFIG_PATH="$ROCM_PATH/lib/pkgconfig"

Quick check before continuing:

hipcc --version

3.6 Install MaxText, then JAX + the ROCm PJRT/plugin#

Order matters (v26.5). The subsections below are numbered for reference, but the correct install order is: MaxText deps (§3.8) → TensorFlow from source (§3.9) → ROCm JAX/PJRT/plugin (§3.6, right here) → TransformerEngine (§3.7) → Primus (§3.10) → RCCL from source (§3.11). MaxText’s setup.sh pulls in a stock jax/tensorflow, so the ROCm JAX must be installed after MaxText (to override it) and before TE (or a later step overwrites jaxlib). The automated setup.sh enforces this ordering for you.

MaxText install is in Section 3.8 (deps) below; the JAX packages are:

JAX_VERSION=0.10.0
JAX_PJRT_VERSION=0.10.0+rocm7.14.0       # https://repo.amd.com/rocm/whl-multi-arch/jax-rocm7-pjrt/
JAX_PLUGIN_VERSION=0.10.0+rocm7.14.0     # https://repo.amd.com/rocm/whl-multi-arch/jax-rocm7-plugin/

pip install jax==${JAX_VERSION} jaxlib==${JAX_VERSION} scipy==1.16
pip install \
    --index-url https://repo.amd.com/rocm/whl-multi-arch/ \
    --pre jax_rocm7_pjrt==${JAX_PJRT_VERSION} \
    --pre jax_rocm7_plugin==${JAX_PLUGIN_VERSION}

3.7 Install TransformerEngine (JAX) from the prebuilt ROCm wheel#

TransformerEngine is installed as a prebuilt ROCm JAX wheel. Check the staging index for the current pin.

TE_VERSION=2.15.0.dev0+rocm7.15.0a20260707.72d01a0

pip install \
    pybind11==3.0.4 \
    importlib-metadata==8.7.1 \
    pydantic==2.13.4 \
    flax==0.12.2

pip install \
    --index-url https://rocm.frameworks-nightlies.amd.com/whl-staging/device-all/ \
    --pre \
    --no-build-isolation \
    transformer_engine_rocm_jax==${TE_VERSION}

The prebuilt wheel needs glibc ≥ 2.38 (Ubuntu 24.04). Verify it actually loads before continuing:

python -c "import transformer_engine.jax; print('TE JAX OK')"

If you see OSError: ... version 'GLIBC_2.38' not found (typical on Ubuntu 22.04, glibc 2.35), the wheel is incompatible with your host. Build TransformerEngine from source instead so it links against your host’s glibc. This is exactly what the automated te_source stage does (bash setup.sh ... te_source ...); the manual equivalent is:

pip uninstall -y transformer_engine transformer_engine_rocm_jax
export USE_ROCM=1 NVTE_FRAMEWORK=jax NVTE_USE_ROCM=1
export NVTE_ROCM_ARCH="${PYTORCH_ROCM_ARCH}" CMAKE_BUILD_PARALLEL_LEVEL=${MAX_JOBS}
git clone --recursive https://github.com/ROCm/TransformerEngine.git
cd TransformerEngine
git checkout 635d7c085c39a6d9bfe4881c7d3efab7a46d7129   # last known-good ROCm JAX TE source commit
git submodule update --init --recursive
python3 setup.py bdist_wheel && pip install dist/*.whl
cd ..

If you only hit a GLIBCXX_3.4.32 (libstdc++) error but glibc is new enough, you can instead side-load a newer libstdc++.so.6 (e.g. extracted from a newer distro’s libstdc++6 package) via LD_LIBRARY_PATH — no rebuild needed.

3.8 Install MaxText and its dependencies#

Clone the ROCm MaxText fork and install its Python dependencies. The reference image runs MaxText’s src/dependencies/scripts/setup.sh; on bare metal we run the Python portion of that script directly (the apt/gcsfuse steps are the one-time root action from Section 2, and the venv already exists).

MaxText release/v26.5 and the Primus API. MaxText v26.5 uses a 2-value initialize()/run() API (config, recorder). Primus main handles it: MaxTextPretrainTrainer forwards initialize()’s tuple verbatim to run() (fix #912), so v26.5 trains out of the box. Override MAXTEXT_BRANCH only if you deliberately need to pin a different MaxText release.

cd ~/primus-jax-env   # or your $WORKSPACE_DIR
git clone https://github.com/ROCm/maxtext.git
cd maxtext
git checkout release/v26.5   # matches the v26.5 image; Primus main supports its 2-value API

# MaxText installs its deps with uv. The default (tpu) requirements set contains
# the framework-agnostic Python deps WITHOUT any CUDA packages, which is what the
# ROCm image uses.
pip install -U setuptools wheel uv
python -m uv pip install --resolution=lowest \
    -r src/dependencies/requirements/generated_requirements/tpu-requirements.txt
python -m src.dependencies.scripts.install_pre_train_extra_deps
python -m uv pip install --no-deps -e .
cd ..

This pulls in the full tensorflow package (via tensorflow-text); the next step swaps it for the CPU build. It may also nudge jax/jaxlib/scipy within their allowed ranges — the ROCm PJRT/plugin installed in 3.6 remain in place.

3.9 Build TensorFlow (CPU) from source#

v26.5 rebuilds TensorFlow 2.21 (CPU) from ROCm’s fork. The stock PyPI TF wheel bundles an LLVM whose symbols collide with ROCm’s libLLVM in Grain “spawn” workers → SIGSEGV on import tensorflow after import jax. A CPU build has correct symbol visibility and no bundled NCCL (so it also preserves the XLA→RCCL collective fix). This is a heavy bazel build (~30–60 min) and needs a host clang/lld (LLVM 18) plus unzip/zip (Section 2).

# bazelisk (to a user-writable location) auto-picks the bazel version TF pins.
wget -O ~/primus-jax-env/bin/bazel \
    https://github.com/bazelbuild/bazelisk/releases/download/v1.25.0/bazelisk-linux-amd64
chmod +x ~/primus-jax-env/bin/bazel
export PATH="$HOME/primus-jax-env/bin:$PATH"

git clone --depth 1 --branch upstream-v2.21.0 https://github.com/ROCm/tensorflow-upstream.git
cd tensorflow-upstream
bazel --output_user_root=/tmp/primus-jax-build/bazel build //tensorflow/tools/pip_package:wheel \
    --repo_env=WHEEL_NAME=tensorflow_cpu \
    --repo_env=HERMETIC_PYTHON_VERSION=3.12
pip uninstall -y tensorflow tensorflow-cpu tensorflow_cpu
pip install --no-deps bazel-bin/tensorflow/tools/pip_package/wheel_house/tensorflow_cpu-2.21.0-cp312-cp312-linux_x86_64.whl
cd ..

Lighter alternative: if you don’t want the bazel build, pip install --no-deps tensorflow-cpu==$(pip show tensorflow | awk '/^Version:/{print $2}') installs a CPU wheel that avoids the bundled-NCCL clash. It may still hit the LLVM-symbol SIGSEGV on some ROCm 7.14 configs — the from-source build is the robust fix. The automated recipe exposes this as the tf_cpu_fix stage.

3.10 Install Primus#

cd ~/primus-jax-env   # or your $WORKSPACE_DIR
git clone --recurse-submodules https://github.com/AMD-AGI/Primus.git
cd Primus
git checkout main
git submodule update --init third_party/maxtext/

# The JAX path does not install Primus' torch-oriented requirements.txt.
# Remove stale dataclasses backports that conflict on modern Python:
pip uninstall -y dataclasses dataclasses_json

# Primus' JAX runtime deps (also installed by the MaxText pre-train hook at
# launch time):
pip install -r requirements-jax.txt

If you already have a local Primus checkout (e.g. this repository), you can skip the clone and just run the git submodule update, pip uninstall, and pip install -r requirements-jax.txt steps from its root.

Which MaxText does Primus run? At launch, Primus resolves the MaxText backend from the MAXTEXT_PATH environment variable, falling back to its own third_party/maxtext submodule. Set MAXTEXT_PATH to the checkout you installed dependencies into (Section 3.8) so the code and the installed deps match — the automated env.sh does this for you.

3.11 Build RCCL from source (into the ROCm tree)#

v26.5 rebuilds RCCL from rocm-systems and drops the libraries into the ROCm tree so JAX/XLA collectives use it. Requires the ROCm toolchain (hipcc) from Section 3.4.

git clone https://github.com/ROCm/rocm-systems.git
cd rocm-systems
git checkout 9e5e4084a4b8e1e86551b0eb054725c62354a926
cd projects/rccl
./install.sh -l --prefix build/ --amdgpu_targets="${PYTORCH_ROCM_ARCH}"
cp -r build/release/librccl* "$ROCM_PATH/lib/"
cd ../../..

4. Multi-node communication stack (UCX, OpenMPI)#

These are only needed for multi-node distributed training. They build from source and install into user-writable prefixes (no root needed, except the AINIC .deb already handled in Section 2.4).

Does JAX MaxText actually need UCX/OpenMPI? For MaxText itself, no — JAX forms its process group through the JAX distributed coordinator (JAX_COORDINATOR_IP/JAX_COORDINATOR_PORT, which Primus sets from MASTER_ADDR/MASTER_PORT) and runs collectives over RCCL; there is no mpirun launch and no rocSHMEM. UCX/OpenMPI are carried over from the shared reference image (used by MPI-launched / other JAX-based workloads) and are installed here only for parity. You can skip Section 4 entirely for single- and multi-node MaxText pretraining.

Multi-node: make all local GPUs visible to each process. On each node, every rank/process must see all local GPUs, otherwise JAX enumerates only a single device per node. Export CUDA_VISIBLE_DEVICES covering every local GPU (the ROCm PJRT plugin honors the CUDA-named variable) before launching:

export CUDA_VISIBLE_DEVICES=$(seq -s, 0 $((GPUS_PER_NODE - 1)))   # e.g. 0,1,2,3,4,5,6,7

Add it to your activation script (Section 5) or your job launcher. It is intentionally not hard-coded in env.sh, since the right value depends on how ranks are pinned to GPUs on your host/scheduler.

4.1 UCX#

cd ~/primus-jax-env
UCX_VERSION="1.18.0"
wget https://github.com/openucx/ucx/releases/download/v${UCX_VERSION}/ucx-${UCX_VERSION}.tar.gz
mkdir -p ucx-${UCX_VERSION}
tar -zxf ucx-${UCX_VERSION}.tar.gz -C ucx-${UCX_VERSION} --strip-components=1
cd ucx-${UCX_VERSION}
mkdir build && cd build
../configure --prefix=$HOME/primus-jax-env/ucx-${UCX_VERSION}/install --with-rocm=${ROCM_PATH}
make -j 16 && make install
cd ../..

export UCX_INSTALL_DIR=$HOME/primus-jax-env/ucx-${UCX_VERSION}/install

4.2 OpenMPI#

MPI_VERSION="4.1.6"
wget https://download.open-mpi.org/release/open-mpi/v$(echo "${MPI_VERSION}" | cut -d. -f1-2)/openmpi-${MPI_VERSION}.tar.gz
mkdir -p ompi-${MPI_VERSION}
tar -zxf openmpi-${MPI_VERSION}.tar.gz -C ompi-${MPI_VERSION} --strip-components=1
cd ompi-${MPI_VERSION}
mkdir build && cd build
# Install to a user-writable prefix instead of /opt to avoid sudo
../configure --prefix=$HOME/primus-jax-env/openmpi --with-ucx=${UCX_INSTALL_DIR} \
    --disable-oshmem --disable-mpi-fortran
make -j 16 && make install
cd ../..

export PATH="$HOME/primus-jax-env/openmpi/bin:${PATH}"
export LD_LIBRARY_PATH="$HOME/primus-jax-env/openmpi/lib:${LD_LIBRARY_PATH}"

The Dockerfile installs OpenMPI under /workspace. The $HOME/... prefixes above keep it unprivileged.

5. Make the environment reproducible (activation script)#

Many of the variables above (especially the ROCm paths and the NVTE_* / XLA_* runtime flags) must be present in every shell that runs training. Append them to the venv’s activation script so they’re set whenever you source ~/primus-jax-env/bin/activate:

cat >> ~/primus-jax-env/bin/activate <<'EOF'

# ---- Primus JAX host environment ----
export PYTORCH_ROCM_ARCH="gfx942;gfx950"
export ROCM_AMDGPU_TARGETS="gfx942,gfx950"
export HSA_ENABLE_SCRATCH_ASYNC_RECLAIM=0
export HSA_NO_SCRATCH_RECLAIM=1
export ROCPROFILER_QUEUE_INTERPOSITION=0
export DEBUG_HIP_DYNAMIC_QUEUES=0

# v26.5: ROCm lives in the extracted tarball dir (Section 3.4), NOT a pip wheel.
export ROCM_PATH=$HOME/primus-jax-env/rocm
export ROCM_HOME=$ROCM_PATH
export HIP_PLATFORM=amd
export HIP_PATH=$ROCM_PATH
export HIP_CLANG_PATH=$ROCM_PATH/llvm/bin
export HIP_INCLUDE_PATH=$ROCM_PATH/include
export HIP_LIB_PATH=$ROCM_PATH/lib
export HIP_DEVICE_LIB_PATH=$ROCM_PATH/lib/llvm/amdgcn/bitcode
export PATH="$ROCM_PATH/lib:$ROCM_PATH/bin:$HIP_CLANG_PATH:$HOME/primus-jax-env/openmpi/bin:$PATH"
export LD_LIBRARY_PATH="$ROCM_PATH/lib:$ROCM_PATH/lib/rocm_sysdeps/lib:$ROCM_PATH/lib64:$ROCM_PATH/llvm/lib:$HOME/primus-jax-env/openmpi/lib"
export LIBRARY_PATH="$HIP_LIB_PATH:$ROCM_PATH/lib64"
export CPATH=$HIP_INCLUDE_PATH
export PKG_CONFIG_PATH="$ROCM_PATH/lib/pkgconfig"

# Point Primus at the MaxText checkout whose deps we installed
export MAXTEXT_PATH=$HOME/primus-jax-env/maxtext

# TransformerEngine (ROCm) runtime flags for JAX
export NVTE_ROCM_ARCH="$PYTORCH_ROCM_ARCH"
export NVTE_USE_ROCM=1
export NVTE_USE_HIPBLASLT=1
export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1
export NVTE_FUSED_ATTN=1
export NVTE_CK_USES_BWD_V3=1
export NVTE_CK_USES_FWD_V3=1
export NVTE_CK_IS_V3_ATOMIC_FP32=1
export NVTE_CK_HOW_V3_BF16_CVT=2

# AMD GPU runtime knobs
export GPU_MAX_HW_QUEUES=2
export HIP_FORCE_DEV_KERNARG=1
export HSA_FORCE_FINE_GRAIN_PCIE=1
export NCCL_DEBUG=VERSION
# Bare-metal only: force RCCL to use its built-in ROCm IB/RoCE transport. On a
# host, /usr/local/lib/librccl-net.so is on the default loader path and is
# ABI-incompatible with the from-source RCCL (undefined symbol
# ncclNetPlugin_v11/_v10 -> falls back to v9 -> segfault at clique init). The
# v26.5 image ships no librccl-net.so, so this matches it.
export NCCL_NET_PLUGIN=none

# XLA / JAX runtime settings (v26.5 uses .9)
export XLA_PYTHON_CLIENT_MEM_FRACTION=.9
export XLA_FLAGS="--xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=0 --xla_gpu_enable_all_gather_combine_by_dim=FALSE --xla_gpu_enable_command_buffer=''"
# ---- end Primus JAX host environment ----
EOF

The automated tools/installation-jax/env.sh sets all of the above (and auto-detects the GPU arch); prefer sourcing it over hand-editing activate.

6. Verify the installation#

source ~/primus-jax-env/bin/activate   # or: source tools/installation-jax/env.sh

# GPUs visible to ROCm?
rocm-smi || ls -l /dev/kfd /dev/dri

# JAX sees the GPUs?
python -c "import jax; print('jax', jax.__version__); \
print('backend:', jax.default_backend()); \
print('devices:', jax.devices())"

# Key libraries import cleanly? (import transformer_engine.jax — this actually
# loads TE's shared lib, which is what fails on glibc < 2.38 with the prebuilt wheel)
python -c "import jax, jaxlib, flax, transformer_engine.jax; print('JAX/flax/TE OK')"

# Run a Primus MaxText training directly (no container). Use the config dir that
# matches your GPU: MI300X/ for gfx942 (MI300X/MI325X), MI355X/ for gfx950 (MI350X/MI355X).
cd ~/primus-jax-env/Primus   # or your Primus checkout
./primus-cli direct -- train pretrain \
  --config examples/maxtext/configs/MI300X/llama2_7B-pretrain.yaml
  # gfx950: --config examples/maxtext/configs/MI355X/llama2_7B-pretrain.yaml

jax.default_backend() should report gpu (ROCm), and jax.devices() should list your AMD GPUs. Use primus-cli direct (not container) since you are running on bare metal with everything installed in your environment.

7. Other important considerations#

  • Python version: MaxText requires Python ≥ 3.12. If your venv is older, the build will fail; get a 3.12 interpreter with uv (no sudo — see Section 2.1) and recreate it. The automated setup.sh does this for you (it provisions Python 3.12 via uv and recreates a too-old venv automatically).

  • GPU device access without root: the user running training must be able to read/write /dev/kfd and /dev/dri/* — usually via membership in the video and render groups (sudo usermod -aG video,render $USER, then re-login).

  • Hugging Face access: for gated models/tokenizers, export your token (export HF_TOKEN=hf_xxx and/or huggingface-cli login).

  • Install order is load-bearing (v26.5): MaxText → TensorFlow (from source) → ROCm JAX/PJRT/plugin → TransformerEngine → RCCL (from source). MaxText’s setup.sh pulls in a stock jax/tensorflow, so the ROCm JAX must be installed after MaxText (to override it) and before TE (or jaxlib gets clobbered). The automated setup.sh stage order enforces this.

  • RDMA / multi-node limits: high-performance networking typically requires ulimit -l unlimited and possibly hugepages, configured in /etc/security/limits.conf (admin help). Verify NICs with ibv_devinfo / ibstat. JAX uses the distributed coordinator (JAX_COORDINATOR_IP / JAX_COORDINATOR_PORT), which Primus sets from MASTER_ADDR / MASTER_PORT.

  • Version drift: the ROCm release tarball, the JAX/PJRT/plugin versions, the TransformerEngine wheel, the TensorFlow/RCCL source revisions, and the MaxText branch are all pinned to one release (see the table in Section 1.1). If you change one, you may need to update the others. The Docker image is the authoritative, tested combination — match its Dockerfile ARGs when in doubt.

  • Automated scripts: the manual steps in Section 3 are automated by tools/installation-jax/ (see Quick path above). The multi-node networking stack in Section 4 is still manual.

8. Quick reference: minimal vs. full install#

If you only need single-node MaxText pretraining, you can skip the multi-node components:

Component

Needed for

ROCm (tarball), JAX + PJRT/plugin, TransformerEngine (JAX)

Core MaxText training (install these)

MaxText + its deps, TensorFlow (from source), RCCL (from source)

Core MaxText training (install these)

Primus + third_party/maxtext

Running MaxText via Primus

gcsfuse

Reading data from GCS buckets only

UCX, OpenMPI, AINIC

Multi-node distributed training

Install the core rows first, validate with Section 6, then add the optional components as your workload requires.

Contents