AMD demonstrates end-to-end DeepSeek-V4 Flash RL training on 24 MI308X GPUs with AITER stack

AMD's ROCm team has published results from the first end-to-end reinforcement learning run of DeepSeek-V4 Flash on AMD Instinct MI308X GPUs, using an open-source AITER-backed software stack. The experiment spans 24 GPUs across three OAM nodes and covers the complete RL workflow: long-context FP8 rollout generation, BF16 mixture-of-experts actor training, online RDMA weight synchronization, evaluation, and checkpoint recovery. Training and inference operate as disaggregated services coordinated by Ray, with two nodes running a 16-worker Megatron-LM BF16 actor for numerical stability during gradient computation, while a third node runs an eight-way tensor-parallel vLLM FP8 rollout engine to maximize KV cache memory for longer sequences.

The reported run covers the first 75 optimizer steps of a GRPO training job with 16,384-token context length. Mean rollout throughput held stable at 311 generated tokens per second, while online reward accuracy improved from 45.3% to 59.2% across the run. Held-out AIME-2024 evaluation accuracy rose from 30.0% to 49.2% over the same period. The configuration uses TP4/PP4/EP4 parallelism for the actor and TP8 for rollout, with AITER providing ROCm-optimized kernels for both paths.

Architecture and disaggregated training design

• FP8/BF16 format split: The rollout engine uses FP8 weights to reduce per-parameter memory footprint, leaving more HBM for KV cache and enabling longer generated sequences; the actor retains BF16 for gradient-stable policy updates and optimizer steps.
• Hybrid attention schedule: DeepSeek-V4 Flash interleaves three layer types across 43 decoder layers—dense sliding-window attention (128 tokens) in layers 1-2, then alternating Compressed Sparse Attention (4:1 pooling, top-512 entries) and Heavily Compressed Attention (128:1 pooling) layers for long-context efficiency.
• Manifold-constrained hyper-connections: The model replaces standard residual connections with four-stream mHC mixing using pre-mixers and post-mixers, which both the vLLM rollout and Megatron-LM actor engines must preserve despite running different kernel sets.
• MoE routing complexity: Each token activates 6 of 256 routed experts per layer, with the first three layers using deterministic hash routing and later layers using learned routing; the actor topology uses TP4/PP4/EP4 parallelism to distribute this load.
• AITER kernel integration: The AMD-native training backend relies on AITER-backed ROCm kernels for attention and projection operations, including specialized handling of the 64-query-head, 512-dimensional shared KV latent and rank-1,024 low-rank query paths.

Why it matters

This validation gives AMD's ROCm stack a concrete, benchmarked reference for large-scale MoE RL training that previously lacked published end-to-end results on Instinct hardware. For organizations running or evaluating open-source AI infrastructure, the disaggregated FP8/BF16 design offers a transferable pattern for memory-constrained long-context training that can be adapted to other models and ROCm-supported GPU generations. The reliance on open-source components—vLLM, Megatron-LM, Ray, and AITER—rather than proprietary orchestration reduces deployment lock-in and provides independent developers with inspectable, modifiable paths for similar workloads.