PaperlayerView exact revision on arXiv

Source-linked AI summary

Training Foundation Models on a Full-Stack AMD Platform: Compute, Networking, and System Design

Quentin Anthony, Yury Tokpanov, Skyler Szot, Srivatsan Rajagopal, Praneeth Medepalli, Anna Golubeva, Vasu Shyam, Robert Washbourne, Rishi Iyer, Ansh Chaurasia, Tomas Figliolia, Xiao Yang, Abhinav Sarje, Drew Thorstensen, Amartey Pearson, Zack Grossbart, Jason van Patten, Emad Barsoum, Zhenyu Gu, Yao Fu, Beren Millidge

arXiv:2511.17127v2cs.CLcs.AIcs.DC

TL;DR

Large-scale LLM pretraining depends on cluster throughput and practical systems choices, yet operational challenges across compute, networking, and software require detailed characterization. This paper studies pretraining on MI300X and Pollara, with system measurements and hardware-aware model and training-stack design. ZAYA1-base is reported to perform competitively with a strong 4B dense model despite having more than four times fewer active parameters.

  • Problem

    Base-model quality is substantially influenced by pretraining, whose quality is primarily bottlenecked by cluster throughput and the training software stack.

    Source: §I · p. 1

  • Method

    The paper presents a production-scale AMD pretraining case study, measuring system behavior and describing hardware-aware model design and training-stack choices.

    Source: §I · pp. 1–2

  • Results

    ZAYA1-base, with approximately 760M active and 8.3B total parameters, performs competitively with Qwen3-4B despite having more than four times fewer active parameters.

    Source: §VI · p. 18

  • Takeaways & Limitations

    The authors conclude that AMD hardware, software, and networking are sufficiently mature and robust for large-scale LLM pretraining.

    Source: §VII · p. 20

Abstract

from arXiv · show

