Installation and setup

This guide covers supported platforms, prerequisites, and how to set up the training environment: container (recommended) and bare metal, plus multi-node distributed training (Slurm recommended).

---

Supported platforms

Requirement	Notes
OS	Linux (ROCm-supported distributions per AMD documentation).
ROCm	≥ 7.0 recommended.
GPUs	AMD Instinct™ MI300X, MI325X, MI355X (or other ROCm-supported Instinct SKUs your site supports).

---

Prerequisites

Prerequisite	Purpose
AMD Instinct GPUs	Training and benchmarks execute on GPU.
ROCm drivers and user-space stack	Required for HIP, RCCL, and ML frameworks.
Docker ≥ 24.0 (or Podman with compatible GPU passthrough)	Container mode and reproducible environments.
git	Clone the repository and submodules.

Quick environment checks

rocm-smi
docker --version

rocm-smi should list your GPUs; docker --version should report 24.0 or newer.

---

Container setup (recommended)

AMD publishes training Docker images monthly, providing a consistent, ready-to-run environment optimized for AMD GPUs. It is recommended to use the AMD-published training Docker images together with this Primus-LM repository to run your training jobs. The images support pre-training and post-training workflows with multiple backends including Megatron-LM, TorchTitan, and JAX MaxText, alongside ROCm-optimized components.

Check the AMD-published training Docker images here:

• For Megatron-LM and TorchTitan backends: https://hub.docker.com/r/rocm/primus/tags

• For MaxText backend: https://hub.docker.com/r/rocm/jax-training/tags

1. Pull the image

# For Megatron-LM and TorchTitan backends
docker pull rocm/primus:v26.4
# For MaxText backend
docker pull rocm/jax-training:maxtext-v26.4-jax0.9.1-te2.12.0

2. Clone the repository

Submodules are required for third-party backends and tools:

git clone --recurse-submodules https://github.com/AMD-AGI/Primus.git
cd Primus
# checkout the branch for the specific release
git checkout release/v26.4
git submodule update --init --recursive

3. Run a verification benchmark

From the repository root:

./primus-cli container --image rocm/primus:v26.4 -- \
  benchmark gemm --M 4096 --N 4096 --K 4096

A successful run validates the GPU stack and Primus CLI wiring without launching a full training job.

---

Bare-metal (host) setup

The container setup above is strongly recommended. The AMD-published training Docker image is the tested, reproducible, and best-supported path. Build the full stack on a bare-metal host only when containers are not an option (for example, due to policy or operational constraints).

What to expect

Reproducing the training environment on the host means building the same stack the Docker image ships, mostly from source. This is a long, build-heavy process that requires:

• Several source-built kernel libraries (Flash Attention, TransformerEngine, AITER, Primus-Turbo, Grouped GEMM, causal-conv1d, Mamba) compiled against ROCm.

• A machine with many CPU cores, ample RAM, and tens of GB of free disk.

• A long time (expect a multi-hour first build).

What a complete host environment needs

Layer	What it provides	How it is installed
Kernel/hardware	AMD GPU driver (amdgpu KMD) and device access (/dev/kfd, /dev/dri)	OS/administrator (root, one-time)
OS libraries	Build toolchain and runtime libraries (g++, git, RDMA, hwloc, etc.)	apt (root, one-time)
ROCm user-space	rocm-sdk-devel + device wheels—no system-wide ROCm install required	pip (TheRock wheels, in venv)
Deep learning framework	ROCm-enabled PyTorch (torch, torchvision, torchaudio, apex)	pip (TheRock wheels, in venv)
Accelerated kernels	Flash Attention, TransformerEngine, AITER, Primus-Turbo, Grouped GEMM, Mamba	build from source (in venv)
Multi-node communications	UCX, OpenMPI, rocSHMEM, AMD AINIC—only for distributed (multi-node) training	build from source or apt
Primus + Python dependencies	Primus, submodules, and training libraries (datasets, transformers, wandb, etc.)	git + pip (in venv)

General approach

1. System packages (root, one-time): install the build toolchain and, for multi-node, the RDMA/networking libraries via apt. The GPU kernel driver must already be loaded.

2. Python virtual environment (no root): create a venv, then install ROCm and PyTorch from AMD’s TheRock multi-arch wheels—this replaces a system ROCm install and keeps everything unprivileged.

3. Build the accelerated kernels from source against the ROCm in venv and your GPU architecture (gfx942 for MI300X/MI325X, gfx950 for MI350X/MI355X).

4. Install Primus and its Python dependencies, then persist the required environment variables (ROCm paths, NVTE_* flags) in your venv activation script.

5. (Optional) Build the multi-node communication stack (UCX, OpenMPI, rocSHMEM) only if you require RDMA-based distributed training.

Detailed instructions

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)

Verify

After the build, validate the environment and run a benchmark directly on the host (no container):

./primus-cli direct -- benchmark gemm --M 4096 --N 4096 --K 4096

---

Multi-node distributed training (Slurm recommended)

For training jobs that span multiple nodes, we recommend using Slurm. Slurm is a cluster workload manager and job scheduler: it allocates nodes and GPUs, places your job on them, launches one task per node, and injects the topology information (node list, node count, per-node rank) that distributed PyTorch needs. primus-cli has a built-in slurm mode that wraps srun/sbatch and wires this topology into the training launcher for you.

Cluster baseline

Before launching distributed jobs, ensure every participating node has:

• The same software stack—use the same container image on all nodes (recommended), or an identical bare-metal install (see sections above).

• A shared filesystem for code, datasets, checkpoints, and logs (e.g. NFS/Lustre), mounted at the same path on every node.

