Suvendu will present a novel service-oriented architecture for scaling Reinforcement Learning from Human Feedback (RLHF) across heterogeneous accelerators, specifically targeting the migration from GPUs to AWS Trainium. Our approach addresses key challenges in implementing complex RLHF pipelines that orchestrate multiple models (SFT, Actor, Critic, and Reward Model) while supporting models up to 1T parameters for dense architectures and 2T parameters for Mixture-of-Experts (MoE).

The proposed system — which we have successfully applied this approach for for a general-purpose conversational AI product — introduces three key innovations: (1) A microservices-based architecture that separates non-Actor components across nodes, enabling flexible scaling and efficient resource utilization across GPU and Trainium accelerators; (2) A novel pipelined generation method that overlaps Actor inference and training using a consumer-producer buffer, significantly reducing training latency; and (3) Support for advanced techniques including LoRA fine-tuning, Grouped Query Attention, and multiple reward models.

This implementation leverages Neuron Nemo Megatron for cross-platform compatibility while incorporating optimizations for both GPU and Trainium backends. The architecture enables efficient handling of long contexts (>200k tokens) and provides a pathway for future optimizations including slop buffer implementation and lightweight multi-tenant reward modeling.