We report on the first large-scale mixture-of-experts (MoE) pretraining study on pure AMD hardware, utilizing both MI300X GPUs and Pollara networking. We distill practical guidance for both systems and model design. On the systems side, we deliver a comprehensive cluster and networking characterization: microbenchmarks for all core collectives (all-reduce, reduce-scatter, all-gather, broadcast) across message sizes and GPU counts over Pollara. To our knowledge, this is the first at this scale. We further provide MI300X microbenchmarks on kernel sizing and memory bandwidth to inform model design. On the modeling side, we introduce and apply MI300X-aware transformer sizing rules for attention and MLP blocks and justify MoE widths that jointly optimize training throughput and inference latency. We describe our training stack in depth, including often-ignored utilities such as fault-tolerance and checkpoint-reshaping, as well as detailed information on our training recipe. We also provide a preview of our model architecture and base model - ZAYA1 (760M active, 8.3B total parameters MoE, available at https://huggingface.co/Zyphra/ZAYA1-base) - which will be further improved upon in forthcoming papers. ZAYA1-base achieves performance comparable to leading base models such as Qwen3-4B and Gemma3-12B at its scale and larger, and outperforms models including Llama-3-8B and OLMoE across reasoning, mathematics, and coding benchmarks. Together, these results demonstrate that the AMD hardware, network, and software stack are mature and optimized enough for competitive large-scale pretraining.

I. INTRODUCTION

The paper presents a measurement-driven case study of large-scale LLM pretraining on AMD hardware and networking, focusing on practical system behavior and design choices. Its contributions span Pollara and MI300X characterization, model sizing, training infrastructure, and the cluster architecture.

I. INTRODUCTION

  • The study measures how AMD hardware, networking, and software behave under LLM pretraining workloads, including where time and bandwidth go and which configuration choices matter.It also addresses pitfalls encountered when porting mature NVIDIA pipelines to ROCm.Source: §I · pp. 1–2

A. Contributions

  • The production case study documents kernel selection, runtime configuration, ROCm/RCCL nuances, fused optimizer operations, and normalization kernels affecting throughput and stability.It also discusses Muon settings for stable large-batch training.Source: A. Contributions · p. 2
  • Pollara benchmarks sweep message sizes and GPU counts for major training collectives, relating results to optimizer, gradient, and CCA context-parallel communication.The characterized collectives include AllReduce, ReduceScatter, and AllGather.Source: A. Contributions · p. 2
  • MI300X-specific sizing guidance addresses attention and MLP blocks, while memory-bandwidth tests measure sustained performance for kernels relevant to LLMs.The sizing recommendations account for compute–memory balance and NIC characteristics and motivate small-MoE choices for training and inference efficiency.Source: A. Contributions · p. 2

II. CLUSTER SETUP

  • The documented cluster topology includes compute and storage-node internals and identifies bandwidth and contention points that affect scaling.Each node has eight MI300X GPUs and a dedicated Pollara NIC per GPU, arranged in a rails-only network.Source: A. Contributions · p. 2, Figure 2 · p. 3

A. Contributions

  • The training-stack account covers checkpointing, reshaping across parallelism regimes, fault tolerance, and a context-parallelism design co-tuned with CCA.The CCA parallelism recipe aims to keep activation memory and communication predictable as context length increases.Source: A. Contributions · p. 2
  • An iteration-time breakdown separates attention, MLP, and normalization compute from gradient and optimizer communication and IO.Source: A. Contributions · p. 2

A. Node Architecture

The section characterizes the AMD cluster and its hardware to guide system and model design, from node topology to GPU memory and compute behavior.

A. Node Architecture

B. Cluster Architecture

  • The rails-only cluster topology trades path diversity and routing flexibility for lower switch cost and simpler wiring, with a slight penalty unless communication avoids cross-rail traffic.The design uses four physical leaf switches and four spine switches per leaf.Source: B. Cluster Architecture · pp. 3–4

III. HARDWARE CHARACTERIZATION

  • Hardware characterization reduces deep-learning operations to their memory accesses and GEMMs, then assesses GPU compute, memory, and interconnects for training and inference.The analysis covers InfinityFabric within nodes and Pollara between nodes.Source: §III · p. 4

A. GPU HBM Bandwidth

  • MI300X HBM bandwidth can bottleneck intermediate operations and short-context attention; the compute-bound crossover for attention tends to occur around 4–16k tokens.The PyTorch benchmark mirrors memory-access patterns and tunes memcpy thread-block sizes to better reflect real workloads, though it does not capture every pattern.Source: A. GPU HBM Bandwidth · p. 4, Figure 3 · p. 5

B. GPU Compute

  • Approximately 200 GFLOPs are required for MI300X GEMMs to reach peak throughput.The authors searched matrix shapes and tuned rocBLAS and hipBLASLt algorithms to select efficient configurations.Source: B. GPU Compute · pp. 4–5, Figure 4 · p. 5

C. InfinityFabric Bandwidth

The section explains how intra-node InfinityFabric and inter-node Pollara shape communication and parallelism choices, alongside ZAYA1’s architectural innovations.

C. InfinityFabric Bandwidth

  • InfinityFabric xGMI requires all GPUs to participate in a collective for full bandwidth, so the authors favor parallelism degrees that span a full node or are not used.ZAYA1 uses ZeRO-1 initially and adds context parallelism at longer training contexts; long-context inference relies on tree attention to avoid xGMI bandwidth bottlenecks.Source: C. InfinityFabric Bandwidth · pp. 5–6, Figure 5 · p. 6

D. Pollara Bandwidth

  • Small interconnect messages are latency-bound, while larger messages are bandwidth-bound; ZAYA1’s gradient fusion buffer is set at the bandwidth saturation point to preserve overlap.Making the buffer larger would make overlap with fixed-size computation operations more difficult.Source: D. Pollara Bandwidth · pp. 6–7, Figure 7 · p. 7, Figure 6 · p. 7

A. Architecture

  • ZAYA1 combines Compressed Convolutional Attention, an expressive router, and residual scaling; CCA improves training speed versus GQA and MLA while reducing prefill FLOPs and maintaining equivalent KV-cache compression.The architecture figure identifies CCA in the attention block and a router with downprojection, EDA, and three sequential MLPs per expert.Source: A. Architecture · p. 8, Figure 8 · p. 8

1) Compressed Convolutional Attention (CCA):

