Scalare l'apprendimento per rinforzo MoE su Amazon EKS con EFA e DeepEP ottenendo il 40% di throughput in più
aws ·
Questo articolo descrive un'architettura ottimizzata per scalare il reinforcement learning con Mixture-of-Experts (MoE) su Amazon EKS, affrontando le sfide del coordinamento tra risorse di calcolo eterogenee e delle comunicazioni ad alta throughput. La soluzione combina Amazon Elastic Kubernetes Service (Amazon EKS), Elastic Fabric Adapter (EFA) e DeepEP per bilanciare la generazione dei rollout con l'addestramento della policy. Un'innovazione chiave consiste nell'utilizzo di DeepEP su EFA, che sostituisce i collettivi all-to-all standard con kernel specializzati per ridurre l'overhead di comunicazione nel parallelismo degli esperti. I test eseguiti su 48 istanze P5en hanno dimostrato un aumento del 40% nel throughput aggregato dei rollout RL rispetto alle configurazioni di base. L'architettura sfrutta anche le istanze Amazon EC2 Spot per la generazione dei rollout, riducendo i costi senza interrompere i carichi di lavoro di addestramento della policy strettamente accoppiati. Vengono utilizzati gruppi di nodi separati per le attività accelerate da GPU, gli ambienti basati su CPU e i buffer di esperienza ottimizzati per la memoria, garantendo una scalabilità indipendente. Questo approccio consente una co-progettazione efficiente a tutti i livelli dello stack, bilanciando vincoli di calcolo, memoria e larghezza di banda. Il risultato è un percorso flessibile ed economico per carichi di lavoro RL distribuiti su larga scala che coinvolgono centinaia o migliaia di acceleratori.
aws