Multimodal RL at Scale: SkyRL on Amazon SageMaker HyperPod with LoRA, FSDP and vLLM

Accelerate multimodal RL training with SkyRL on Amazon SageMaker HyperPod

One example run started with a vision and language agent solving 28 out of 64 mazes (43.75%). After a few hundred GRPO steps on a three-node HyperPod cluster, the same agent peaked at 62/64 solved, 96.875%. Those numbers come from a single experiment; they show what’s possible with the right topology and resilient infra, not a guaranteed outcome for every workload.

Why this pattern matters for business

Products that require planning and multi-turn reasoning (robotics orchestration, multimodal assistants, game agents) face two big frictions: fitting large policy models into memory for RL updates, and protecting long jobs from losing days of work on a node failure. This design combines three practical scaling levers:

Parameter-efficient tuning (LoRA) to produce small adapter artifacts instead of shipping full model weights for every variant.
Sharded training (PyTorch FSDP) so an 8B vision and language policy can be trained across GPUs without full replication.
High-throughput rollout engines (vLLM) to generate multi-turn trajectories efficiently; in the example these engines are colocated with policy shards to reduce idle GPUs.

Together, these let you produce deployable LoRA adapters that are cheap to store and quick to swap at inference time, provided you invest in checkpointing, monitoring, and careful memory profiling.

How the pieces connect (one-line flow)

vLLM rollouts → group-relative advantage estimator (GRPO) computes per-episode advantages → LoRA adapter updates applied to frozen base model → FSDP coordinates sharded optimizer/checkpointing → adapters exported and staged to S3 → Ray Serve dynamically loads adapters for inference.

Pipeline at a glance

