diff --git a/docs/01-getting-started/bare-metal-installation-jax.md b/docs/01-getting-started/bare-metal-installation-jax.md new file mode 100644 index 000000000..379f97f0f --- /dev/null +++ b/docs/01-getting-started/bare-metal-installation-jax.md @@ -0,0 +1,862 @@ +# 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](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](#37-install-transformerengine-jax-from-the-prebuilt-rocm-wheel). +> 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/](https://github.com/AMD-AGI/Primus/tree/main/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): + +```bash +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: + +```bash +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`](https://docs.astral.sh/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: + + ```bash + 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: + +```bash +export PRIMUS_PYTHON=/path/to/python3.12 +``` + +### Build the environment + +```bash +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). + +```bash +# 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 + +```bash +# 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 +``` + +### 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](https://github.com/AMD-AGI/Primus/blob/main/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). +- **UCX + OpenMPI** — point-to-point transport and the launcher. +- **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`. + +> Unlike the PyTorch stack, the JAX MaxText image does **not** build rocSHMEM. +> JAX distributes over RCCL + the JAX distributed coordinator (`JAX_COORDINATOR_IP`). + +--- + +## 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 + +```bash +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: +> +> ```bash +> 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): +> ```bash +> 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): +> ```bash +> 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: + +```bash +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: + +```bash +# 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) + +```bash +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. + +```bash +# 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 + +```bash +# 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 + +```bash +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 + +```bash +# 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. + +```bash +# 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). + +```bash +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: + +```bash +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: + +```bash +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](https://rocm.frameworks-nightlies.amd.com/whl-staging/device-all/transformer-engine-rocm-jax/) +for the current pin. + +```bash +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: +> +> ```bash +> 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: +> +> ```bash +> 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. + +```bash +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). + +```bash +# 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 + +```bash +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. + +```bash +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). + +### 4.1 UCX + +```bash +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 + +```bash +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`: + +```bash +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 + +# 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 + +```bash +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) +cd ~/primus-jax-env/Primus # or your Primus checkout +./primus-cli direct -- train pretrain \ + --config examples/maxtext/configs/MI300X/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/](https://github.com/AMD-AGI/Primus/tree/main/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. diff --git a/docs/01-getting-started/bare-metal-installation.md b/docs/01-getting-started/bare-metal-installation.md index 0a8d5cc97..f9f1af45d 100644 --- a/docs/01-getting-started/bare-metal-installation.md +++ b/docs/01-getting-started/bare-metal-installation.md @@ -6,6 +6,8 @@ It is derived from the official training `Dockerfile` and installs the same comp > **Important**: This is a long, build-heavy process. A full from-source build (Flash Attention, TransformerEngine, aiter, Primus-Turbo, FBGEMM, rocSHMEM, etc.) can take **several hours** and needs a machine with many CPU cores, plenty of RAM, and tens of GB of free disk. The Docker image remains the recommended and best-supported path. Use this guide only when containers are not an option. +> **Using the JAX MaxText backend instead?** This guide builds the PyTorch stack (Megatron-LM / TorchTitan). For the JAX / MaxText backend, follow the leaner [JAX bare-metal installation guide](bare-metal-installation-jax.md) instead. + --- ## Quick path: automated install scripts diff --git a/docs/01-getting-started/installation.md b/docs/01-getting-started/installation.md index 0e54187bc..ca054197b 100644 --- a/docs/01-getting-started/installation.md +++ b/docs/01-getting-started/installation.md @@ -117,7 +117,8 @@ Reproducing the training environment on the host means building the **same stack Follow the full, step-by-step guide here, which includes the exact pinned versions, environment variables, and automated install scripts: -- **[Bare-metal installation: build the Primus training stack from source (no Docker)](./bare-metal-installation.md)** +- **[Bare-metal installation (PyTorch: Megatron-LM / TorchTitan): build the Primus training stack from source (no Docker)](./bare-metal-installation.md)** +- **[Bare-metal installation (JAX / MaxText): build the Primus JAX training stack from source (no Docker)](./bare-metal-installation-jax.md)** ### Verify diff --git a/tools/installation-jax/README.md b/tools/installation-jax/README.md new file mode 100644 index 000000000..e9a60e967 --- /dev/null +++ b/tools/installation-jax/README.md @@ -0,0 +1,163 @@ +# Primus JAX / MaxText environment in a venv (no docker, no sudo) — v26.5 + +Reproduces the Primus **v26.5 JAX training Dockerfile** in a Python virtual +environment. Same package pins as the Dockerfile, adapted for a bare-metal host +with no root and no containers. + +This is the JAX/MaxText counterpart to `tools/installation/` (the PyTorch / +Megatron / TorchTitan stack). Use this one if you want to run the **JAX MaxText** +training backend of Primus. + +## Why it differs from the Dockerfile + +| Constraint | Dockerfile | Here | +|---|---|---| +| ROCm | release **tarball** → `/opt/rocm` | same tarball → **`$ROCM_DIR`** (`$PRIMUS_JAX_BASE/rocm`, user-writable, **no sudo**) | +| GPU arch | gfx942 + gfx950 | **auto-detected** from the host (`rocminfo`/KFD sysfs) for source builds/TE | +| Build dir | container FS | venv on **`$PRIMUS_JAX_BASE`** (persistent); transient sources on local `/tmp` | +| System deps | `apt install ...` | **skipped** (no sudo); documented in the guide's Section 2 | +| MaxText setup | `setup.sh` (runs `apt`, prompts for a venv) | Python steps only (apt is a one-time root action; venv already made) | + +The key reason no sudo is required: the ROCm tarball is extracted into a +user-writable dir (`$ROCM_DIR`) instead of `/opt/rocm`, so we don't depend on +system ROCm or apt for the ROCm toolchain itself. + +> **Python 3.12+ required (3.12 preferred).** MaxText requires Python ≥ 3.12 (the +> PyTorch recipe works on 3.10; this one does not). 3.12 is preferred because the +> prebuilt ROCm wheels are `cp312`. You do **not** need `sudo` or a PPA: +> - The scripts use `python3.12` (preferred) or `python3.13` if on PATH. +> - Otherwise, if [`uv`](https://docs.astral.sh/uv/) is installed, `env.sh` +> auto-detects a uv-managed `>= 3.12` interpreter and `setup.sh` runs +> `uv python install 3.12` for you when none exists yet. +> - No `uv`? Install it once (no root) and re-run: +> `python3 -m pip install --user uv` (or `curl -LsSf https://astral.sh/uv/install.sh | sh`). +> - To force a specific interpreter: `export PRIMUS_PYTHON=/path/to/python3.12`. +> +> On Ubuntu 22.04, `apt install python3.12` fails (jammy has no such package) — +> use the `uv` path above instead. + +> **Host OS: Ubuntu 24.04 / glibc ≥ 2.38 recommended, but 22.04 works.** The +> prebuilt `transformer_engine_rocm_jax` wheel is built against the Dockerfile's +> `ubuntu:24.04` base and needs `glibc ≥ 2.38` + `GLIBCXX_3.4.32` (GCC 13/14). +> On Ubuntu 22.04 (glibc 2.35) the prebuilt TE wheel can't load +> (`version 'GLIBC_2.38' not found`). **`setup.sh` handles this automatically:** +> the `te` stage detects the host glibc and builds TransformerEngine from source +> when it is < 2.38, so `bash setup.sh` works on both 22.04 and 24.04. (Only the +> prebuilt TE wheel needs glibc ≥ 2.38; the ROCm JAX/PJRT wheels load on glibc +> 2.35.) Check with `ldd --version`. + +## Run it + +```bash +cd tools/installation-jax +# PRIMUS_JAX_BASE is REQUIRED (no default). Point it at a directory you can write +# to with tens of GB free — the venv, ROCm tarball, and checkouts all live here: +export PRIMUS_JAX_BASE=/some/big/disk/primus-jax-env +bash setup.sh # all default stages +``` + +If any stage fails the script stops immediately and prints which stage failed; +fix the cause and re-run just that stage. Stages are idempotent: + +```bash +bash setup.sh --list # show stages +bash setup.sh te # reinstall just TransformerEngine +bash setup.sh venv rocm jax # venv + ROCm + JAX only +``` + +## Use the environment afterward + +```bash +# Set the SAME PRIMUS_JAX_BASE you built with (required — env.sh errors without it) +export PRIMUS_JAX_BASE=/some/big/disk/primus-jax-env +source tools/installation-jax/env.sh # activates venv + sets ROCm / NVTE / XLA env vars +python -c "import jax; print(jax.devices())" + +# Primus is checked out at $WORKSPACE_DIR/Primus; MaxText at $MAXTEXT_DIR. +cd "$WORKSPACE_DIR/Primus" +./primus-cli direct -- train pretrain \ + --config examples/maxtext/configs/MI300X/llama2_7B-pretrain.yaml +``` + +`env.sh` exports `MAXTEXT_PATH=$MAXTEXT_DIR`, so Primus runs the same MaxText +checkout we installed the dependencies for. + +## Stages (default order, v26.5) + +`venv` → `rocm` → `maxtext` → `tf_source` → `jax` → `te` → `primus` +→ `jaxreqs` → `rccl` → `manifest` + +- **venv** — create the venv (Python ≥ 3.12) and bootstrap `cmake`/`ninja`/`uv`. +- **rocm** — download + extract the TheRock ROCm release tarball into `$ROCM_DIR` + (+ `amdsmi`). +- **maxtext** — clone ROCm/MaxText (`release/v26.5`) and install its deps (the + Python part of MaxText's `setup.sh`) + the editable MaxText package. +- **tf_source** — build **tensorflow-cpu 2.21 from source** (bazel, ~30–60 min); + fixes the ROCm-vs-TF LLVM symbol clash (SIGSEGV) and drops bundled NCCL. +- **jax** — `jax`/`jaxlib` 0.10.0 + the ROCm `jax_rocm7_pjrt` / `jax_rocm7_plugin` + (installed after MaxText to override its stock jax). +- **te** — prebuilt `transformer_engine_rocm_jax` wheel (+ `flax`, `pydantic`, ...); + auto-falls-back to a from-source build (`te_source`) on glibc < 2.38 hosts. +- **primus** — clone Primus, init the `third_party/maxtext` submodule, drop the + stale `dataclasses` backports. +- **jaxreqs** — install Primus' `requirements-jax.txt` (loguru, wandb, ...). +- **rccl** — build **RCCL from source** and install it into `$ROCM_PATH/lib`. +- **manifest** — dump `pip list` / `env` for reproducibility. + +Optional / alternative stages: + +- **te_source** — force the from-source TransformerEngine build regardless of + glibc. Normally unnecessary: the default `te` stage **auto-detects the host + glibc** and builds from source when it is < 2.38 (e.g. Ubuntu 22.04), where the + prebuilt wheel fails to load with `version 'GLIBC_2.38' not found`. Heavy build + (~30–60 min, compiles CK fused-attention kernels). Only the prebuilt TE wheel + needs glibc ≥ 2.38 — the ROCm JAX/PJRT wheels load fine on glibc 2.35. +- **tf_cpu_fix** — lighter alternative to `tf_source`: `pip install tensorflow-cpu` + instead of the bazel build (avoids the bundled-NCCL clash; may still hit the + LLVM-symbol SIGSEGV on some ROCm 7.14 configs). + +```bash +# Ubuntu 22.04 (glibc < 2.38): `bash setup.sh` already builds TE from source +# automatically. To also skip the heavy tf_source bazel build, swap in tf_cpu_fix: +bash setup.sh venv rocm maxtext tf_cpu_fix jax te primus jaxreqs rccl manifest +``` + +> **MaxText v26.5 and Primus.** MaxText v26.5 uses a **2-value** +> `initialize()`/`run()` API. Primus `main` handles it — `MaxTextPretrainTrainer` +> forwards `initialize()`'s tuple verbatim to `run()` (fix #912) — so this +> recipe's pinned `release/v26.5` trains out of the box. Override +> `MAXTEXT_BRANCH` only if you deliberately need a different MaxText release. + +## What is SKIPPED (needs sudo / apt — not reproducible here) + +- **System packages** (build toolchain, `numactl`, `gcsfuse`, RDMA/verbs libs): + a one-time root action, documented in Section 2 of the guide. +- **AINIC** (`add-apt-repository`, `libionic-dev`): apt-only. Skipped. +- **UCX + OpenMPI**: autotools source builds needed only for multi-node RDMA. + Single-node training works without them; see Section 4 of the guide. +- **gcsfuse**: only needed to mount GCS buckets for data; not required for + synthetic-data or local-data runs. + +## Notes / gotchas vs. the Dockerfile + +- **Stage order matters (v26.5):** `maxtext` → `tf_source` → `jax` → `te`. MaxText's + `setup.sh` pulls in a stock `jax`/`tensorflow`; TF is then rebuilt from source and + the ROCm JAX/plugin is installed after (overriding MaxText's), and TE must come + after JAX or `jaxlib` gets clobbered. The stage order enforces this. +- **TensorFlow + RCCL are built from source** in the default flow (matching the + v26.5 image). These are the two long builds; use `tf_cpu_fix` if you want to skip + the TF bazel build. +- **No FlashAttention/aiter/torch stack.** The JAX MaxText path does not build the + PyTorch kernel libraries; attention fusion comes from `transformer_engine_rocm_jax` + + XLA. +- **Two MaxText checkouts in the Docker image** (`/workspace/maxtext` and + `Primus/third_party/maxtext`) are collapsed here: we install deps from one + checkout and point `MAXTEXT_PATH` at it. +- **Runtime-validated (single-node).** The 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**. + On Ubuntu 24.04 (glibc ≥ 2.38) the prebuilt `te` wheel works directly. The + multi-node networking stack (UCX/OpenMPI/AINIC) is still not exercised here. + +Treat the reference `Dockerfile` as the authoritative, tested version +combination; if you bump one pin you may need to bump the others. diff --git a/tools/installation-jax/env.sh b/tools/installation-jax/env.sh new file mode 100644 index 000000000..b9646ec78 --- /dev/null +++ b/tools/installation-jax/env.sh @@ -0,0 +1,193 @@ +#!/usr/bin/env bash +# env.sh — Primus JAX/MaxText venv environment (v26.5). +# Source this both during the build (setup.sh does it) and every time you +# want to USE the environment: source env.sh +# +# This mirrors the v26.5 JAX training Dockerfile, adapted for a no-docker / +# no-sudo bare-metal install. In v26.5 ROCm is delivered as a RELEASE TARBALL +# (TheRock dist) extracted to a user-writable dir ($ROCM_DIR) — NOT the pip +# `rocm-sdk-devel` wheels used by the v26.4 recipe. No system ROCm is required. +# +# NOTE: PRIMUS_JAX_BASE is REQUIRED (no default) — export it to a big-disk dir +# before sourcing this file. MaxText requires Python >= 3.12; see +# PRIMUS_PYTHON below. + +# ---- Install location (persistent) + transient build sources ---- +# PRIMUS_JAX_BASE is REQUIRED (there is intentionally no default): it is where +# the venv, the extracted ROCm tarball, and the kept checkouts live. Point it at +# a directory you can write to with tens of GB free, then (re-)source this file: +# export PRIMUS_JAX_BASE=/big/disk/primus-jax-env +if [ -z "${PRIMUS_JAX_BASE:-}" ]; then + echo "[env] ERROR: PRIMUS_JAX_BASE is not set (no default). Set it to a directory" >&2 + echo "[env] you can write to with tens of GB free, then re-run:" >&2 + echo "[env] export PRIMUS_JAX_BASE=/big/disk/primus-jax-env" >&2 + # env.sh is always sourced (by setup.sh and to use the env); stop sourcing + # without killing an interactive shell. Fall back to exit if run directly. + # shellcheck disable=SC2317 # 'exit 1' is reachable when this file is executed, not sourced + return 1 2>/dev/null || exit 1 +fi +export PRIMUS_JAX_BASE +export VENV_DIR="${VENV_DIR:-$PRIMUS_JAX_BASE/venv}" +export WORKSPACE_DIR="${WORKSPACE_DIR:-$PRIMUS_JAX_BASE/workspace}" # kept checkouts (maxtext, Primus) +# Transient build sources: put on fast LOCAL /tmp (NFS is slow for compile I/O; +# these dirs are deleted after each build anyway). +export SRC_DIR="${SRC_DIR:-/tmp/primus-jax-build}" +# ROCm release tarball (TheRock dist) is extracted here. This replaces the +# Dockerfile's /opt/rocm and keeps the whole install no-sudo / user-writable. +export ROCM_DIR="${ROCM_DIR:-$PRIMUS_JAX_BASE/rocm}" + +# MaxText clone (deps + editable install live here). Primus resolves the +# MaxText backend from MAXTEXT_PATH (else its own third_party/maxtext), so we +# point it at the checkout we actually installed the deps for. +export MAXTEXT_DIR="${MAXTEXT_DIR:-$WORKSPACE_DIR/maxtext}" +export MAXTEXT_PATH="${MAXTEXT_PATH:-$MAXTEXT_DIR}" + +# Python interpreter used to CREATE the venv. MaxText needs >= 3.12. +# 3.12 is PREFERRED: it is the version the reference image (ubuntu:24.04) uses, +# and the prebuilt ROCm wheels (transformer_engine_rocm_jax, jax_rocm7_*) are +# built for cp312 — a 3.13 venv would need those built from source instead. +# Resolution order: +# 1. An explicit python3.12 / python3.13 on PATH (3.12 first). +# 2. A uv-managed interpreter (>= 3.12), if `uv` is installed. This is the +# no-sudo path on distros whose apt has no python3.12 (e.g. Ubuntu 22.04): +# `uv python install 3.12` provides one without root. setup.sh will run +# that install automatically if none is found yet. +# 3. Fall back to python3 (setup.sh errors clearly if that is < 3.12). +if [ -z "${PRIMUS_PYTHON:-}" ]; then + for _py in python3.12 python3.13; do + if command -v "$_py" >/dev/null 2>&1; then export PRIMUS_PYTHON="$_py"; break; fi + done + if [ -z "${PRIMUS_PYTHON:-}" ] && command -v uv >/dev/null 2>&1; then + _uv_py="$(uv python find '>=3.12' 2>/dev/null || true)" + [ -n "$_uv_py" ] && [ -x "$_uv_py" ] && export PRIMUS_PYTHON="$_uv_py" + fi + export PRIMUS_PYTHON="${PRIMUS_PYTHON:-python3}" +fi + +# Build parallelism. +export MAX_JOBS="${MAX_JOBS:-128}" + +# ---- Target GPU architecture (auto-detected) ---- +# PYTORCH_ROCM_ARCH controls the gfx targets used for the source builds (TE, +# RCCL) and TE's runtime NVTE_ROCM_ARCH. The name is kept for parity with the +# Dockerfile / TE which read it. If you set it yourself it is respected as-is: +# export PYTORCH_ROCM_ARCH="gfx942;gfx950" +# Otherwise it is auto-detected from the GPUs on this host. Detection needs no +# sudo and no system ROCm: it uses `rocminfo` when available (e.g. after the +# ROCm tarball is extracted) and otherwise falls back to the kernel KFD sysfs +# topology, so it works even on a fresh machine before any install. + +# Decode a KFD gfx_target_version integer (e.g. 90402) into a gfx name (gfx942). +_primus_gfxver_to_arch() { + local v="$1" + [ -n "$v" ] && [ "$v" != "0" ] || return 0 + printf 'gfx%d%x%x\n' "$(( v / 10000 ))" "$(( (v / 100) % 100 ))" "$(( v % 100 ))" +} + +# Echo a ";"-separated, de-duplicated list of detected gfx architectures. +_primus_detect_gpu_arch() { + local arches="" + if command -v rocminfo >/dev/null 2>&1; then + arches="$(rocminfo 2>/dev/null \ + | awk '/^[[:space:]]*Name:[[:space:]]*gfx/{print $2}' \ + | sort -u | paste -sd';' -)" + fi + if [ -z "$arches" ] && [ -d /sys/class/kfd/kfd/topology/nodes ]; then + local f v a list="" + for f in /sys/class/kfd/kfd/topology/nodes/*/properties; do + [ -f "$f" ] || continue + v="$(awk '/^gfx_target_version/{print $2}' "$f" 2>/dev/null)" + a="$(_primus_gfxver_to_arch "$v")" + [ -n "$a" ] && list="${list}${a}"$'\n' + done + arches="$(printf '%s' "$list" | sort -u | sed '/^$/d' | paste -sd';' -)" + fi + printf '%s' "$arches" +} + +if [ -z "${PYTORCH_ROCM_ARCH:-}" ]; then + _DETECTED_ARCH="$(_primus_detect_gpu_arch)" + if [ -n "$_DETECTED_ARCH" ]; then + export PYTORCH_ROCM_ARCH="$_DETECTED_ARCH" + echo "[env] detected GPU arch: PYTORCH_ROCM_ARCH=$PYTORCH_ROCM_ARCH" >&2 + else + export PYTORCH_ROCM_ARCH="gfx942;gfx950" + echo "[env] WARNING: no GPU detected; defaulting PYTORCH_ROCM_ARCH=$PYTORCH_ROCM_ARCH (override by exporting it)" >&2 + fi +fi + +# Comma-separated variant for the tools/vars that expect commas. +_ARCH_CSV="${PYTORCH_ROCM_ARCH//;/,}" +export ROCM_AMDGPU_TARGETS="${ROCM_AMDGPU_TARGETS:-$_ARCH_CSV}" +export GPU_ARCHS="${GPU_ARCHS:-$PYTORCH_ROCM_ARCH}" +export HCC_AMDGPU_TARGET="${HCC_AMDGPU_TARGET:-$_ARCH_CSV}" +export HIP_ARCHITECTURES="${HIP_ARCHITECTURES:-$_ARCH_CSV}" + +# Keep caches OFF the tiny home quota. Use local /tmp. +export PIP_CACHE_DIR="${PIP_CACHE_DIR:-/tmp/primus-jax-cache/pip}" +export XDG_CACHE_HOME="${XDG_CACHE_HOME:-/tmp/primus-jax-cache/xdg}" +export UV_CACHE_DIR="${UV_CACHE_DIR:-/tmp/primus-jax-cache/uv}" +mkdir -p "$PIP_CACHE_DIR" "$XDG_CACHE_HOME" "$UV_CACHE_DIR" 2>/dev/null || true + +# Activate the venv if it exists +if [ -f "$VENV_DIR/bin/activate" ]; then + # shellcheck disable=SC1091 + source "$VENV_DIR/bin/activate" +fi + +# NOTE: the two HSA_* vars below are NOT in the v26.5 Dockerfile — they are a +# carryover workaround for HSA_STATUS_ERROR_OUT_OF_RESOURCES. Harmless, but +# remove them if you want an exact match to the image. +export HSA_ENABLE_SCRATCH_ASYNC_RECLAIM=0 +export HSA_NO_SCRATCH_RECLAIM=1 + +# Flags to fix the ROCm profiler hang issue (from the JAX Dockerfile base stage) +export ROCPROFILER_QUEUE_INTERPOSITION=0 +export DEBUG_HIP_DYNAMIC_QUEUES=0 + +# ---- ROCm path: from the extracted TheRock release tarball (v26.5) ---- +# setup.sh's `rocm` stage downloads therock-dist-*.tar.gz and extracts it into +# $ROCM_DIR. Once present, wire the HIP/ROCm env at it (mirrors the Dockerfile's +# /opt/rocm layout, whose LD path is lib + lib/rocm_sysdeps/lib). +if [ -d "$ROCM_DIR/lib" ] || [ -x "$ROCM_DIR/bin/hipcc" ]; then + export ROCM_PATH="$ROCM_DIR" + export ROCM_HOME="$ROCM_DIR" # Primus uses ROCM_HOME; set both + 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" + # PATH: the Dockerfile puts /opt/rocm/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:${LD_LIBRARY_PATH:-}" + export LIBRARY_PATH="$ROCM_PATH/lib:$ROCM_PATH/lib64" + export CPATH="$HIP_INCLUDE_PATH" + export PKG_CONFIG_PATH="$ROCM_PATH/lib/pkgconfig" +fi + +# ---- TransformerEngine (ROCm) runtime settings for JAX ---- +# The transformer_engine_rocm_jax wheel is prebuilt, so these are runtime knobs. +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 (from the JAX Dockerfile) ---- +export GPU_MAX_HW_QUEUES=2 +export HIP_FORCE_DEV_KERNARG=1 +export HSA_FORCE_FINE_GRAIN_PCIE=1 +export NCCL_DEBUG=VERSION + +# NOT in the v26.5 Dockerfile: avoids a NaN-loss issue when training on gfx950 +# (no-op on gfx942). Keep it if you train on MI350/MI355; drop for exact parity. +export RCCL_WARP_SPEED_AUTO=0 + +# ---- XLA / JAX runtime settings ---- +export XLA_PYTHON_CLIENT_MEM_FRACTION=.9 +export XLA_FLAGS="${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=''}" diff --git a/tools/installation-jax/setup.sh b/tools/installation-jax/setup.sh new file mode 100644 index 000000000..a388450a4 --- /dev/null +++ b/tools/installation-jax/setup.sh @@ -0,0 +1,412 @@ +#!/usr/bin/env bash +# setup.sh — Reproduce the Primus JAX/MaxText training environment (v26.5) in a +# Python venv (no sudo, no docker). Mirrors the v26.5 JAX training Dockerfile, +# adapted for bare metal: +# * ROCm from a RELEASE TARBALL (TheRock dist) extracted to $ROCM_DIR, under +# $PRIMUS_JAX_BASE (no sudo, no /opt/rocm) — replaces the v26.4 pip wheels +# * builds/checkouts on a big disk (home quota is usually tiny) +# * GPU arch auto-detected (gfx942 and/or gfx950); apt/sudo steps skipped +# * MaxText's own setup.sh apt/interactive steps skipped (system packages are +# a one-time root action, documented in the guide's Section 2) +# * TensorFlow (2.21 CPU) and RCCL are built from source in the default flow +# (matching the v26.5 image) +# +# Usage: +# bash setup.sh # run all default stages in order +# bash setup.sh ... # run only specific stage(s), e.g. +# bash setup.sh venv rocm jax +# +# Stages are re-runnable. List them with: bash setup.sh --list + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +# shellcheck disable=SC1091 +source "$SCRIPT_DIR/env.sh" + +log() { echo -e "\n\033[1;36m[setup] $*\033[0m"; } +die() { echo -e "\033[1;31m[setup][ERROR] $*\033[0m" >&2; exit 1; } +# shellcheck disable=SC1091 +reload_env() { source "$SCRIPT_DIR/env.sh"; } + +# ---- pinned versions / commits (from the v26.5 JAX Dockerfile) ---- +# ROCm: TheRock RELEASE TARBALL (not pip wheels). Extracted into $ROCM_DIR. +# See: https://repo.amd.com/rocm/tarball-multi-arch/ +THEROCK_TARBALL="${THEROCK_TARBALL:-https://repo.amd.com/rocm/tarball-multi-arch/therock-dist-linux-multiarch-7.14.0.tar.gz}" + +# JAX + ROCm PJRT/plugin +# See: https://repo.amd.com/rocm/whl-multi-arch/jax-rocm7-pjrt/ +# https://repo.amd.com/rocm/whl-multi-arch/jax-rocm7-plugin/ +JAX_VERSION="0.10.0" +JAX_PJRT_VERSION="0.10.0+rocm7.14.0" +JAX_PLUGIN_VERSION="0.10.0+rocm7.14.0" +JAX_INDEX="https://repo.amd.com/rocm/whl-multi-arch/" + +# TransformerEngine (prebuilt ROCm JAX wheel) +# See: https://rocm.frameworks-nightlies.amd.com/whl-staging/device-all/transformer-engine-rocm-jax/ +TE_VERSION="2.15.0.dev0+rocm7.15.0a20260707.72d01a0" +TE_INDEX="https://rocm.frameworks-nightlies.amd.com/whl-staging/device-all/" +# From-source TransformerEngine (used by the `te_source` stage). Needed on hosts +# whose glibc is older than the prebuilt wheel requires (< 2.38, e.g. Ubuntu +# 22.04): building against the host toolchain links TE to the local glibc so it +# actually loads. (Commit is the last known-good ROCm TE JAX source build; bump +# to match the prebuilt wheel above when a matching source tag is published.) +TE_REPO="https://github.com/ROCm/TransformerEngine.git" +TE_SOURCE_COMMIT="${TE_SOURCE_COMMIT:-635d7c085c39a6d9bfe4881c7d3efab7a46d7129}" + +# TensorFlow (CPU) built from source — replaces the stock PyPI wheel, whose +# bundled LLVM collides with ROCm's libLLVM in Grain workers (SIGSEGV on +# `import tensorflow` after `import jax`). A CPU build also drops the bundled +# NCCL, preserving the XLA->RCCL collective fix. +TF_REPO="https://github.com/ROCm/tensorflow-upstream.git" +TF_BRANCH="upstream-v2.21.0" +BAZELISK_VERSION="v1.25.0" + +# RCCL built from source (rocm-systems), installed into the ROCm tree. +RCCL_REPO="https://github.com/ROCm/rocm-systems.git" +RCCL_COMMIT="9e5e4084a4b8e1e86551b0eb054725c62354a926" + +# MaxText (ROCm fork) +# v26.5 switched MaxText's initialize()/run() to a 2-value API (config, recorder); +# v26.4 and earlier returned 3 values (adds diagnostic_config). Primus main +# supports BOTH: MaxTextPretrainTrainer forwards initialize()'s tuple verbatim to +# run() (fix #912), so v26.5 trains out of the box with Primus main. Override +# MAXTEXT_BRANCH only if you deliberately need a different MaxText release. +MAXTEXT_REPO="https://github.com/ROCm/maxtext.git" +MAXTEXT_BRANCH="${MAXTEXT_BRANCH:-release/v26.5}" +# Which MaxText requirements set to install. The reference Docker image runs +# MaxText's setup.sh with defaults (DEVICE=tpu), which — on ROCm — pulls the +# framework-agnostic deps WITHOUT any CUDA packages. Override to `cuda12` only +# if you specifically need that set. +MAXTEXT_DEVICE="${MAXTEXT_DEVICE:-tpu}" + +# Primus +PRIMUS_REPO="https://github.com/AMD-AGI/Primus.git" +PRIMUS_BRANCH="main" + +PIP="python -m pip" +UVPIP="python -m uv pip" + +# Fresh clone helper: clone into transient SRC_DIR, build, then optionally clean. +fresh_clone() { # fresh_clone [extra git clone args...] + local url="$1"; local dir="$2"; shift 2 + rm -rf "${SRC_DIR:?}/$dir" + git clone "$@" "$url" "$SRC_DIR/$dir" +} + +# Return 0 if the host glibc is >= 2.38 (what the prebuilt TE/JAX wheels need). +# On failure to parse (unknown libc) we conservatively return non-zero so the +# caller falls back to the always-works from-source build. +_glibc_ge_238() { + local v; v="$(ldd --version 2>/dev/null | head -1 | grep -oE '[0-9]+\.[0-9]+$')" + [ -n "$v" ] || return 1 + awk -v v="$v" 'BEGIN{split(v,a,"."); exit !(a[1]>2 || (a[1]==2 && a[2]>=38))}' +} + +# ============================ STAGES ============================ + +stage_venv() { + # MaxText requires Python >= 3.12. + _venv_python_ok() { # _venv_python_ok -> 0 if >= 3.12 + local v; v="$("$1" -c 'import sys; print("%d.%d" % sys.version_info[:2])' 2>/dev/null)" || return 1 + case "$v" in 3.1[2-9]|3.[2-9][0-9]|[4-9].*) return 0 ;; *) return 1 ;; esac + } + + # If the resolved interpreter is too old, try to provision Python 3.12 with + # uv (no sudo) — the standard fix on Ubuntu 22.04 where apt has no 3.12. + if ! _venv_python_ok "$PRIMUS_PYTHON"; then + if command -v uv >/dev/null 2>&1; then + log "PRIMUS_PYTHON ($PRIMUS_PYTHON) is < 3.12; provisioning Python 3.12 via uv (no sudo)" + uv python install 3.12 || die "uv python install 3.12 failed" + PRIMUS_PYTHON="$(uv python find '>=3.12' 2>/dev/null || true)" + if [ -z "$PRIMUS_PYTHON" ] || [ ! -x "$PRIMUS_PYTHON" ]; then + die "could not locate a uv-managed Python >= 3.12 after install" + fi + else + die "MaxText requires Python >= 3.12 but no suitable interpreter was found. Options (no sudo): install uv (pip install --user uv) then re-run — setup.sh will fetch Python 3.12 automatically; or install python3.12 yourself and re-run with PRIMUS_PYTHON=python3.12" + fi + fi + log "Creating venv at $VENV_DIR (interpreter: $PRIMUS_PYTHON -> $("$PRIMUS_PYTHON" --version 2>&1))" + mkdir -p "$PRIMUS_JAX_BASE" "$SRC_DIR" "$WORKSPACE_DIR" || die "could not create the install dirs under PRIMUS_JAX_BASE=$PRIMUS_JAX_BASE (SRC_DIR=$SRC_DIR, WORKSPACE_DIR=$WORKSPACE_DIR). Set PRIMUS_JAX_BASE to a directory you can write to with tens of GB free, e.g. export PRIMUS_JAX_BASE=/big/disk/primus-jax-env" + # Recreate the venv if it exists but was built with a too-old Python. + if [ -f "$VENV_DIR/bin/activate" ] && ! _venv_python_ok "$VENV_DIR/bin/python"; then + log "Existing venv at $VENV_DIR is < Python 3.12; recreating it" + rm -rf "$VENV_DIR" + fi + [ -f "$VENV_DIR/bin/activate" ] || "$PRIMUS_PYTHON" -m venv "$VENV_DIR" + reload_env + $PIP install --upgrade pip + # Build front-end tooling (matches the Dockerfile; add uv, used by MaxText). + $PIP uninstall -y wheel || true + $PIP install \ + cmake==3.31.6 \ + ninja==1.11.1.3 \ + wheel==0.45.1 \ + packaging==25.0 \ + setuptools==69.5.1 \ + uv + rm -rf /root/.cache 2>/dev/null || true +} + +stage_rocm() { + reload_env + # v26.5: ROCm comes from a TheRock RELEASE TARBALL extracted into $ROCM_DIR + # (user-writable, no sudo, no /opt/rocm). This also sidesteps the pip + # rocm-sdk-* version-skew that broke hipBLASLt GEMMs in the v26.4 recipe. + log "Installing ROCm from tarball into $ROCM_DIR" + mkdir -p "$ROCM_DIR" "$SRC_DIR" + if [ -x "$ROCM_DIR/bin/hipcc" ] || [ -d "$ROCM_DIR/lib" ]; then + log "ROCm already present at $ROCM_DIR (skipping download; rm -rf it to force)" + else + local tb="$SRC_DIR/therock-dist.tar.gz" + log "Downloading $THEROCK_TARBALL" + wget -O "$tb" "$THEROCK_TARBALL" || die "ROCm tarball download failed" + log "Extracting to $ROCM_DIR" + tar -xf "$tb" -C "$ROCM_DIR" || die "ROCm tarball extract failed" + rm -f "$tb" + fi + reload_env + [ -n "${ROCM_PATH:-}" ] || die "ROCM_PATH not resolved after extraction (check $ROCM_DIR layout: expected $ROCM_DIR/bin, $ROCM_DIR/lib)" + log "ROCM_PATH=$ROCM_PATH" + # amdsmi (the Dockerfile installs it via pip after the ROCm tarball) + $PIP install amdsmi==7.0.2 || log "WARNING: amdsmi install failed (non-fatal)" + ( "$ROCM_PATH/bin/hipcc" --version || hipcc --version ) 2>/dev/null || log "hipcc not found yet; re-source env.sh" +} + +stage_jax() { + reload_env + log "Installing JAX ${JAX_VERSION} + ROCm PJRT/plugin (v26.5)" + # Note: JAX and related libraries need to be installed BEFORE TE and AFTER + # MaxText (whose setup.sh pulls in a stock jax/tensorflow we override here). + $PIP install "jax==${JAX_VERSION}" "jaxlib==${JAX_VERSION}" scipy==1.16 + $PIP install \ + --index-url "$JAX_INDEX" \ + --pre "jax_rocm7_pjrt==${JAX_PJRT_VERSION}" \ + --pre "jax_rocm7_plugin==${JAX_PLUGIN_VERSION}" + python -c "import jax; print('jax', jax.__version__); print('devices:', jax.devices())" || \ + log "WARNING: jax.devices() failed (expected if no GPU is visible on this build host)" +} + +stage_te() { + reload_env + # The prebuilt wheel targets the Dockerfile's ubuntu:24.04 base (glibc>=2.38). + # On older hosts (e.g. Ubuntu 22.04, glibc 2.35) it cannot load, so fall back + # transparently to the from-source build, which links against the host glibc. + # This keeps `bash setup.sh` (defaults) working on both 22.04 and 24.04. + if ! _glibc_ge_238; then + log "host glibc < 2.38 (prebuilt TE needs >= 2.38): building TransformerEngine from source instead (te_source)" + stage_te_source + return + fi + log "Installing TransformerEngine (JAX) ${TE_VERSION} (prebuilt ROCm wheel)" + $PIP install \ + pybind11==3.0.4 \ + importlib-metadata==8.7.1 \ + pydantic==2.13.4 \ + flax==0.12.2 + $PIP install \ + --index-url "$TE_INDEX" \ + --pre \ + --no-build-isolation \ + "transformer_engine_rocm_jax==${TE_VERSION}" + + # The prebuilt wheel targets the Dockerfile's ubuntu:24.04 base (glibc>=2.38, + # GCC 13/14 libstdc++). Verify it actually loads NOW — otherwise training + # dies later with a silent exit (the launcher swallows the OSError). + log "Verifying TransformerEngine loads" + if ! python -c "import transformer_engine.jax" 2>/tmp/te_import_err; then + cat /tmp/te_import_err >&2 || true + if grep -q "GLIBC_2" /tmp/te_import_err 2>/dev/null; then + die "TransformerEngine failed to import: your host glibc ($(ldd --version 2>/dev/null | head -1 | grep -oE '[0-9]+\.[0-9]+$')) is older than the prebuilt wheel requires (needs glibc>=2.38 / Ubuntu 24.04). Options: run on Ubuntu 24.04, or build TransformerEngine from source (see docs Section 3.7). glibc cannot be side-loaded via LD_LIBRARY_PATH." + fi + die "TransformerEngine failed to import (see error above). See docs Section 3.7." + fi + log "TransformerEngine import OK" +} + +stage_te_source() { + reload_env + # Build TransformerEngine (JAX) from source, linking against the HOST glibc. + # Use this instead of `te` on hosts whose glibc is older than the prebuilt + # wheel needs (< 2.38, e.g. Ubuntu 22.04). NOTE: this is a heavy build (CK + # fused-attention kernels for your arch) — expect ~30-60 min. + log "Building TransformerEngine (JAX) from source @ $TE_SOURCE_COMMIT" + $PIP install pybind11==3.0.4 importlib-metadata==8.7.1 pydantic==2.13.4 flax==0.12.2 + # Remove ALL prebuilt/stale TE variants so TE's install sanity-check sees a + # single, self-consistent from-source package. + $PIP uninstall -y transformer_engine transformer-engine transformer_engine_rocm7 \ + transformer_engine_rocm_jax transformer_engine_jax 2>/dev/null || true + fresh_clone "$TE_REPO" TransformerEngine --recursive + ( cd "$SRC_DIR/TransformerEngine" \ + && git checkout "$TE_SOURCE_COMMIT" \ + && git submodule update --init --recursive \ + && USE_ROCM=1 NVTE_FRAMEWORK=jax NVTE_USE_ROCM=1 \ + HIP_PATH="$ROCM_PATH" NVTE_ROCM_ARCH="$PYTORCH_ROCM_ARCH" \ + CMAKE_BUILD_PARALLEL_LEVEL="$MAX_JOBS" \ + PYTHONPATH="$SRC_DIR/TransformerEngine/3rdparty/hipify_torch:${PYTHONPATH:-}" \ + python3 setup.py bdist_wheel \ + && $PIP install --no-deps --force-reinstall dist/*.whl ) || die "TransformerEngine source build failed" + log "Verifying TransformerEngine (from source) loads" + python -c "import transformer_engine.jax" || die "TransformerEngine (source) import failed" + log "TransformerEngine (source) import OK" + rm -rf "$SRC_DIR/TransformerEngine" +} + +stage_maxtext() { + reload_env + log "Installing MaxText @ $MAXTEXT_BRANCH (kept in $MAXTEXT_DIR)" + rm -rf "$MAXTEXT_DIR" + git clone "$MAXTEXT_REPO" "$MAXTEXT_DIR" || die "MaxText clone failed" + ( cd "$MAXTEXT_DIR" && git checkout "$MAXTEXT_BRANCH" ) || die "MaxText checkout failed" + + # Replicate the Python portion of MaxText's src/dependencies/scripts/setup.sh + # for MODE=stable, WORKFLOW=pre-training (the reference image's default). + # The apt/gcsfuse and interactive-venv steps of that script are skipped: + # system packages are a one-time root action (see the guide's Section 2) + # and the venv already exists here. + local req="src/dependencies/requirements/generated_requirements/${MAXTEXT_DEVICE}-requirements.txt" + [ -f "$MAXTEXT_DIR/$req" ] || die "MaxText requirements not found: $MAXTEXT_DIR/$req" + ( cd "$MAXTEXT_DIR" \ + && $PIP install -U setuptools wheel uv \ + && $UVPIP install --resolution=lowest -r "$req" \ + && python -m src.dependencies.scripts.install_pre_train_extra_deps \ + && $UVPIP install --no-deps -e . ) || die "MaxText dependency install failed" +} + +stage_tf_cpu_fix() { + reload_env + # Fix: TF 2.20 sets RTLD_GLOBAL which exposes its bundled CUDA-targeting NCCL + # symbols globally, causing XLA's ncclCommInitRankConfig to resolve to TF's + # NCCL instead of the system ROCm RCCL. tensorflow-cpu has no bundled NCCL. + local tfver + tfver="$($PIP show tensorflow 2>/dev/null | awk '/^Version:/{print $2}')" + if [ -z "$tfver" ]; then + log "tensorflow not installed; skipping tensorflow-cpu override (run 'maxtext' stage first)" + return 0 + fi + log "Overriding tensorflow with tensorflow-cpu==$tfver (--no-deps) to avoid NCCL symbol clash" + $PIP install --no-deps "tensorflow-cpu==${tfver}" +} + +stage_tf_source() { + reload_env + # v26.5: build tensorflow-cpu 2.21 from ROCm's fork (bazel). 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. HEAVY: ~30-60 min. + # Needs a host clang/lld (clang-18) and unzip/zip — see the guide's Section 2. + # Idempotent: if the CPU build is already installed, skip the ~30-60 min + # bazel rebuild (pip uninstall tensorflow-cpu first to force a rebuild). + if $PIP show tensorflow-cpu >/dev/null 2>&1; then + log "tensorflow-cpu already installed ($($PIP show tensorflow-cpu 2>/dev/null | awk '/^Version:/{print $2}')); skipping source build" + return 0 + fi + log "Building tensorflow-cpu ($TF_BRANCH) from source with bazel (~30-60 min)" + local bz="$PRIMUS_JAX_BASE/bin/bazel" + mkdir -p "$PRIMUS_JAX_BASE/bin" + if [ ! -x "$bz" ]; then + log "Fetching bazelisk $BAZELISK_VERSION -> $bz" + wget -O "$bz" "https://github.com/bazelbuild/bazelisk/releases/download/${BAZELISK_VERSION}/bazelisk-linux-amd64" \ + || die "bazelisk download failed" + chmod +x "$bz" + fi + export PATH="$PRIMUS_JAX_BASE/bin:$PATH" + fresh_clone "$TF_REPO" tensorflow-upstream --depth 1 --branch "$TF_BRANCH" + # Build the wheel for the venv's ACTUAL Python (the reference image uses 3.12; + # deriving the version keeps the recipe working on hosts whose interpreter is + # 3.13). A cp312 wheel will not install into a cp313 venv and vice versa. + local pyver pytag + pyver="$(python -c 'import sys; print("%d.%d" % sys.version_info[:2])')" + pytag="cp${pyver//./}" # 3.12 -> cp312, 3.13 -> cp313 + # Keep bazel's (large) output tree off the tiny home quota. + ( cd "$SRC_DIR/tensorflow-upstream" \ + && "$bz" --output_user_root="$SRC_DIR/bazel" build //tensorflow/tools/pip_package:wheel \ + --repo_env=WHEEL_NAME=tensorflow_cpu \ + --repo_env=HERMETIC_PYTHON_VERSION="$pyver" ) || die "TensorFlow bazel build failed" + $PIP uninstall -y tensorflow tensorflow-cpu tensorflow_cpu || true + $PIP install --no-deps "$SRC_DIR"/tensorflow-upstream/bazel-bin/tensorflow/tools/pip_package/wheel_house/tensorflow_cpu-*-"${pytag}"-"${pytag}"-linux_x86_64.whl \ + || die "TensorFlow wheel install failed" + python -c "import tensorflow as tf; print('tensorflow', tf.__version__)" || log "WARNING: tensorflow import failed" + # bazel marks its external/downloaded trees read-only (dirs lose the write + # bit), so a plain `rm -rf` fails with "Permission denied". Make them + # writable first, and never let cleanup abort the install (TF is already + # built + installed at this point). + chmod -R u+w "$SRC_DIR/tensorflow-upstream" "$SRC_DIR/bazel" 2>/dev/null || true + rm -rf "$SRC_DIR/tensorflow-upstream" "$SRC_DIR/bazel" 2>/dev/null || true +} + +stage_rccl() { + reload_env + # v26.5: build RCCL from source (rocm-systems) and drop the libraries into the + # ROCm tree so JAX/XLA's collectives use it. Requires the ROCm toolchain + # (hipcc) from the extracted tarball -> run AFTER the `rocm` stage. + [ -n "${ROCM_PATH:-}" ] || die "ROCM_PATH not set; run the 'rocm' stage first" + log "Building RCCL from source @ $RCCL_COMMIT -> $ROCM_PATH/lib" + fresh_clone "$RCCL_REPO" rocm-systems + ( cd "$SRC_DIR/rocm-systems" \ + && git checkout "$RCCL_COMMIT" \ + && cd projects/rccl \ + && ./install.sh -l --prefix build/ --amdgpu_targets="$PYTORCH_ROCM_ARCH" \ + && cp -r build/release/librccl* "$ROCM_PATH/lib/" ) || die "RCCL build failed" + rm -rf "$SRC_DIR/rocm-systems" +} + +stage_primus() { + reload_env + log "Installing Primus @ $PRIMUS_BRANCH (kept in $WORKSPACE_DIR/Primus)" + rm -rf "$WORKSPACE_DIR/Primus" + git clone --recurse-submodules "$PRIMUS_REPO" "$WORKSPACE_DIR/Primus" || die "Primus clone failed" + ( cd "$WORKSPACE_DIR/Primus" \ + && git checkout "$PRIMUS_BRANCH" \ + && git submodule update --init third_party/maxtext/ ) || die "Primus checkout failed" + # The JAX Dockerfile does NOT pip install Primus' (torch-oriented) + # requirements.txt; the JAX runtime deps live in requirements-jax.txt + # (installed by the `jaxreqs` stage). It also removes stale dataclasses + # backports that conflict on modern Python. + $PIP uninstall -y dataclasses dataclasses_json || true +} + +stage_jaxreqs() { + reload_env + # Primus' JAX/MaxText runtime deps (normally installed by the Primus + # MaxText pre-train hook at launch time). Front-load them here. + local req="$WORKSPACE_DIR/Primus/requirements-jax.txt" + if [ -f "$req" ]; then + log "Installing Primus JAX requirements from $req" + $PIP install -r "$req" + else + log "Primus requirements-jax.txt not found (skipping); run 'primus' stage first" + fi +} + +stage_manifest() { + reload_env + log "Writing manifest to $WORKSPACE_DIR/.manifest" + mkdir -p "$WORKSPACE_DIR/.manifest" + env > "$WORKSPACE_DIR/.manifest/env.txt" + $PIP list > "$WORKSPACE_DIR/.manifest/requirements.txt" + cp "$SCRIPT_DIR/env.sh" "$WORKSPACE_DIR/.manifest/env.sh" +} + +# v26.5 order mirrors the Dockerfile: ROCm tarball -> MaxText -> TF-from-source +# -> JAX (overrides what MaxText pulled) -> TE -> Primus -> RCCL-from-source. +DEFAULT_STAGES=(venv rocm maxtext tf_source jax te primus jaxreqs rccl manifest) + +run_stage() { local s="$1"; local fn="stage_$s"; declare -F "$fn" >/dev/null || die "unknown stage: $s"; "$fn"; } + +main() { + if [ "${1:-}" = "--list" ]; then + echo "default: ${DEFAULT_STAGES[*]}" + echo "note: the 'te' stage auto-falls-back to a from-source build on glibc < 2.38 hosts (e.g. Ubuntu 22.04)" + echo "optional: te_source (force the from-source TransformerEngine build regardless of glibc)" + echo "optional: tf_cpu_fix (lighter alternative to tf_source: pip tensorflow-cpu instead of the ~30-60 min bazel build)" + exit 0 + fi + local stages=("$@"); [ ${#stages[@]} -eq 0 ] && stages=("${DEFAULT_STAGES[@]}") + log "Base dir: $PRIMUS_JAX_BASE | arch: $PYTORCH_ROCM_ARCH | stages: ${stages[*]}" + df -h "$PRIMUS_JAX_BASE" 2>/dev/null | tail -1 || true + for s in "${stages[@]}"; do run_stage "$s"; done + log "DONE. Activate later with: source $SCRIPT_DIR/env.sh" +} + +main "$@"