The section describes CCA, the ZAYA1 router, and residual scaling as architectural changes to MoE models. It reports efficiency and model-performance benefits for these design choices.

1) Compressed Convolutional Attention (CCA):

  • CCA mixes sequences in a compressed latent space, reducing training and prefill compute and KV-cache size while matching MLA and GQA.The authors say ZAYA1-base provides some evidence that CCA can support complex reasoning, in-context learning, and long-range recall.Source: 1) Compressed Convolutional Attention… · p. 8

2) ZAYA1 Router:

  • The ZAYA1 router combines an MLP with Exponential Depth Averaging (EDA), which mixes router information with the previous layer’s routing choices.Its PID-inspired balancing scheme enforces balance across a microbatch; the authors report improved performance, balancing, and expert specialization with negligible iteration-time impact.Source: 2) ZAYA1 Router · pp. 8–9

3) ZAYA1 Residual Scaling:

  • Residual scaling applies learned biases and gating coefficients to residual streams and layer inputs, controlling unwanted inputs, forgetting, and residual-norm growth.The authors report benefits matching Qwen’s attention gating without its gating-matrix parameter and FLOP overhead; residual scaling adds only 2 × L × D parameters.Source: 3) ZAYA1 Residual Scaling · p. 9
  • Sixteen experts with a hidden-dimension expansion factor of 2 improved performance at fixed parameter sizing, while top-1 routing outperformed higher top-k in FLOP-matched experiments.The authors attribute avoiding residual experts to the router’s expressivity and expert specialization, and hypothesize that its lower-entropy routing makes additional experts less useful.Source: 3) ZAYA1 Residual Scaling · p. 9

B. Training

The section details ZAYA1-base’s phased training recipe and its use of MI300X memory capacity, alongside sizing guidance and an iteration-time breakdown. It links training and model choices to hardware and workload constraints.

B. Training

  • 8T tokens of web-heavy data formed phase 1, followed by 4T tokens with more code, math, reasoning, and SFT-formatted data in phase 2.Phase 3 extended context and increased STEM emphasis; a final annealing phase prepared the model for post-training, with a replay of phase 2 data before annealing.Source: B. Training · pp. 9–10, Figure 9 · p. 10
  • MI300X’s 192 GB HBM enabled data-parallel training with ZeRO-1, while Muon reduced memory needs relative to AdamW.For phase 1, minibatch size 5 fit in HBM; at 32k context the minibatch size was 1 with context-parallel size 2.Source: B. Training · p. 10
  • During the 1T-token context-extension phase, context grew from 4k to 32k; CCA reduced training FLOPs and activation memory by approximately 8× versus MLA.The authors report training at 32k with approximately the same efficiency as at 4k; at 128k, context-parallel size 8 did not require inter-node communication.Source: B. Training · pp. 9–10
  • Global batch size rose from 16M tokens during phases 1–2 to 30M during mid-training; the authors note MoE and Muon require larger batches for peak efficiency.With 16 experts and a 16M-token batch, each expert sees approximately 1M tokens, which the authors say may be below the critical batch for those parameters.Source: B. Training · p. 10

A. Model Sizing

  • The sizing guidance favors vocabulary sizes divisible by 64, the largest feasible microbatch, and minimal tensor and expert parallelism.The authors recommend parallelism only when HBM is insufficient and tune MoE GEMMs across expected input-length bands because their sizes depend on balancing.Source: A. Model Sizing · p. 11
  • At sequence length 4096, core forward and backward computation takes most iteration time; Muon compute and communication overhead motivated kernel and tensor-fusion tuning.The breakdown comes from GPU operations tagged in a PyTorch profile; optimizer communication and compute include distributed Muon communication and its Newton–Schulz iterations.Source: A. Model Sizing · pp. 11–12

B. Kernels

