在 Amazon EKS 上使用 EFA 和 DeepEP 扩展 MoE 强化学习,吞吐量提升 40%
aws ·
本文详细介绍了在 Amazon EKS 上扩展混合专家(MoE)强化学习的优化架构,旨在解决异构计算协调和高吞吐量通信方面的挑战。该解决方案结合了 Amazon Elastic Kubernetes Service(Amazon EKS)、弹性网卡适配器(EFA)和 DeepEP,以平衡 rollout 生成与策略训练。一项关键创新是在 EFA 之上使用 DeepEP,通过专用内核替换标准的 all-to-all 集合操作,从而降低专家并行化中的通信开销。在 48 个 P5en 实例上的测试表明,与基线配置相比,聚合 RL rollout 吞吐量提高了 40%。该架构还利用 Amazon EC2 Spot 实例进行 rollout 生成,以降低在不干扰紧密耦合的策略训练工作负载情况下的成本。系统为 GPU 加速任务、基于 CPU 的环境以及内存优化的经验缓冲区分配了独立的节点组,以确保独立扩展。这种方法实现了跨技术栈的高效协同设计,平衡了计算、内存和带宽约束。最终结果为涉及数百或数千个加速器的大规模分布式 RL 工作负载提供了一条灵活且具成本效益的路径。
aws