Core framework: SkyRL, repository: https://github.com/NovaSky-AI/SkyRL (see the repo/README for GRPO details and code).
Base model: Qwen3-VL-8B (example env var: QWEN_VL_MODEL=Qwen/Qwen3-VL-8B-Instruct).
Starting checkpoint: VisGym SFT checkpoint on Hugging Face (example: https://huggingface.co/VisGym/visgym_model).
Estimator: SkyRL’s Group Relative Policy Optimization (GRPO), a group-relative advantage estimator used in this run (critic-free in this implementation; see the SkyRL README for the algorithm sketch).
Adapter format: LoRA (Low-Rank Adaptation) producing HF-compatible safetensors you can stage to S3 for dynamic loading.
Sharded training: PyTorch FSDP to shard model state across GPUs.
Rollouts: vLLM engines used as high-throughput inference engines; example run colocated one engine per GPU to avoid idle resources.
Shared storage: Amazon FSx for Lustre mounted at /shared for checkpoints and adapter staging; adapters are synced to S3 for serving.
Orchestration: Amazon SageMaker HyperPod on EKS running Ray (KubeRay), with observability and a sagemaker_ray:// Ray Jobs submission path (provided by toolkit-for-ray-on-sagemaker-ai).

Reproducibility checklist (run this before you scale)

Pin the exact code commits and Docker image used (see “Technical pins” section below) and export the full training config file.
Fix RNG seeds and maze generation seeds. Publish or snapshot your 64-maze eval set so others can reproduce the evaluation.
Record region and instance availability; verify that ml.g7e.12xlarge and ml.r5d.16xlarge map to expected GPUs in your region before provisioning.
Run a single-node smoke test (FSDP + colocated vLLM) to capture peak memory and FSx/S3 IO; only then scale to multi-node.
Decide checkpoint policy and retention: save full checkpoints frequently enough to meet your RPO, a practical heuristic is ckpt_interval ≈ (expected time-to-failure) / 3; then measure rollback time from the saved checkpoint.
Verify model license and usage constraints on the Hugging Face model cards for Qwen3-VL-8B and VisGym before using them in commercial products.

Example cluster and resource mapping used in the walkthrough

Topological example used to post-train a LoRA-adapted Qwen3-VL-8B policy with GRPO:

Head node: ml.r5d.16xlarge (512 GB RAM), cluster coordination and memory-heavy orchestration.
Worker nodes: 3 × ml.g7e.12xlarge (example had 2 GPUs per instance, NVIDIA RTX PRO 6000 Blackwell in our region at test time; verify current instance specs in your region).
Total GPUs in the example: 6 (FSDP sharded policy across 6 GPUs).
Rollout engines: 6 colocated vLLM instances (one per GPU in this example).
Shared file system: Amazon FSx for Lustre mounted at /shared for fast checkpoint IO and adapter staging.

Note on instance claims: Day1HPC reported that EC2 G7e instances with NVIDIA RTX PRO 6000 Blackwell deliver improved inference performance, which supports choosing g7e family hardware for rollout-heavy workloads. Still, always confirm instance specs and pricing in your region before committing.

Technical pins and artifact locations (exact values used)

Base image: novaskyai/skyrl-train-ray-2.57.0-py3.12-cu13.0
Ray version pinned: "ray[default]==2.57.0"
SKYRL_REF: 4298730b55bb01fe1b711662df53dca42a3b7615
VISGYM_REF: 184fbd5e5dc81e32c8b944d9e40ac54dad62e3f2
Environment variable example: QWEN_VL_MODEL=Qwen/Qwen3-VL-8B-Instruct
VisGym SFT checkpoint: https://huggingface.co/VisGym/visgym_model

“skyrl stack OK”

This short validation string is printed during the Dockerfile build validation in the example pipeline, a quick sanity check that required components were installed into the image.

Key hyperparameters and operational knobs (exact values excerpted from train_job.sh)

trainer.epochs = 20
trainer.train_batch_size = 24
trainer.policy_mini_batch_size = 12
trainer.micro_forward_batch_size_per_gpu = 1
trainer.micro_train_batch_size_per_gpu = 1
trainer.update_epochs_per_batch = 1
trainer.max_prompt_length = 2048
generator.sampling_params.max_generate_length = 1024
generator.sampling_params.temperature = 0.7
generator.max_turns = 15
generator.max_input_length = 8192
generator.n_samples_per_prompt = 8
generator.inference_engine.num_engines = 6
generator.inference_engine.tensor_parallel_size = 1
generator.inference_engine.gpu_memory_utilization = 0.45
generator.inference_engine.engine_init_kwargs.max_model_len = 16000
trainer.algorithm.advantage_estimator = “grpo”
trainer.policy.model.lora.rank = 32
trainer.policy.model.lora.alpha = 32
trainer.placement.colocate_all = true
trainer.strategy = fsdp
trainer.placement.policy_num_nodes = 3
trainer.placement.policy_num_gpus_per_node = 2
trainer.placement.ref_num_nodes = 3
trainer.placement.ref_num_gpus_per_node = 2
trainer.ref.fsdp_config.cpu_offload = false
trainer.eval_interval = 10
trainer.eval_before_train = true
trainer.ckpt_interval = 20
trainer.hf_save_interval = 20
trainer.resume_mode = latest
trainer.ckpt_path example: "s3://<your-bucket>/skyrl-visgym/ckpts/sft-grpo"
trainer.export_path example: "/shared/runs/${RUN_ID}-train"
ENV_ID used in example: maze_2d/easy
Dataset generation: train num_rows = 256, eval num_rows = 64

We set gpu_memory_utilization=0.45 intentionally to leave headroom when alternating between vLLM inference and FSDP training phases. Tune this after a colocated smoke run to avoid OOMs.

Result snapshot and how to interpret it

Baseline (pre-GRPO) solve rate: 43.75% (28/64)
~75% solve rate by step ~100
Peak 96.875% solve rate at step 160 (62/64)

Important context: these numbers come from a single experimental trace on a fixed 64-maze evaluation set. They illustrate GRPO’s potential when combined with LoRA and FSDP, but variability across random seeds, maze distributions, and hyperparameter tweaks is expected. To build production confidence, run multiple seeds, publish mean ± standard deviation, and snapshot the exact eval set and config.

Checkpoints, adapter exports, and serving

Full checkpoints (model weights, optimizer state, LR schedule, dataloader position) saved every 20 steps to the configured trainer.ckpt_path (example: S3 prefix).
HF-compatible LoRA adapter saves at trainer.hf_save_interval=20 are written to /shared/runs/<run_id>/global_step_<N>/policy/adapter_model.safetensors.
Adapters should be synced to an S3 prefix such as s3://<your-bucket>/lora-adapters/maze-grpo/ for serving.

“=== Downloading SFT checkpoint ===”

This echo appears during job setup when the training job fetches the supervised starting point.

Serving pattern and cold-start considerations

Ray Serve (via a RayService / KubeRay pattern) can dynamically download and apply LoRA adapters from S3 at request time. Example configuration keys used in the walkthrough include:

model_source pointing at a shared SFT base model on FSx
dynamic_lora_loading_path pointing at an S3 prefix
max_lora_rank (32 in example) and max_loras limits
enable_lora = true and max_model_len = 16000

Request ID pattern is base-model-id:adapter-name (for example visgym-qwen3vl:maze-grpo). On first request for an adapter, the replica downloads and caches it locally, so first-request latency includes S3 download plus adapter apply. Mitigations: warm replicas, pre-warm adapters, or use background pre-fetching based on expected traffic.

Job submission, monitoring, and observability

Remote job submission via Ray Jobs CLI with address scheme sagemaker_ray://skyrl-visgym/default (registered by the toolkit-for-ray-on-sagemaker-ai package).
Common job control commands used in the example: aws eks update-kubeconfig ..., ray job submit --address sagemaker_ray://... -- bash train_job.sh, and ray job logs ... --follow.
Ray Dashboard is reachable from SageMaker Studio Tasks → Open Ray Dashboard (example presigned URL validity shown as up to six hours in the walkthrough; confirm current behavior in your environment).
HyperPod Observability EKS add-on provisions Amazon Managed Grafana with pre-built Ray dashboards (Ray Core, Ray Data, Ray Train, Ray Serve) in the walkthrough; confirm availability in your HyperPod release.

Essential metrics to monitor during long RL runs: rollout throughput (tokens/s), episodes per minute, episode return distribution, advantage distribution, training step latency, GPU utilization and memory headroom, FSx IOPS and throughput, checkpoint upload time, and adapter load latency on first request. Alert on failed checkpoint uploads and sudden drops in FSx/S3 throughput.

Required cluster operators and packages

KubeRay operator (for Ray/Kubernetes reconciliation)
HyperPod Observability EKS add-on (Managed Grafana dashboards)
HyperPod Ray Endpoint Operator (for remote job endpoints)
Amazon FSx for Lustre CSI driver (mount /shared)
toolkit-for-ray-on-sagemaker-ai Python package (registers sagemaker_ray:// for Ray Jobs CLI)
AWS Deep Learning Container for Ray Serve LLM is recommended for serving (bundles Ray Serve, ray[llm], and vLLM)

Practical cautions, trade-offs and operational guidance

Cost estimation: the walkthrough does not publish total GPU-hours or a dollar figure. Estimate cost as: total GPUs × wall-clock hours × hourly instance rate (plus FSx and S3 IO). Use SageMaker/EC2 pricing pages for hourly rates and budget for checkpoint and FSx throughput costs.
Single-run variance: treat the 64-maze numbers as illustrative. Run multiple random seeds, snapshot the eval set, and report mean ± std for production claims.
Checkpoint frequency vs IO: frequent large checkpoints reduce lost work but increase FSx/S3 IO and cost. A practical rule: set ckpt interval ≈ expected time-to-failure / 3, then test restore time to validate your RPO.
Colocation trade-offs: colocating vLLM and FSDP shards increases GPU utilization but complicates memory management and scheduling. Profile memory fragmentation and peak usage in both inference and training phases; keep conservative headroom (example used 0.45).
Adapter rank choice: LoRA rank 32 was used in the example. Higher ranks increase expressivity but also adapter size and runtime memory during fusion; test lower ranks if adapter storage or serving latency matters.
Security and compliance: use IAM roles for service accounts (IRSA) or least-privilege pod IAM policies for S3/FSx access. Sanitize logs and avoid persisting sensitive environment observations unless you have explicit consent and retention policies.

Next experiments for teams that want to productionize

Run N ≥ 5 seeds with the same hyperparameters and report mean ± std on the fixed eval set.
Compare colocated vs. separated pools (separate vLLM inference instances) for cost, throughput, and ease of ops.
Measure cold-start latency for dynamic adapters at different adapter sizes and S3 regions; add pre-warming if SLA requires low latency.
Test incremental checkpoints (delta/sparse checkpoints) to reduce IO and speed restore where supported.

Logs and small signals you’ll see

"skyrl stack OK", printed during Dockerfile validation in the image build pipeline.
"=== Downloading SFT checkpoint ===", echoed when the job fetches the supervised fine-tuning checkpoint at startup.

Key takeaways, questions a curious reader would ask

Can SkyRL + HyperPod realistically train a multimodal agent like Qwen3-VL-8B with RL?

Yes, with caveats: the walkthrough demonstrates a feasible pattern using LoRA for parameter-efficient updates, FSDP to shard an 8B policy across 6 GPUs, and colocated vLLM rollouts. The example run improved solve rate from 43.75% to 96.875% on a fixed 64-maze eval set, but this is a single run, production confidence requires multi-seed experiments, eval snapshotting, and cost analysis.
What infrastructure is required?

Example topology used a head node (ml.r5d.16xlarge) and three ml.g7e.12xlarge workers (6 RTX PRO 6000 Blackwell GPUs total in the test region), plus Amazon FSx for