• Working inter-node networking: for best performance, use a high-speed RDMA fabric (InfiniBand, RoCE, or AMD AINIC) and ensure RCCL can select the right interface.

Setting up Slurm

Setting up Slurm itself is a cluster-administration task and is outside Primus’s scope. Follow the official documentation:

• Slurm Quick Start (users)

• Slurm Quick Start Administrator Guide (install & configure)

Once sinfo and srun work on your login node, Primus can submit jobs to it. If you don’t administer the cluster, your site administrator typically provides the partition, account, and reservation names you need.

Launching with primus-cli slurm

The Slurm wrapper uses a single -- separator:

• Before the --: put the launcher (srun or sbatch, default srun) and Slurm flags (-N, -p, --nodelist, --account, --qos, --reservation, …).

• After the --: put the Primus command to run (train / benchmark / …). It runs inside the container image (see Selecting the container image below).

cd /path/to/Primus

# Pretrain on 2 nodes via srun
./primus-cli slurm srun -N 2 \
  -- train pretrain \
  --config examples/megatron/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml

Add the global --dry-run flag before the launcher to print the exact command without executing it:

./primus-cli --dry-run slurm srun -N 2 \
  -- train pretrain --config <your-config>.yaml

See runner/README.md and the CLI reference for the full set of launcher flags and scenarios.

Selecting the container image

The image used for the job is resolved according to the following priority order (highest first):

1. DOCKER_IMAGE environment variable—overrides everything else. This is the simplest way to switch images and it propagates to all nodes:

export DOCKER_IMAGE=rocm/primus:v26.4
./primus-cli slurm srun -N 2 \
  -- train pretrain --config <your-config>.yaml

2. --image CLI flag—place it immediately after the --, before the Primus command (ignored if DOCKER_IMAGE is set):

./primus-cli slurm srun -N 2 \
  -- --image rocm/primus:v26.4 train pretrain --config <your-config>.yaml

3. Config file default—container.options.image in runner/.primus.yaml (or your ~/.primus.yaml), which is set to rocm/primus:v26.4 by default.

Distributed environment variables

Primus launches training with torchrun, which needs to know the cluster topology. These are the key variables:

Variable	Role	Default
MASTER_ADDR	Hostname/IP of rank 0; all ranks rendezvous here.	localhost
MASTER_PORT	Port on the master used for rendezvous.	1234
NNODES	Number of nodes in the job.	1
NODE_RANK	Index of this node (0-based, unique per node).	0
GPUS_PER_NODE	GPUs (processes) to launch per node.	8

The total number of training processes (world size) is NNODES × GPUS_PER_NODE.

Under Slurm, the values of these variables are derived automatically. primus-cli slurm reads Slurm’s own variables and sets the Primus ones for you: NNODES from SLURM_NNODES, NODE_RANK from SLURM_NODEID, and MASTER_ADDR from the first host in SLURM_NODELIST (with MASTER_PORT defaulting to 1234). You normally only need to set GPUS_PER_NODE if your nodes don’t have 8 GPUs. You can still override MASTER_PORT (e.g. to avoid a port clash) via --env.

Without Slurm: Kubernetes or parallel-ssh

Slurm is recommended but not required. The mechanism underneath is simple: set the same distributed environment variables on every node, point them all at the same MASTER_ADDR, give each node a unique NODE_RANK, and run the same primus-cli direct command on each node. Any tool that can run a command across nodes works—for example Kubernetes (e.g. a PyTorchJob / indexed Job) or parallel-ssh/pdsh.

For a 2-node job you would run, on the master (rank 0):

export NNODES=2 GPUS_PER_NODE=8 NODE_RANK=0 MASTER_ADDR=<node0-host> MASTER_PORT=1234
./primus-cli direct -- train pretrain --config <your-config>.yaml

and on the worker (rank 1), the same command with NODE_RANK=1 and the same MASTER_ADDR:

export NNODES=2 GPUS_PER_NODE=8 NODE_RANK=1 MASTER_ADDR=<node0-host> MASTER_PORT=1234
./primus-cli direct -- train pretrain --config <your-config>.yaml

With Kubernetes, inject these as container environment variables (deriving NODE_RANK from the pod’s index); with parallel-ssh, pass them per host. The training command itself is identical on every node.

Other important considerations

• Cluster validation. Run the built-in preflight check across your nodes before a long job: ./primus-cli slurm srun -N <N> -- preflight.

• Networking/RCCL. On RDMA fabrics, make sure the correct interface is selected (Primus should auto-detect this; if it doesn’t, set NCCL_SOCKET_IFNAME / NCCL_IB_HCA). Use NCCL_DEBUG=INFO (passed via --env) to diagnose hangs at startup.

• RDMA limits. High-performance networking usually needs locked-memory limits raised (ulimit -l unlimited) and sometimes hugepages—configured by your admin.

• MASTER_PORT must be free on the master node and reachable from all workers; firewalls between nodes will cause rendezvous timeouts.

• Hugging Face access. If your configuration downloads gated models or tokenizers, export HF_TOKEN (and ensure it is propagated to all nodes and into the containers).

---

Post-installation verification checklist (for all setup approaches)

Step	Check
ROCm	rocm-smi shows expected GPUs and no driver errors.
Container engine	docker run --rm ... rocm/primus:v26.4 (or your site’s GPU test) succeeds.
GEMM benchmark	./primus-cli container or direct benchmark completes (see sections above).
Preflight	Run preflight diagnostics: ./primus-cli direct -- preflight (single node) or ./primus-cli slurm srun -N <N> -- preflight (cluster).

If your training pulls models or tokenizers from Hugging Face Hub, configure tokens (for example HF_TOKEN) in the environment or container flags as required by your configuration.

---

Related documentation

• Overview

• Quickstart

• CLI reference