The section describes custom HIP kernels for LayerNorm and Muon operations, with a symmetric multiplication kernel to reduce Newton–Schulz work. The reported savings make Muon’s Newton–Schulz phase more bandwidth-friendly without altering its algorithm or stability.

B. Kernels

  • Custom HIP kernels were written for LayerNorm and several Muon optimizer subroutines.The Newton–Schulz step replaces generic Gram-matrix multiplications with a symmetric matrix multiplication kernel.Source: B. Kernels · p. 12, 1) Muon · p. 13

1) Muon:

  • Fused HIP multi-tensor kernels update momentum and weights in memory-linear passes, helping make iteration time compute- and communication-dominated rather than optimizer-step-dominated.The momentum pass emits BF16 Newton–Schulz inputs; the weight-update pass applies decoupled weight decay and updates FP32 master weights.Source: 1) Muon · p. 13
  • Approximately 2× less arithmetic forms A and A2, with approximately 2× fewer off-diagonal output-store writes; the authors report benign FP32 reduction-order differences only.The symmetric kernel computes upper-triangle tiles and stores their transposes, while materializing the full matrix before use.Source: 1) Muon · p. 13
  • With five Newton–Schulz steps, symmetric multiplication savings dominate optimizer overhead, shrinking its iteration-time share without altering the algorithm or its stability.The kernel avoids extra reads and writes, making the Newton–Schulz phase bandwidth-friendly.Source: 1) Muon · p. 13

2) LayerNorm:

The section presents a fused HIP LayerNorm kernel and describes its forward statistics and backward use. It also notes an RMSNorm variant using the same fusion and reduction structure.

2) LayerNorm:

  • The HIP kernel fuses residual addition, statistics, normalization, and affine transformation over (B·T) × E rows.It accumulates statistics in FP32 with single-pass Welford; I/O supports BF16 or FP32, and compute is FP32.Source: 2) LayerNorm · p. 14
  • Saved row means and inverse standard deviations support the backward pass, while parameter gradients use standard reductions over rows.The backward receives the upstream gradient and saved statistics.Source: 2) LayerNorm · p. 14
  • The RMSNorm variant drops mean subtraction while retaining the same fusion and reduction structure.Source: 2) LayerNorm · p. 14

C. Parallelism

Context-parallelism combines Ring Attention with CCA-specific sequence-mixing exchanges to scale longer contexts. For Muon, neighbor-to-neighbor SendRecv reconstructs only split parameters, avoiding a costly full-parameter AllGather.

1) Context-Parallelism:

  • Sequence shards are distributed across 2 × context-parallel world size GPUs, and Ring Attention overlaps local attention computation with key/value communication around a logical ring.Each GPU holds sequence length divided by 2 × context-parallel world size tokens; the shards are shuffled for load balancing.Source: 1) Context-Parallelism · p. 14
  • CCA’s narrow convolutions exchange only two boundary tokens per sequence chunk, while value shifting uses an all-gather in compressed space before re-sharding.The backward pass reverses the boundary send/receive operations and uses ReduceScatter instead of AllGather for the value projection.Source: 1) Context-Parallelism · p. 14, Figure 11 · p. 15

2) Muon:

  • For ZAYA1, AllGather in the Muon optimizer accounts for 99% of total optimizer memory, so the sharding scheme makes full-parameter gathering unnecessary.Flattened parameters are divided into contiguous chunks, leaving at most two parameters split across rank boundaries.Source: 2) Muon · p. 15
  • Muon exchanges only incomplete parameter portions with neighboring ranks via SendRecv, then concatenates them locally; rank-ordered send/receive avoids deadlock.Complete parameters are used directly, avoiding the full-parameter AllGather used by the baseline procedure.Source: 2) Muon · pp. 15–16

D. Checkpointing

Distributed checkpointing writes optimizer shards independently and overlaps asynchronous disk writes with training. Its reshape-friendly layout and frequent saves support faster recovery from failures and world-size changes.

