From e87244b9afdc2b37033c29c98c1fb1080d2e02af Mon Sep 17 00:00:00 2001 From: Fuyuan Jing Date: Thu, 23 Jul 2026 17:52:35 +0000 Subject: [PATCH 1/2] docs(installation-jax): add bare-metal JAX/MaxText install guide and scripts (v26.5) Add a no-Docker, no-sudo bare-metal installation path for the Primus JAX / MaxText training backend, mirroring the v26.5 JAX training Dockerfile. New tools/installation-jax/: - env.sh: sources the venv and exports ROCm (TheRock tarball), NVTE, and XLA runtime env; auto-detects the GPU arch (rocminfo -> KFD sysfs) so it works on gfx942/gfx950 without hand-editing; prefers Python 3.12 (prebuilt wheels are cp312). - setup.sh: re-runnable staged installer (ROCm tarball, MaxText, TensorFlow-CPU from source, ROCm JAX/PJRT/plugin, TransformerEngine, Primus, RCCL from source). The 'te' stage auto-falls-back to a from-source build on glibc < 2.38 hosts (e.g. Ubuntu 22.04); the tf_source build derives its Python wheel tag from the venv interpreter (cp312/cp313). - README.md: rationale and stage reference. New docs/01-getting-started/bare-metal-installation-jax.md: full manual guide with pinned versions, a version-requirements table (Python/glibc/toolchain/GPU prerequisites), and single- vs multi-node breakdown. Cross-linked from installation.md and bare-metal-installation.md (PyTorch guide relabeled). Validated end-to-end for single-node MaxText pretraining on gfx942 / Ubuntu 22.04 (glibc 2.35) / Python 3.12. Primus main already handles MaxText v26.5's 2-value initialize()/run() API (#912). Co-authored-by: Cursor --- .../bare-metal-installation-jax.md | 868 ++++++++++++++++++ .../bare-metal-installation.md | 2 + docs/01-getting-started/installation.md | 3 +- tools/installation-jax/README.md | 169 ++++ tools/installation-jax/env.sh | 179 ++++ tools/installation-jax/setup.sh | 412 +++++++++ 6 files changed, 1632 insertions(+), 1 deletion(-) create mode 100644 docs/01-getting-started/bare-metal-installation-jax.md create mode 100644 tools/installation-jax/README.md create mode 100644 tools/installation-jax/env.sh create mode 100644 tools/installation-jax/setup.sh 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..25dd6a997 --- /dev/null +++ b/docs/01-getting-started/bare-metal-installation-jax.md @@ -0,0 +1,868 @@ +# 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` +(`unified-training-dockers/jax_training_the_rock/docker/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). The +**default `PRIMUS_JAX_BASE` is site-specific** and may not exist on your host, so +set it to a directory you can write to that has tens of GB free: + +```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. + +> **Changed from v26.4:** the v26.4 recipe installed ROCm from pip +> `rocm-sdk-devel` **nightly** wheels. v26.5 uses a pinned **release tarball**, +> which is more stable and avoids the pip version-skew that could break +> hipBLASLt GEMMs. RCCL is then rebuilt from source (see Section 3.9) 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 v26.4/v26.5 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** (the v26.4 recipe's +> `rocm-sdk-*` nightly wheels could drift out of sync and break hipBLASLt GEMMs; +> the pinned tarball avoids that entirely). + +### 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.** v26.5 switched MaxText's +> `initialize()`/`run()` to a **2-value** API (`config, recorder`); v26.4 and +> earlier used a **3-value** API (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. You only need +> `MAXTEXT_BRANCH=release/v26.4 bash setup.sh maxtext` if you deliberately want +> to pin an older 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..c4a538f76 --- /dev/null +++ b/tools/installation-jax/README.md @@ -0,0 +1,169 @@ +# 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. + +> **v26.5 vs v26.4:** ROCm moved from pip `rocm-sdk-devel` **nightly** wheels to a +> pinned release **tarball** (`repo.amd.com/.../therock-dist-linux-multiarch-7.14.0.tar.gz`); +> JAX 0.10.2 → **0.10.0** (`repo.amd.com/rocm/whl-multi-arch/`); TE → **2.15.0.dev0**; +> TensorFlow is now **built from source** (2.21 CPU, bazel) and **RCCL** is rebuilt +> from source; MaxText → **release/v26.5** (Primus `main` supports its 2-value API). + +> **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 +# Default base is /it-share-4/envs/primus-jax-env (persistent shared disk). +# This default is site-specific; override it for your host: +# 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 +# Use the SAME PRIMUS_JAX_BASE you built with (export it first if you overrode the default) +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.** v26.5 MaxText uses a **2-value** +> `initialize()`/`run()` API; v26.4 and earlier used a **3-value** API. Primus +> `main` supports **both** — `MaxTextPretrainTrainer` forwards `initialize()`'s +> tuple verbatim to `run()` (fix #912) — so this recipe's pinned `release/v26.5` +> trains out of the box. Use `MAXTEXT_BRANCH=release/v26.4 bash setup.sh maxtext` +> only if you deliberately want to pin an older 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..e7acd439b --- /dev/null +++ b/tools/installation-jax/env.sh @@ -0,0 +1,179 @@ +#!/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: choose a BASE on a big disk. Override by exporting PRIMUS_JAX_BASE first. +# MaxText requires Python >= 3.12; see PRIMUS_PYTHON below. + +# ---- Install location (persistent) + transient build sources ---- +export PRIMUS_JAX_BASE="${PRIMUS_JAX_BASE:-/it-share-4/envs/primus-jax-env}" # persistent shared disk +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..919a6ea36 --- /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" + # 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 "$@" From 44842f34412bc3ac5f89c4dcf9e2e2dfa01a08f2 Mon Sep 17 00:00:00 2001 From: Fuyuan Jing Date: Fri, 24 Jul 2026 14:19:39 +0000 Subject: [PATCH 2/2] make PRIMUS_JAX_BASE required; update documents --- .../bare-metal-installation-jax.md | 36 ++++++++----------- tools/installation-jax/README.md | 26 ++++++-------- tools/installation-jax/env.sh | 20 +++++++++-- tools/installation-jax/setup.sh | 2 +- 4 files changed, 43 insertions(+), 41 deletions(-) diff --git a/docs/01-getting-started/bare-metal-installation-jax.md b/docs/01-getting-started/bare-metal-installation-jax.md index 25dd6a997..379f97f0f 100644 --- a/docs/01-getting-started/bare-metal-installation-jax.md +++ b/docs/01-getting-started/bare-metal-installation-jax.md @@ -5,8 +5,7 @@ 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` -(`unified-training-dockers/jax_training_the_rock/docker/Dockerfile`) and +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 @@ -70,9 +69,9 @@ There are two files: ### 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). The -**default `PRIMUS_JAX_BASE` is site-specific** and may not exist on your host, so -set it to a directory you can write to that has tens of GB free: +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) @@ -199,11 +198,9 @@ extracted into a user-writable directory (`$ROCM_DIR`, default `jax_rocm7_plugin` wheels (from `repo.amd.com/rocm/whl-multi-arch/`) provide GPU-accelerated JAX built against that ROCm. -> **Changed from v26.4:** the v26.4 recipe installed ROCm from pip -> `rocm-sdk-devel` **nightly** wheels. v26.5 uses a pinned **release tarball**, -> which is more stable and avoids the pip version-skew that could break -> hipBLASLt GEMMs. RCCL is then rebuilt from source (see Section 3.9) and dropped -> into this ROCm tree. +> 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: @@ -248,7 +245,7 @@ scripts use the same pins; change one and you may have to change the others. | 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 v26.4/v26.5 compatibility shim. | +| 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). | @@ -452,9 +449,8 @@ 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** (the v26.4 recipe's -> `rocm-sdk-*` nightly wheels could drift out of sync and break hipBLASLt GEMMs; -> the pinned tarball avoids that entirely). +> 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 @@ -567,13 +563,11 @@ 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.** v26.5 switched MaxText's -> `initialize()`/`run()` to a **2-value** API (`config, recorder`); v26.4 and -> earlier used a **3-value** API (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. You only need -> `MAXTEXT_BRANCH=release/v26.4 bash setup.sh maxtext` if you deliberately want -> to pin an older MaxText release. +> **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 diff --git a/tools/installation-jax/README.md b/tools/installation-jax/README.md index c4a538f76..e9a60e967 100644 --- a/tools/installation-jax/README.md +++ b/tools/installation-jax/README.md @@ -22,12 +22,6 @@ 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. -> **v26.5 vs v26.4:** ROCm moved from pip `rocm-sdk-devel` **nightly** wheels to a -> pinned release **tarball** (`repo.amd.com/.../therock-dist-linux-multiarch-7.14.0.tar.gz`); -> JAX 0.10.2 → **0.10.0** (`repo.amd.com/rocm/whl-multi-arch/`); TE → **2.15.0.dev0**; -> TensorFlow is now **built from source** (2.21 CPU, bazel) and **RCCL** is rebuilt -> from source; MaxText → **release/v26.5** (Primus `main` supports its 2-value API). - > **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: @@ -56,9 +50,9 @@ system ROCm or apt for the ROCm toolchain itself. ```bash cd tools/installation-jax -# Default base is /it-share-4/envs/primus-jax-env (persistent shared disk). -# This default is site-specific; override it for your host: -# export PRIMUS_JAX_BASE=/some/big/disk/primus-jax-env +# 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 ``` @@ -74,7 +68,8 @@ bash setup.sh venv rocm jax # venv + ROCm + JAX only ## Use the environment afterward ```bash -# Use the SAME PRIMUS_JAX_BASE you built with (export it first if you overrode the default) +# 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())" @@ -127,12 +122,11 @@ Optional / alternative stages: bash setup.sh venv rocm maxtext tf_cpu_fix jax te primus jaxreqs rccl manifest ``` -> **MaxText v26.5 and Primus.** v26.5 MaxText uses a **2-value** -> `initialize()`/`run()` API; v26.4 and earlier used a **3-value** API. Primus -> `main` supports **both** — `MaxTextPretrainTrainer` forwards `initialize()`'s -> tuple verbatim to `run()` (fix #912) — so this recipe's pinned `release/v26.5` -> trains out of the box. Use `MAXTEXT_BRANCH=release/v26.4 bash setup.sh maxtext` -> only if you deliberately want to pin an older MaxText release. +> **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) diff --git a/tools/installation-jax/env.sh b/tools/installation-jax/env.sh index e7acd439b..b9646ec78 100644 --- a/tools/installation-jax/env.sh +++ b/tools/installation-jax/env.sh @@ -8,11 +8,25 @@ # (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: choose a BASE on a big disk. Override by exporting PRIMUS_JAX_BASE first. -# MaxText requires Python >= 3.12; see PRIMUS_PYTHON below. +# 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 ---- -export PRIMUS_JAX_BASE="${PRIMUS_JAX_BASE:-/it-share-4/envs/primus-jax-env}" # persistent shared disk +# 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; diff --git a/tools/installation-jax/setup.sh b/tools/installation-jax/setup.sh index 919a6ea36..a388450a4 100644 --- a/tools/installation-jax/setup.sh +++ b/tools/installation-jax/setup.sh @@ -127,7 +127,7 @@ stage_venv() { 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" + 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"