Architecting for Hardware Diversity: Meta's End-to-End Co-Design for High-Performance AI Training on AMD GPUs
Meta 面向 AMD GPU 的端到端协同设计与高性能 AI 训练
🧭 异构 GPU 训练协同优化 ★ 深度解读
Meta · AMD
通过硬件-模型协同设计、Origami 分析性能模型与 agentic kernel 优化，使 MI350X 在 DLRM 训练吞吐达 B200 的 101%
📌 工作参考价值（4 点）
- 硬件-模型协同设计：若业务有 AMD MI350X 或其他大 HBM 加速器，不要直接移植 CUDA kernel / 模型配置；应利用更大 HBM 扩大 per-GPU batch size、减少 activation checkpointing，并尝试扩大 embedding 表、序列长度或 hidden dim，可同时获得吞吐和模型质量收益（论文报告 20% 提速、>0.1% NE 提升）。
- 用分析性能模型替代离线 autotuning：推荐模型 GEMM 形状多且不规则，Origami 用硬件参数化 tile 模型微秒级选择 kernel，精度达 93%；可将类似思路用于自有训练框架的 kernel 选择或 Triton autotune 配置，避免维护海量 lookup table。
- 把 AMD 多 chiplet 的坑写进 agent skills：atomics 在 MI350X 上高竞争下可慢 10×，应优先用 shared memory 树归约 / 寄存器 cache 替换全局 atomic；把这类优化编码为 agentic workflow 技能，可自动修复长尾 kernel（论文中 JFA 提速 2.3×、端到端 25%）。
- Triton kernel 的硬件适配经验：AMD 64 线程 wavefront 不同于 NVIDIA warp，建议更大的 num_warps、浅 pipeline（num_stages≈2）、用 PID swizzle 改善 L2 locality，并启用 matrix_instr_nonkdim=16 / fast_expf 等；这些可作为跨 GPU 平台的 kernel 调优 checklist。