As Large Language Models (LLMs) evolve, Reinforcement Learning (RL) is becoming the crucial technique for aligning powerful models with human preferences and complex task objectives.

However, enterprises that need to implement and scale RL for LLMs are facing infrastructure challenges. The primary hurdles include the memory contention from concurrently hosting multiple large models (such as the actor, critic, reward, and reference models), iterative switching between high latency inference generation, and high throughput training phases.

This blog details Google Cloud’s full-stack, integrated approach, from custom TPU hardware to the GKE orchestration layer — and shares how you can solve the hybrid, high-stakes demands of RL at scale.

A quick primer: Reinforcement Learning (RL) for LLMs

RL is a continuous feedback loop that combines elements of both training and inference. At a high level, the RL loop for LLMs functions as follows:

The LLM generates a response to a given prompt. 1.

A “reward model” (often trained on human preferences) assigns a quantitative score, or reward, to the output. 1.

An RL algorithm (e.g., DPO, GRPO) uses this reward signal to update the LLM’s parameters, adjusting its policy to generate higher-rewarding outputs in subsequent interactions.

This generation, evaluation, and optimization continually improves the LLM’s performance based on predefined objectives.

RL workloads are hybrid and cyclical. The main goal of RL is not to minimize error (training) or fast prediction (inference), but to maximize reward through iterative interaction. The primary constraint for the RL workload is not just the computational power, but also system-wide efficiency, specifically minimizing aggregate sampler latency and maximizing the speed of weight copying for efficient end-to-end step time.

Google Cloud’s full-stack approach to RL

Solving these system-wide challenges requires an integrated approach. You can’t just have fast hardware or a good orchestrator; you need every layer of the stack to work together. Here is how our full-stack approach is built to solve the specific demands of RL:

1. Flexible, high-performance compute (TPUs and GPUs): Instead of locking customers into one path, we provide two high-performance options. Our TPU stack is a vertically integrated, JAX-native solution where our custom hardware (excelling at matrix operations) is co-designed with our post-training libraries (MaxText and Tunix). In parallel, we fully support the NVIDIA GPU ecosystem, partnering with NVIDIA on optimized NeMo RL recipes so customers can leverage their existing expertise directly on GKE.

2. Holistic, full-stack optimization: We integrate optimization from the bare metal up. This includes our custom TPU accelerators, high-throughput storage (Managed Lustre, Google Cloud Storage), and — critically — the orchestration and scheduling that GKE provides. By optimizing the entire stack, we can attack the system-wide latencies that bottleneck hybrid RL workloads.

3. Leadership in open-source: RL infrastructure is complex and built on a wide range of tools. Our leadership starts with open-sourcing Kubernetes and extends to active partnerships with orchestrators like Ray. We contribute to key projects like vLLM, develop open-source solutions like llm-d for cost-effective serving, and open-source our own high-performance MaxText and Tunix libraries. This helps ensure you can integrate the best tools for the job, not just the ones from a single vendor.

4. Proven, mega-scale orchestration: Post-training RL can require compute resources that rival pre-training. This requires an orchestration layer that can manage massive, distributed jobs as a single unit. GKE AI mega-clusters support up to 65,000 nodes today, and we are heavily investing in multi-cluster solutions like MultiKueue to scale RL workloads beyond the limits of a single cluster.

Running RL workloads on GKE

Existing GKE infrastructure is well-suited for demanding RL workloads and provides several infrastructure-level efficiencies.

The image below outlines the architecture and key recommendations for implementing RL at scale.

Figure : GKE infrastructure for running RL

At the base, the infrastructure layer provides the foundational hardware, including supported compute types (CPUs, GPUs, and TPUs). You can use the Run:ai model streamer to accelerate the model streaming for all three compute types. High performance storage (Managed Lustre, Cloud Storage) can be used for storage needs for RL.

The middle layer is the managed K8s layer powered by GKE, which handles the resource orchestration, resource obtainability using Spot or Dynamic Workload Scheduler, autoscaling, placement, job queuing and job scheduling and more at mega scale.

Finally, the open frameworks layer runs on top of GKE, providing the application and execution environment. This includes the managed support for open-source tools such as KubeRay, Slurm and gVisor sandbox for secure isolated task execution.

Building RL workflow

Before creating an RL workload, you must first identify a clear use case. With that objective defined, you then architect the core components: selecting the algorithm (e.g, DPO, GRPO), the model server (like vLLM or SGLang), the target GPU/TPU hardware, and other critical configurations.

Next, you can provision a GKE cluster configured with Workload Identity, GCS Fuse, and DGCM metrics. For robust batch processing, install the Kueue and JobSet APIs. We recommend deploying Ray as the orchestrator on top of this GKE stack. From there, you can launch the Nemo RL container, configure it for your GRPO job, and begin monitoring its execution. For the detailed implementation steps and source code, please refer to this repository.

Getting started with RL

Run RL on GPUs: Try the RL recipe on TPUs using MaxText and Pathways for GRPO algorithm, or if you use GPUs, try the NemoRL recipes. 1.

Partner with the open-source ecosystem: Our leadership in AI is built on open standards like Kubernetes, llm-d, Ray, MaxText or Tunix. We invite you to partner with us to build the future of AI together. Come contribute to llm-d! Join the llm-d community, check out the repository on GitHub, and help us define the future of open-source LLM serving.

Posted in

• Compute
• AI & Machine Learning