AWS ускорила обучение MoE-моделей с подкреплением на 40% благодаря DeepEP и EFA

AWS показала архитектуру на базе Amazon EKS, EFA и DeepEP, которая ускоряет крупномасштабное обучение с подкреплением моделей с архитектурой Mixture-of-Experts. По данным компании, сквозной рост пропускной способности rollout-генерации составил около 40%.

Инженеры AWS описали связку Amazon EKS, Elastic Fabric Adapter (EFA) и DeepEP для масштабирования RL-пост-тренинга MoE-моделей, включая пайплайны RLHF и GRPO. Ключевая проблема таких сценариев — балансировка гетерогенных нагрузок: распределённой генерации rollout-ов и плотно связанного обучения политики, а также высокая нагрузка на межузловую коммуникацию из-за Expert Parallelism с динамической маршрутизацией токенов.

В архитектуре EKS управляет отдельными группами узлов для rollout-генерации, обучения политики, reward-моделей и вспомогательных сервисов; GPU-узлы используют NVLink внутри ноды и EFA между нодами, а долговременное хранение чекпоинтов вынесено в Amazon S3. Для экономии rollout-узлы предлагается запускать на Amazon EC2 Spot Instances с возможностью прерывания без остановки всего задания.

Основной прирост производительности даёт DeepEP over EFA: стандартные all-to-all коллективы NCCL заменяются специализированными dispatch и combine-ядрами, которые для межузлового обмена используют libfabric поверх EFA с поддержкой NVIDIA GPUDirect RDMA. На конфигурации из 48 инстансов P5en (16 на обучение, 32 на инференс) с супер-разреженной MoE-моделью включение DeepEP over EFA повысило совокупную пропускную способность rollout-генерации на 40% по сравнению со стеком на базе Slime без EFA-ускоренного Expert Parallelism. Описанный стек рассчитан на инстансы P5/P6 (p5.48xlarge, p5e.48xlarge, p6-b200.48xlarge, p6-b300.48xlarge) и компоненты CUDA 13.0, PyTorch 2.12, NCCL 2.31, DeepEP 2.0, SGLang 0.5.17 и Miles 0.1.0.

AWS также отдельно подчёркивает, что конфигурация масштабируется примерно до тысячи ускорителей и сокращает время итерации политики.