The presentation introduces AMD’s Primus reference architecture designed for training large-scale AI models on AMD GPUs. It begins by highlighting the rapid growth in model sizes, with parameters increasing 15-fold over the past three years, driven by the surge in open-source AI models since 2023. This growth presents significant challenges for AI infrastructure, especially as training now often requires clusters with over 100,000 GPUs. The speaker emphasizes that simply scaling hardware is insufficient; instead, a holistic approach involving hardware-software co-optimization and efficient coordination among compute, communication, and storage components is essential.
AMD’s solution focuses on four key pillars: user-friendliness, performance, scalability, and stability. The goal is to enable users to start training with a single command, supporting a wide range of models including language, multi-modality, diffusion, and recommendation systems. Performance optimizations include support for low-precision training formats like FP8 and FP4, and scalable cluster management that can handle the complexity of hyperscale AI factories. Stability is critical given the long training durations, often spanning months, requiring fault tolerance and efficient resource utilization to minimize downtime.
The Primus software stack offers an end-to-end training lifecycle solution, from system design and trial runs to performance tuning and production-scale training. It includes tools for scaling projections, benchmarking, and tuning, all built on AMD’s ROCm platform. The architecture supports popular AI frameworks and provides over 80 YAML configuration files for quick deployment of open-source models. Continuous improvements are delivered through bi-weekly Docker releases and non-Docker installation options, ensuring users have access to the latest performance enhancements and model support.
Significant performance gains have been achieved through advanced sharding strategies, kernel optimizations, and memory management techniques. Innovations such as sync-free Mixture of Experts (MOE) models, the YokoEP communication protocol for better NIC compatibility, and mega-kernels that pipeline computation and communication have led to up to 2.4x performance improvements since ROCm 7.0. Memory visualization tools help users optimize pipeline parallelism and reduce bottlenecks, enabling strong scaling up to 64,000 GPUs with minimal performance loss.
Finally, the presentation addresses the importance of system stability and reliability for long-running training jobs. AMD is developing an agent-based telemetry and observability framework to monitor cluster health and predict potential failures. Additionally, dynamic fault tolerance mechanisms allow the training cluster to elastically shrink or expand without stopping the training process, minimizing wasted time due to hardware or network issues. The speaker concludes by reaffirming AMD’s commitment to evolving the Primus architecture to meet the demands of increasingly large AI models, highlighting recent hardware releases and ongoing software optimizations as key enablers for future advancements.
Useful Links
