How to Scale Your Model on AMD

An AMD companion to How to Scale Your Model, performed on MI355X (Chapter 0: Intro | Chapter 1: MI355X architecture and system topology)

How MI355X hardware, ROCm kernels, JAX compilation, and MaxText configuration affect training throughput.

Motivation

The renowned JAX Scaling Book teaches a comprehensive understanding on TPU architecure, rooflines, Transformer analysis, sharding behaviour, training parallelism, and much more.

This ROCm book serves as complementary material for readers to understand the ROCm stack for JAX, including ROCm experimental features, and in general how to extract the best training performance out of AMD GPU’s with JAX.

But perhaps more importantly, readers should be aware of the performance features available within the JAX/XLA/MaxText ecosystem, understand their underlying mechanisms, and know where to find further information.

Given the pace of active development, the JAX/ROCm software stack evolves rapidly, with new optimizations and capabilities often arriving faster than documentation can be updated. Consequently, the goal is not to present current performance recommendations, but to provide readers with the context necessary to discover new developments, evaluate their applicability, and incorporate optimizations suited to their own training environments.

Reader and Prerequisites

The intended reader knows Python, basic JAX, and Transformer block architecture. The book introduces AMD hardware, the ROCm execution path, and the sharding and profiling details needed by the experiments. Readers who want first-principles explanations and derivations should read the JAX Scaling Book chapters linked throughout.

Chapters

1. MI355X architecture and system topology explains CDNA 4 matrix execution, memory, native precision formats, and AMD’s Infinity Fabric topology.
2. The JAX/ROCm Stack follows a JAX program through tracing, StableHLO, XLA, ROCm libraries, FFI, and device execution.
3. Profiling a Training Step connects roofline estimates to XProf, rocprofv3, hardware counters, and rocprof-compute.
4. Training in Mixed Precision explains training in BF16, FP16, FP8, MXFP8, MXFP6, and MXFP4, accompanied with case studies on Llama 70B.
5. Sharding and Parallelism explains different parallelism strategies through sharding, accompanied with case studies on Mixtral 8x22B
6. Memory and Kernel Optimizations covers rematerialization, ROCm attention backends, and grouped GEMM lowering for MoE models.

Attribution

The book reuses concepts and, where noted, adapted material from the MIT-licensed JAX Scaling Book. Citations accompany reused derivations and figures. AMD, JAX, OpenXLA, and OCP specifications are cited where their facts are used.

Next: Chapter 1, MI355X architecture and system topology.