# ZAYA 1 MoE: Training Foundation Models on a Full-Stack AMD Platform

Source: https://www.youtube.com/watch?v=t-AIepeWxGI
Recap page: https://rapidrecap.app/video/t-AIepeWxGI
Generated: 2025-11-27T01:04:49.032+00:00

---
## Quick Overview

AMD achieved significant training efficiency gains for foundation models on their full-stack platform by employing a novel, custom-designed kernel architecture called ZAYA 1 MoE, which optimizes data movement and reduces redundant computations compared to standard methods.

**Key Points:**
- The ZAYA 1 MoE architecture achieved a 40% increase in iteration time compared to the standard Mixtral 8x7B model when training on the AMD platform.
- The model achieved 3.33 teraflops per second (TFLOPS/s) aggregate bandwidth, significantly higher than the 8x7B model's performance.
- The authors used custom-fused HIP kernels and implemented a strategy to reduce latent space memory usage by having each GPU node act as a dedicated page cache for the most recent checkpoint.
- The ZAYA 1 MoE design proved superior to the standard linear gate approach, effectively eliminating the bottleneck caused by excessive data movement in the standard setup.
- The paper highlights the success of the custom architecture in achieving peak throughput (around 300 GFLOPS) for a roughly 4B parameter dense model equivalent, validating the design choices.
- The authors specifically replaced the standard linear gate with a compact MLP featuring MoE-style routing, which smartly distributes computation across the 8 expert paths.
- The final implementation successfully avoided common pitfalls like silent data corruption or relying on slower residual stream checkpoints.

![Screenshot at 00:14: The speakers introduce the paper which focuses on training foundation models on a full-stack AMD platform, emphasizing the need to prove the viability of their custom architecture.](https://ss.rapidrecap.app/screens/t-AIepeWxGI/00-00-14.png)

**Context:** The video details the technical implementation and results of training large language foundation models, specifically focusing on AMD's hardware and software stack, known as the full-stack approach. The core concept revolves around overcoming communication and memory bottlenecks inherent in Mixture-of-Experts (MoE) models like Mixtral, particularly when deploying them across multiple GPUs.

## Detailed Analysis

The presenters discuss a major engineering paper detailing the training of foundation models on a full-stack AMD platform, specifically utilizing the MI300X GPUs and the ROCm software stack. The key achievement is the introduction of a novel architecture, ZAYA 1 MoE, designed to address the major bottleneck in MoE models: data movement and memory usage. The authors implemented custom-fused HIP kernels and a strategy where each GPU node acts as a dedicated page cache for the most recent checkpoint, avoiding slow residual stream transfers. This custom approach, which replaced the standard linear gate with a compact MLP featuring MoE-style routing, resulted in significant performance gains. For example, the ZAYA 1 MoE model achieved an iteration time improvement of 40% over the standard Mixtral 8x7B baseline and sustained 3.33 TFLOPS/s aggregate bandwidth. The design also proved highly efficient for both large-scale clusters and smaller consumer GPUs, delivering peak performance (around 300 GFLOPS) for a 4B parameter dense model equivalent while requiring significantly less memory overhead than older methods. The engineers essentially implemented a proactive automation system to manage memory and communication efficiently, validating their co-design methodology.

### Paper Focus

- Training foundation models on a full-stack AMD platform (MI300X GPUs, ROCm)
- Proving viability of custom architecture
- Deep dive into engineering solutions for MoE scaling

### Architectural Innovations

- Replaced standard linear gate with compact MLP featuring MoE-style routing
- Implemented dedicated page cache per GPU node for recent checkpoints
- Used custom-fused HIP kernels for efficient data handling

### Performance Metrics

- Achieved 40% faster iteration time vs. Mixtral 8x7B baseline
- Sustained 3.33 TFLOPS/s aggregate bandwidth
- Model is 4x denser than a 4B parameter dense model equivalent

### Key Takeaways

- The co-design of hardware and software (especially the custom kernels) is crucial for efficiency
- The new architecture allows for better context handling (up to 32K) and lower memory usage
- The design proves effective across large clusters and consumer hardware

![Screenshot at 00:00: Introductory graphic displaying the podcast theme and a call to action to become a member.](https://ss.rapidrecap.app/screens/t-AIepeWxGI/00-00-00.png)
![Screenshot at 00:25: The speaker explicitly mentions the hardware and software stack under discussion: MI300X GPUs and the ROCm software stack.](https://ss.rapidrecap.app/screens/t-AIepeWxGI/00-00-25.png)
![Screenshot at 01:38: Visual representation of the tight coupling between the dedicated network interface card \(NIC\) and the GPU, highlighting the importance of dedicated communication paths.](https://ss.rapidrecap.app/screens/t-AIepeWxGI/00-01-38.png)
![Screenshot at 04:48: A diagram or slide illustrating the concept of 'rigid parallelism' enforced by the hardware structure, which the new design circumvents.](https://ss.rapidrecap.app/screens/t-AIepeWxGI/00-04-48.png)
![Screenshot at 08:35: The speaker contrasts the new Muon optimizer with the older AdamW, noting Muon's lighter memory footprint.](https://ss.rapidrecap.app/screens/t-AIepeWxGI/00-08-35.png)