Accelerate multimodal RL training with SkyRL on Amazon SageMaker HyperPod

來源原標題

Accelerate multimodal RL training with SkyRL on Amazon SageMaker HyperPod

作者

Nilesh PS

原始標籤

Advanced (300) · Amazon SageMaker HyperPod · Technical How-to

原始連結

https://aws.amazon.com/blogs/machine-learning/accelerate-multimodal-rl-training-with-skyrl-on-amazon-sagemaker-hyperpod/

摘要原文

Reinforcement learning (RL) post-training is becoming a standard step in building capable language model agents. Models learn to reason and act across sequences of steps by generating trajectories, receiving rewards, and updating their policy based on outcomes. Running this at scale, across multiple nodes with hundreds of GPU-hours of rollouts per training run, requires persistent cluster infrastructure. That infrastructure needs to sustain long jobs, recover from hardware failures without losing progress, and provide visibility into training dynamics as they unfold. Amazon SageMaker HyperPod provides this infrastructure for large-scale machine learning (ML) workloads on Amazon Elastic Kubernetes Service (Amazon EKS) . Through its cluster resiliency features , it continuously monitors node health and automatically replaces faulty nodes, so a hardware failure does not take the cluster down with it. Paired with checkpointing, a training job can pick up from its last saved step instead of restarting from scratch. This matters for long multi-node RL runs, where a single hardware failure would otherwise cost hours of rollout progress. Combined with the Ray capabilities on HyperPod , you can create Ray clusters from SageMaker Studio, submit jobs remotely using secure connections, and monitor training through pre-built Amazon Managed Grafana dashboards that the HyperPod Observability EKS add-on provisions for you. In this post, we show how to use these capabilities to run SkyRL , an open-source RL framework, to train a Qwen3-VL-8B vision-language model to navigate visual mazes using Group Relative Policy Optimization (GRPO) on SageMaker HyperPod. Starting from the VisGym SFT checkpoint , a supervised fine-tuning (SFT) starting point, GRPO post-training on HyperPod improves the maze solve rate from 43.75% to more than 95% on a fixed 64-maze evaluation set. To follow this walkthrough, you need: This section reviews the reinforcement learning concepts behind the training and the cluster topology the walkthrough uses. Standard single-turn RL assigns a reward to a single model output. Multi-turn RL instead trains an agent over a whole sequence of steps, where it observes a state, acts, gets feedback, and moves on to the next state. The policy learns from the reward accumulated over the entire episode rather than from any one step. Consider the example problem of navigating a 2D maze. One episode is a single run at a maze, and each turn is one move: the model looks at the current picture of the maze, chooses a direction or decides to stop, and the environment sends back the updated view. Rewards are sparse, so the model earns 1.0 only when it actually reaches the goal within the move limit and nothing otherwise. There is no move-by-move answer key to train against, since whether a move was good depends on the moves around it. This is where SkyRL’s Group Relative Policy Optimization (GRPO) comes in. For each starting position, the agent runs the maze several times under the current policy, and GRPO grades those runs against one another, reinforcing the ones that beat the group’s average and pushing down the ones that trail it. That within-group comparison is the whole training signal, which lets GRPO work without a separate critic or value model. The solution discussed here runs SkyRL on a HyperPod Ray cluster with three GPU worker nodes and a CPU head node. SkyRL colocates inference and training on the same GPUs: vLLM engines generate rollouts (complete maze episodes) while a policy model sharded with Fully Sharded Data Parallel (FSDP) handles gradient updates. After each optimizer step, updated LoRA adapter weights sync from the training ranks to the inference engines through Amazon FSx for Lustre shared storage. These are the instance types we used. Other GPU instances and cluster sizes work as well, provided the workers have enough GPU memory for the model. HyperPod provides the cluster infrastructure: the Ray cluster is created from SageMaker Studio, job submission uses the sagemaker_ray:// protocol, and training metrics flow automatically into pre-built Amazon Managed Grafana dashboards through the HyperPod Observability add-on. Figure 1: RayCluster topology on Amazon SageMaker HyperPod, with one CPU head node and three GPU worker nodes that colocate FSDP policy shards and vLLM rollout engines over a shared Amazon FSx for Lustre filesystem The following steps walk through preparing the training environment, launching the cluster, running the job, monitoring progress, and hosting the trained model. To get started quickly, use the following Dockerfile to build a container image with SkyRL, VisGym, and their dependencies pre-installed. This is the image you will specify when launching your Ray cluster on HyperPod in the next step. It builds on the official NovaSky-AI SkyRL base and pins both SkyRL and VisGym to specific commit SHAs so the build is reproducible: Build the image and push it to an Amazon Elastic Container Registry (Amazon ECR) repository in your account. Note the full image URI, as you will use it when creating the Ray cluster in the next step: Navigate to SageMaker Studio, choose HyperPod , select your cluster, then go to the Tasks tab. From the task type list, choose RayCluster , then choose Create Ray Cluster . In the creation form, give the cluster the name skyrl-visgym , set the head instance type to ml.r5d.16xlarge , and add three workers using ml.g7e.12xlarge . Set the container image to the IMAGE_URI you pushed in Step 1. The instance types listed here are what we used for this walkthrough. Other instance types will work, but keep one constraint in mind for the head node: it needs large memory. The head consolidates LoRA adapter shards from the GPU workers at each checkpoint save, which briefly loads the full adapter weight set into CPU memory. We used ml.r5d.16xlarge for its large memory capacity (512 GB RAM) to accommodate this. To attach your Amazon FSx filesystem, choose the YAML button in the top-right corner of the creation form to switch to the raw manifest editor, then add the volume and mount to both the head and worker pod specs. The relevant section for each pod looks like this: Replace <your-fsx-pvc-name> with the name of the PersistentVolumeClaim backed by your Amazon FSx filesystem. With this in place, /shared is available on every node in the cluster and the training job can read and write checkpoints, LoRA weights, and evaluation output from the pods. Turn on Remote endpoints so you can submit jobs and open dashboards without a local kubectl port-forward . The cluster generates IAM-authenticated URLs for both. Figure 2: The Create Ray Cluster form in SageMaker Studio with remote endpoints turned on for job submission and dashboard access Once the cluster reaches Running status, the Actions menu in the Tasks tab offers Open Ray Dashboard, Open Grafana, and cluster management options. Figure 3: The skyrl-visgym Ray cluster at Running status in the SageMaker Studio Tasks tab, with the Actions menu open Save the following as train_job.sh in your working directory. The script downloads the SFT checkpoint and generates datasets on first run (both go to Amazon FSx, so they persist across runs), then launches the GRPO training job. The SFT checkpoint gives GRPO a strong starting point: Qwen3-VL-8B pre-trained on VisGym demonstrations already knows how to parse a maze image and emit structured move actions, so GRPO only needs to refine which sequences reach the goal. With trainer.placement.colocate_all=true , the vLLM rollout engines and FSDP policy workers share the same GPUs. During rollout, GPUs run inference in parallel. During the policy update step, they run FSDP training collectively. The lora_sync_path points to Amazon FSx so updated adapter weights are immediately visible to the inference engines on the nodes after each optimizer step.