D. Checkpointing

  • Checkpointing reduces save time by more than 10× versus naive checkpointing, allowing more frequent saves and reducing the time-cost of failures and restarts.Each rank writes its optimizer shard; the root rank writes model weights, and asynchronous disk writes overlap with post-checkpoint training.Source: D. Checkpointing · pp. 16–17
  • Muon and AdamW optimizer states are distributed across data-parallel ranks, while model weights are written once by the root rank.Muon is used for selected parameters and AdamW for the remainder; the optimizer states dominate checkpoint size.Source: D. Checkpointing · pp. 16–17
  • Optimizer shards can be reshaped offline for a different world size using deterministic unpad–remap–repad, enabling recovery after arbitrary node failures with minor global-batch adjustments.The framework pads shard sizes so the global parameter vector is divisible by alignment and GPU count.Source: D. Checkpointing · p. 17

E. Storage and I/O

Storage needs differ by model size: small-model training stresses IOPS, while large checkpoints stress bandwidth. Aegis detects failures and automates actions for known cases, while novel failures still require human intervention.

E. Storage and I/O

  • Small-model training stresses storage IOPS because fast iterations can outpace dataloading; sequential data shards, non-blocking loading, and sufficient storage bandwidth address the bottleneck.The authors report that simply adding CPU dataloader workers only delays GPU starvation.Source: E. Storage and I/O · p. 17, F. Fault Tolerance · p. 17
  • Large models stress storage bandwidth through massive checkpoints, so the authors use a 1 TB page cache to make subsequent checkpoint loads nearly instant from RAM.The cache can retain several recent checkpoints, which supports rewinding after loss spikes or divergence.Source: E. Storage and I/O · p. 17

F. Fault Tolerance

  • Aegis polls logs and observability systems for failures, then can notify operators or automatically act in clear-cut cases such as node failures and network issues.It connects to the IBM cloud platform to create, reboot, or delete nodes.Source: F. Fault Tolerance · p. 17, Figure 13 · p. 18
  • Aegis automates known failure cases, but newly detected failures still require human oversight and intervention.The authors describe seamless node or NIC replacement without restarting training as a longer-term goal.Source: F. Fault Tolerance · p. 18

VI. RESULTS

ZAYA1-base performs competitively against leading dense and MoE base models, with strong mathematics and coding comparisons despite its active parameter count. The authors also report architectural efficiency gains and conclude that AMD infrastructure supports large-scale pretraining.

VI. RESULTS

  • ZAYA1-base, with 760M active and 8.3B total parameters, competes with Qwen3-4B and larger prior-generation dense models, outperforms Llama3-8B across benchmarks, and exceeds Gemma3-12B in challenging mathematics and coding.It outperforms similar-scale MoE models on mathematics and general knowledge, while lagging slightly in coding; the table covers general knowledge, mathematics, and coding evaluations.Source: §VI · pp. 18–19, Table IV · p. 19
  • At high pass@k, ZAYA1-base performs strongly against full-fledged reasoning models and nearly matches Qwen3-4B-Thinking despite having no instruct or RL post-training.The table reports mathematics and general-knowledge evaluations using chain-of-thought reasoning at best-of-64 pass@k.Source: Table V · p. 20

VII. DISCUSSION

  • CCA, the ZAYA1 router, and residual scaling improve loss per FLOP and per parameter over contemporary MoE architectures; CCA also reduces prefill compute and memory costs while matching MLA decode speed.The router improves expert specialization, while residual scaling provides fine-grained information-flow control.Source: §VII · p. 20, §VIII · p. 21
  • The study focuses on AMD pretraining and a base model; the authors leave the model’s full potential after instruction-following and reasoning post-training for future work.The authors identify post-training, especially reasoning RL, as an avenue they plan to pursue.Source: §VI · p. 19, §VII · p. 20
  • The authors report that Muon added relatively small optimizer-step overhead with five Newton–Schulz iterations, while its optimal learning rate was sensitive to critical batch size.They also note that batch size becomes increasingly important as MoE sparsity increases.Source: §VII · p. 21

