Масштабирование обучения с подкреплением MoE на Amazon EKS с использованием EFA и DeepEP для увеличения пропускной способности на 40%
aws ·
В статье описана оптимизированная архитектура для масштабирования обучения с подкреплением (RL) на основе смеси экспертов (MoE) в среде Amazon EKS, направленная на решение проблем координации гетерогенных вычислений и обеспечения высокоскоростной передачи данных. Решение объединяет Amazon Elastic Kubernetes Service (Amazon EKS), адаптер эластичной сети (EFA) и DeepEP для балансировки между генерацией роллов и обучением политики. Ключевой инновацией является использование DeepEP поверх EFA, которое заменяет стандартные коллективные операции all-to-all специализированными ядрами, снижая накладные расходы на коммуникацию при экспертном параллелизме. Тестирование на 48 экземплярах P5en показало увеличение совокупной пропускной способности RL-роллов на 40% по сравнению с базовыми конфигурациями. Архитектура также использует спотовые экземпляры Amazon EC2 для генерации роллов, что снижает затраты без прерывания тесно связанных рабочих нагрузок по обучению политики. Для обеспечения независимого масштабирования применяются отдельные группы узлов: для задач с ускорением GPU, сред на базе CPU и буферов опыта с оптимизацией по объему памяти. Такой подход позволяет эффективно согласовывать работу всех уровней стека, балансируя ограничения вычислительных мощностей, памяти и пропускной способности. В результате получается гибкое и экономически эффективное решение для распределенных рабочих нагрузок RL крупного масштаба, включающих сотни или тысячи акселераторов.
aws