Posts by Zhen Huang
Using ODC to Accelerate AMD SFT Training
- 12 August 2026
Large-scale training spends a surprising share of its wall-clock time waiting instead of computing. Under Fully Sharded Data Parallel (FSDP), every layer ends in a collective all-gather or reduce-scatter, and every collective is a barrier that the whole data-parallel group has to reach together. Feed that machinery variable-length supervised fine-tuning (SFT) data and the picture gets worse: some ranks draw long documents while others draw short ones, so the fast ranks sit idle waiting for the slow ones.
Dropless MoE Training in JAX with Primus-Turbo
- 10 June 2026
Mixture-of-Experts (MoE) models have become a standard way to scale a transformer’s parameter count without paying the full compute bill — but training them efficiently on GPUs forces an uncomfortable trade-off. The default path in JAX/MaxText keeps every expert’s tensors at a fixed shape and simply drops the tokens that overflow each expert’s capacity, trading model quality for speed. The fully dropless alternative keeps every token, but in pure JAX it hits a memory wall that makes it impractical at production scale.
MoE Training Best Practices on AMD GPUs
- 16 December 2025
This blog covers best practices for training Mixture-of-Experts (MoE) models on AMD Instinct™ MI300/MI355-series[a] GPUs with the ROCm ecosystem. Whether you’re new to MoE distributed architectures or optimizing trillion-parameter models, this guide will help you identify bottlenecks and maximize efficiency on AMD hardware.