VIII. CONCLUSION

  • The paper concludes that MI300X GPUs and Pollara networking are mature enough for large-scale LLM pretraining, supported by detailed cluster benchmarks and training-stack engineering.The reported work includes custom kernels, fault tolerance, checkpoint reshaping, and context-parallelism for CCA.Source: §VII · p. 20, §VIII · p. 21

APPENDIX A CLUSTER DETAILS

The appendix documents the compute, storage, and login-node configurations, then estimates shared-storage I/O needs for training. Under the stated ZAYA1 workload, the 70,000-IOPS budget remains sufficient even with a scatter factor of 8.

A. Compute Nodes

  • The compute nodes each contain eight MI300X GPUs, 2 TB of DDR5 RAM, eight 400-Gbps Pollara NICs, and 25.6 TB of local NVMe storage.Each node also has two 56-core Intel Xeon Platinum 8570 sockets and a 200-GbE cloud NIC for data and checkpoint loading.Source: A. Compute Nodes · p. 26

1) Storage Node:

  • The dedicated storage node has 120 TB of NVMe in a single RAID0 array, 256 GB of DDR5 RAM, and a 100-GbE cloud NIC.Its two Xeon Gold 6426Y sockets each have 16 physical cores; local OS drives are separate.Source: 1) Storage Node · p. 26

2) Login Node:

  • The login node is a virtualized VM with 80 GB of memory, two eight-core Xeon sockets, and 1 TB split across three virtual disks.The disks are listed as 100 GB, 520 GB, and 520 GB.Source: 2) Login Node · p. 26

APPENDIX B STORAGE NODE SIZING AND I/O CALCULATIONS

  • The storage analysis models dataset reads for Megatron-style pretokenized corpora via mmap or buffered reads, excluding throughput-bound checkpoint writes.The scatter factor σ captures departures from contiguous access: σ ∈[1, 2] is typical for well-packed datasets, while fragmented or random-seek workloads can reach σ ∈[2, 8].Source: Appendix B · p. 26
  • At t=2.5 s, the ZAYA1 run remains above tbreak ≈ 0.234 · σ s even at σ=8, supporting the 70,000-IOPS storage budget.Each iteration reads 64 MiB across 16,384 pages and requires approximately 6,554 · σ IOPS.Source: Appendix B · p. 27

APPENDIX C COMPRESSED CONVOLUTIONAL ATTENTION

Compressed Convolutional Attention (CCA) performs attention in a compressed latent space, combining low-rank projections and convolutional mixing with normalization and skip connections. ZAYA1-base uses its GQA variant for 8× KV-cache compression; the appendix also presents intra- and inter-node All-to-All performance results.

APPENDIX C COMPRESSED CONVOLUTIONAL ATTENTION

  • CCA performs attention in a compressed latent space, with reported savings in attention memory and FLOPs and advantages over GQA and MLA in perplexity and training and inference FLOPs.The passage also identifies high KV-cache compression as crucial for rapid decoding.Source: Compressed Convolutional Attention · p. 27
  • CCA combines low-rank projections, sequence-mixing convolutions, a value-head time delay, and query/key mean skip connections with key-only head-wise temperature normalization.The authors identify these additions, especially convolution, as vital to the expressivity and nonlinearity that can let CCA match or exceed full attention with less compute and memory.Source: Compressed Convolutional Attention · p. 27
  • ZAYA1-base uses CCGQA with 2 KV heads for 8 query heads, achieving 8× KV-cache compression versus full multi-head attention.This combines GQA with 2× query compression; the paper describes the compression as perhaps the highest for a model of this scale.Source: Compressed Convolutional Attention · p. 27

APPENDIX D COMMUNICATION RESULTS

  • The appendix includes intra-node and multi-node RCCL All-to-All bandwidth and latency results; future expert-parallel training is intended to co-design with the intra-node fabric.The stated motivation is that later models will exceed HBM capacity for the expert block, requiring expert parallelism and token shuffling to and from experts each layer.Source: Appendix D · p. 27, Figure 15 · p. 28, Figure 16 · p. 28

Related papers

Loading 2511.17127v2…