SageMaker HyperPod上加速多模态RL训练

强化学习（RL）后训练正在成为构建强大语言模型智能体的标准步骤。模型通过生成轨迹、接收奖励并根据结果更新策略来学习跨多步骤的推理和行动。要在多个节点上大规模运行，每次训练需要数百 GPU 小时的 rollout（展开），就需要持久化的集群基础设施。该基础设施需要维持长时间运行的作业，在硬件故障时无需丢失进度即可恢复，并能在训练过程中提供对训练动态的可见性。

Amazon SageMaker HyperPod 在 Amazon Elastic Kubernetes Service（Amazon EKS）上为大规模机器学习（ML）工作负载提供这种基础设施。通过其集群弹性功能，HyperPod 持续监控节点健康状况并自动替换故障节点，因此硬件故障不会导致整个集群宕机。结合检查点（checkpointing）功能，训练作业可以从上次保存的步骤恢复，而无需从头开始。这对于长时间的多节点 RL 运行至关重要，否则单次硬件故障就会损失数小时的 rollout 进度。结合 HyperPod 上的 Ray 能力，你可以从 SageMaker Studio 创建 Ray 集群，使用安全连接远程提交作业，并通过 HyperPod Observability EKS 插件为你预置的 Amazon Managed Grafana 仪表板监控训练过程。

在本文中，我们展示如何使用这些功能来运行 SkyRL（一个开源 RL 框架），在 SageMaker HyperPod 上使用 Group Relative Policy Optimization（GRPO）训练 Qwen3-VL-8B 视觉语言模型，使其能够导航视觉迷宫。从 VisGym SFT 检查点（有监督微调起点）开始，HyperPod 上的 GRPO 后训练使迷宫解决率从 43.75% 提高到固定 64 个迷宫评估集上的 95% 以上。

要跟随本教程，你需要具备以下条件：

• 一个 SageMaker HyperPod 集群，采用 Amazon EKS 编排，至少有 3 个 ml.g7e.12xlarge 实例和 1 个 ml.r5d.16xlarge 实例。
• 集群中安装以下 Kubernetes 运算符：KubeRay 运算符、HyperPod Observability EKS 插件和 HyperPod Ray Endpoint 运算符（用于远程作业提交）。请参阅 Ray on HyperPod 入门指南。
• 集群上安装 Amazon FSx for Lustre CSI 驱动程序。你还需要一个 Amazon FSx for Lustre 文件系统、一个由该文件系统支持的 PersistentVolume，以及一个 pods 可以挂载的 PersistentVolumeClaim（ReadWriteMany）。训练作业将此存储用于 /shared 以进行检查点存储、低秩适应（LoRA）适配器同步和评估输出。
• 一个具有连接到 HyperPod 集群权限的 SageMaker Studio 域。请参阅设置 SageMaker Studio for Ray。
• 已安装 toolkit-for-ray-on-sagemaker-ai Python 包。

本节回顾了训练背后的强化学习概念以及本教程使用的集群拓扑。

多轮 RL 与 GRPO

标准单轮 RL 为单个模型输出分配奖励。而多轮 RL 则在整个步骤序列上训练一个智能体，它观察状态、执行动作、获得反馈，然后进入下一个状态。策略从整个回合中累积的奖励中学习，而不是从任何一个步骤中学习。

以导航 2D 迷宫为例。一个回合就是一次迷宫尝试，每回合就是一步：模型查看迷宫的当前图片，选择一个方向或决定停止，环境返回更新后的视图。奖励是稀疏的，所以只有当模型在实际达到目标（且在步数限制内）时才能获得 1.0，否则为 0。没有逐步的答案来训练，因为某一步是否好取决于它周围的步骤。

这正是 SkyRL 的 Group Relative Policy Optimization（GRPO）的用武之地。对于每个起始位置，智能体在当前策略下运行迷宫数次，GRPO 将这些运行相互比较，强化那些超过小组平均水平的运行，压制那些落后于平均水平的运行。这种组内比较就是整个训练信号，这使得 GRPO 可以在没有单独的评论家或价值模型的情况下工作。

这里讨论的解决方案在具有三个 GPU 工作节点和一个 CPU 头节点的 HyperPod Ray 集群上运行 SkyRL。SkyRL 将推理和训练共同放置在同一 GPU 上：vLLM 引擎生成 rollout（完整的迷宫回合），而使用完全分片数据并行（FSDP）分片的策略模型处理梯度更新。在每个优化器步骤之后，更新后的 LoRA 适配器权重通过 Amazon FSx for Lustre 共享存储从训练排名同步到推理引擎。

以下是我们使用的实例类型。其他 GPU 实例和集群大小也可以工作，只要工作节点有足够的 GPU 内存来运行模型。

• Workers（工作节点）：3x ml.g7e.12xlarge（每个 2 块 NVIDIA RTX PRO 6000 Blackwell GPU，共 6 块 GPU）
• Head（头节点）：ml.r5d.16xlarge（512 GB RAM，管理 Ray GCS、仪表板和 LoRA 适配器整合）
• Policy model（策略模型）：Qwen3-VL-8B + LoRA（rank 32），使用 PyTorch FSDP 在 6 块 GPU 上分片
• Rollout engines（Rollout 引擎）：6 个共置的 vLLM 实例，每 GPU 一个
• Shared storage（共享存储）：Amazon FSx for Lustre，挂载到 /shared，用于 LoRA 同步和评估输出

HyperPod 提供集群基础设施：Ray 集群从 SageMaker Studio 创建，作业提交使用 sagemaker_ray:// 协议，训练指标通过 HyperPod Observability 插件自动流入预置的 Amazon Managed Grafana 仪表板。

图 1：Amazon SageMaker HyperPod 上的 RayCluster 拓扑，包含一个 CPU 头节点和三个 GPU 工作节点，这些工作节点通过共享 Amazon FSx for Lustre 文件系统将 FSDP 策略分片和 vLLM rollout 引擎共置

以下步骤将引导你完成准备训练环境、启动集群、运行作业、监控进度以及托管训练好的模型。

步骤 1：准备容器镜像

要快速开始，请使用以下 Dockerfile 构建一个预装了 SkyRL、VisGym 及其依赖项的容器镜像。这是你在下一步在 HyperPod 上启动 Ray 集群时指定的镜像。它基于官方 NovaSky-AI SkyRL 基础镜像，并将 SkyRL 和 VisGym 固定到特定的提交 SHA 以确保构建可重现：

FROM novaskyai/skyrl-train-ray-2.57.0-py3.12-cu13.0

ENV HF_HUB_ENABLE_HF_TRANSFER=1 \
    QWEN_VL_MODEL=Qwen/Qwen3-VL-8B-Instruct \
    UV_PROJECT_ENVIRONMENT=/home/ray/anaconda3

# SkyRL's full FSDP stack into the system python Ray uses (uv .venv is invisible to `ray start`).
# Pinned to a commit SHA so the build is reproducible.
ARG SKYRL_REF=4298730b55bb01fe1b711662df53dca42a3b7615
RUN git clone https://github.com/NovaSky-AI/SkyRL.git /home/ray/skyrl \
    && cd /home/ray/skyrl && git checkout ${SKYRL_REF} \
    && cd /home/ray/skyrl/skyrl-train \
    && uv sync --active --extra fsdp \
    && /home/ray/anaconda3/bin/python -c \
       "import ray, torch, vllm, transformers, flash_attn; \
        from vllm_router.launch_router import launch_router; \
        print('skyrl stack OK', torch.__version__, vllm.__version__)"

# uv sync prunes Ray's dashboard extras; restore ray[default] so the full dashboard starts.
RUN uv pip install --python /home/ray/anaconda3/bin/python "ray[default]==2.57.0" \
    && /home/ray/anaconda3/bin/python -c \
       "from ray.dashboard.optional_deps import aiohttp"

# VisGym maze environment and dataset generator.
ARG VISGYM_REF=184fbd5e5dc81e32c8b944d9e40ac54dad62e3f2
RUN git clone https://github.com/anyscale/VisGym.git /home/ray/visgym \
    && cd /home/ray/visgym && git checkout ${VISGYM_REF} \
    && uv pip install --python /home/ray/anaconda3/bin/python -e . "pygame==2.6.1"

ENV FI_PROVIDER=efa FI_EFA_USE_DEVICE_RDMA=1 NCCL_PROTO=simple NCCL_DEBUG=INFO
WORKDIR /workspace
CMD ["/bin/bash"]

构建镜像并将其推送到你账户中的 Amazon Elastic Container Registry（Amazon ECR）仓库。记下完整的镜像 URI，因为你将在下一步创建 Ray 集群时使用它：

ACCOUNT=$(aws sts get-caller-identity --query Account --output text)
REGION=us-west-2
IMAGE_URI="${ACCOUNT}.dkr.ecr.${REGION}.amazonaws.com/ray/skyrl-visgym:latest"

aws ecr get-login-password --region ${REGION} | \
  docker login --username AWS --password-stdin ${ACCOUNT}.dkr.ecr.${REGION}.amazonaws.com

docker build -t skyrl-visgym .
docker tag skyrl-visgym ${IMAGE_URI}
docker push ${IMAGE_URI}

echo "Image URI: ${IMAGE_URI}"

步骤 2：从 SageMaker Studio 启动 Ray 集群

导航到 SageMaker Studio，选择 HyperPod，选择你的集群，然后转到 Tasks（任务）选项卡。从任务类型列表中选择 RayCluster，然后选择 Create Ray Cluster。

在创建表单中，给集群命名为 skyrl-visgym，将头节点实例类型设置为 ml.r5d.16xlarge，并添加三个使用 ml.g7e.12xlarge 的工作节点。将容器镜像设置为你在第一步中推送的 IMAGE_URI。

这里列出的实例类型是我们用于此演示的类型。其他实例类型也可以工作，但需要记住头节点的一个约束：它需要大容量内存。头节点在每次保存检查点时会将 LoRA 适配器分片从 GPU 工作节点汇总上来，这会短暂地将完整的适配器权重集加载到 CPU 内存中。我们使用 ml.r5d.16xlarge 是因为其大容量内存（512 GB RAM）可以满足这一需求。

要挂载你的 Amazon FSx 文件系统，请点击创建表单右上角的 YAML 按钮切换到原始清单编辑器，然后将卷和挂载点添加到头节点和工作节点 pod 规范中。每个 pod 的相关部分如下：

containers:
  - name: ray-head  # or ray-worker
    volumeMounts:
      - name: shared
        mountPath: /shared
volumes:
  - name: shared
    persistentVolumeClaim:
      claimName: <your-fsx-pvc-name>

将 <your-fsx-pvc-name> 替换为你的 Amazon FSx 文件系统支持的 PersistentVolumeClaim 的名称。有了这个配置后，/shared 在集群中的每个节点上都可用，训练任务可以从 pod 读取和写入检查点、LoRA 权重和评估输出。

开启 Remote endpoints，这样你就可以在不需要本地 kubectl port-forward 的情况下提交任务和打开仪表板。集群会为两者生成 IAM 认证的 URL。

图 2：SageMaker Studio 中的 Create Ray Cluster 表单，已开启 remote endpoints 用于任务提交和仪表板访问

一旦集群达到 Running 状态，Tasks 选项卡中的 Actions 菜单会提供 Open Ray Dashboard、Open Grafana 和集群管理选项。

图 3：SageMaker Studio Tasks 选项卡中处于 Running 状态的 skyrl-visgym Ray 集群，Actions 菜单已展开

第三步：准备训练脚本

将以下内容保存为你工作目录中的 train_job.sh。该脚本在首次运行时下载 SFT 检查点并生成数据集（两者都写入 Amazon FSx，因此会在多次运行间持久化），然后启动 GRPO 训练任务。

SFT 检查点为 GRPO 提供了一个良好的起点：在 VisGym 演示数据上预训练的 Qwen3-VL-8B 已经知道如何解析迷宫图像并发出结构化的移动动作，因此 GRPO 只需要优化哪些序列能够到达目标。

通过 trainer.placement.colocate_all=true，vLLM rollout 引擎和 FSDP 策略工作进程共享相同的 GPU。在 rollout 期间，GPU 并行运行推理。在策略更新步骤期间，它们集体运行 FSDP 训练。lora_sync_path 指向 Amazon FSx，以便在每个优化器步骤后将更新的适配器权重立即对节点上的推理引擎可见。如果没有共置，你需要为训练和推理分别准备 GPU 池，且它们在各阶段之间处于空闲等待状态，这种 ping-pong 模式会浪费计算资源。共置通过让训练和推理在同一硬件上交替执行来避免这种空闲时间。这就是为什么 gpu_memory_utilization=0.45 设置得比较保守：每个 GPU 需要同时为 FSDP 分片和 vLLM KV 缓存预留空间。

还有几个参数值得注意。n_samples_per_prompt=8 控制 GRPO 为每个迷宫 prompt 生成多少条 rollout 轨迹来计算组优势。max_turns=15 将每个 episode 限制在 15 步。hf_save_interval=20 每 20 步将 LoRA 适配器分片汇总到头节点并保存一个 Hugging Face 兼容的检查点。eval_interval=10 每 10 步运行一次保留的 64 个迷宫评估，以便你跟踪训练过程中的解决率。

任务还会写入完整的训练检查点，以便在中断后恢复。设置 ckpt_interval=20 每 20 步将完整训练状态保存到 ckpt_path。该状态包括模型权重、优化器状态、学习率调度器和数据加载器位置。resume_mode=latest 告诉 SkyRL 在任务启动时从该路径下最近的检查点继续。这与 HyperPod 集群弹性相结合。当节点发生故障时，HyperPod 自动检测并替换它，当你重新提交任务时，它会从最后保存的步骤继续，而不是从头开始。这里写入的完整检查点捕获训练状态用于恢复，而 hf_save_interval 导出的则是推理就绪的 LoRA 适配器，两者并行运行以满足不同目的。

#!/usr/bin/env bash
# GRPO fine-tune from the VisGym SFT checkpoint on maze_2d/easy.
set -euxo pipefail
cd /home/ray/skyrl
export HF_HUB_ENABLE_HF_TRANSFER=0
export SKYRL_RAY_PG_TIMEOUT_IN_S=600
RUN_ID=$(date +%Y%m%d-%H%M%S)
EXPORT_PATH="/shared/runs/${RUN_ID}-train"

ENV_ID=maze_2d/easy
TRAIN_DIR=/tmp/visgym_train_256
EVAL_DIR=/tmp/visgym_eval_seeded
SFT_CKPT=/shared/models/visgym_sft_mixed_qwen3vl

if [ ! -f "$SFT_CKPT/config.json" ]; then
  echo "=== Downloading SFT checkpoint ==="
  mkdir -p "$SFT_CKPT"
  python -c "
from huggingface_hub import snapshot_download
snapshot_download(repo_id='VisGym/visgym_model', allow_patterns='mixed_qwen3vl/*',
    local_dir='$SFT_CKPT', local_dir_use_symlinks=False)
import shutil, os; src='$SFT_CKPT/mixed_qwen3vl'
[shutil.move(os.path.join(src,f),'$SFT_CKPT') for f in os.listdir(src)]; os.rmdir(src)
"
fi

python examples/train/visgym/dataset.py --env_id "$ENV_ID" --num_rows 256 --output_dir "$TRAIN_DIR"
python examples/train/visgym/dataset.py --env_id "$ENV_ID" --num_rows 64  --seed --output_dir "$EVAL_DIR"

python examples/train/visgym/entrypoint.py \
  --env_variant sft \
  data.train_data="['$TRAIN_DIR/train.parquet']" \
  data.val_data="['$EVAL_DIR/train.parquet']" \
  trainer.algorithm.advantage_estimator="grpo" \
  trainer.policy.model.path="$SFT_CKPT" \
  trainer.policy.model.lora.rank=32 \
  trainer.policy.model.lora.alpha=32 \
  trainer.policy.model.lora.lora_sync_path="/shared/lora" \
  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 \
  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 \
  environment.env_class=visgym \
  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.vision_language_generator=true \
  generator.batched=false \
  trainer.remove_microbatch_padding=false \
  trainer.algorithm.use_kl_loss=false \
  trainer.policy.optimizer_config.lr=3.0e-6 \
  trainer.eval_interval=10 \
  trainer.eval_before_train=true \
  trainer.ckpt_interval=20 \
  trainer.hf_save_interval=20 \
  trainer.logger="console" \
  trainer.project_name="vlm_maze_2d_easy" \
  trainer.run_name="sft_grpo_${RUN_ID}" \
  trainer.resume_mode=latest \
  trainer.log_path="/tmp/skyrl-logs" \
  trainer.dump_eval_results=true \
  trainer.export_path="$EXPORT_PATH" \
  trainer.ckpt_path="s3://<your-bucket>/skyrl-visgym/ckpts/sft-grpo"

第四步：远程提交训练任务

当 toolkit-for-ray-on-sagemaker-ai 安装后，Ray 的标准 Jobs CLI 会通过该包注册的 sagemaker_ray:// 地址方案通过集群的安全端点进行认证。该库使用你的 AWS 凭证向 Ray 端点进行认证，因此你无需自己操作。这样，你就可以从笔记本电脑、CI/CD 流水线或具有 AWS 凭证的环境提交和跟踪任务，无需 kubectl port-forward，也无需与集群建立直接网络路径。

首先，对 EKS 集群进行认证：

aws eks update-kubeconfig --name <eks-cluster-name> --region us-west-2

然后提交任务，传递当前目录作为工作目录，以便 train_job.sh 被上传到集群头节点：

# Address format: sagemaker_ray://<ray-cluster-name>/<namespace>
ray job submit \
  --address sagemaker_ray://skyrl-visgym/default \
  --submission-id sft-train \
  --working-dir . \
  -- bash train_job.sh

提交后，使用相同的地址跟踪进度：

Step 5: Monitor training progress

HyperPod provides two monitoring surfaces: the Ray Dashboard for job-level visibility, and Amazon Managed Grafana for infrastructure and training metrics. Both are accessible directly from the Tasks tab in SageMaker Studio.

From the Tasks tab in SageMaker Studio, choose Open Ray Dashboard. This generates a short-lived authenticated URL for you automatically.

You can also generate the URL from the HyperPod CLI:

hyp create ray-dashboard-connection \
  --cluster-name <ray-cluster-name> \
  --namespace <kubernetes-namespace>

The command returns a presigned URL. Open it in a browser to view the Ray dashboard. The session is valid for up to six hours. For more details, see Generating a dashboard connection URL in the HyperPod documentation.

The following screenshot shows the Jobs view with the running job, its current step, and per-worker resource utilization.

Figure 4: The Ray Dashboard Jobs view showing the running training job and per-worker GPU utilization

From the Tasks tab, choose Open Grafana. The HyperPod Observability EKS add-on provisions four pre-built Ray dashboards in Amazon Managed Grafana: Ray Core, Ray Data, Ray Train, and Ray Serve. All four appear under a Ray folder and support filtering by cluster name. Here is a section of the core dashboard showing CPU, GPU, and memory utilization while the training is in progress.

Figure 5: Amazon Managed Grafana Ray Core dashboard panels for CPU, GPU, and memory utilization during training

SkyRL runs an evaluation pass every eval_interval=10 steps against the fixed 64-maze held-out set and logs eval/all/pass_at_1 to the console. You can grep for it in the job logs:

ray job logs sft-train \
  --address sagemaker_ray://skyrl-visgym/default | grep "eval/all/pass_at_1"

In our experiment, the model reached 75% solve rate around step 100 and peaked at 96.875% (62/64 mazes) at step 160, compared to a baseline of 43.75% (28/64 mazes) before GRPO post-training. Your results will vary based on hyperparameters and the maze configuration. Once the solve rate reaches your target, the LoRA adapter at that step is ready for inference. Checkpoints are saved to /shared/runs/<run_id>/global_step_<N>/policy/adapter_model.safetensors every 20 steps.

Step 6: Host the trained model for inference

With training complete, the artifact you deploy is a LoRA adapter rather than a full model, and Ray Serve loads that adapter on demand at request time. The one requirement is where the adapter lives. Ray Serve's dynamic LoRA loader reads adapters from cloud storage such as Amazon S3, so begin by staging the adapter in an S3 prefix. Copy the adapter from the checkpoint step you selected during training into that prefix:

aws s3 cp --recursive \
  /shared/runs/<run_id>/global_step_N/policy/ \
  s3://<your-bucket>/lora-adapters/maze-grpo/

Each adapter occupies its own subdirectory beneath this prefix, and that subdirectory name (maze-grpo in this example) is the name you will use to request the adapter once the endpoint is live.

To host the adapter, run Ray Serve on a Ray cluster built from the AWS Deep Learning Container for Ray Serve LLM, which bundles Ray Serve, the ray[llm] stack, and vLLM. With it, you can stand up an OpenAI-compatible endpoint using the built-in ray.serve.llm:build_openai_app builder and no custom image. We recommend running inference on a RayService, which KubeRay reconciles into its own Ray cluster (see the KubeRay docs).

A RayService describes its Serve application through a serveConfigV2 spec, and within that spec the ray.serve.llm:build_openai_app builder accepts one llm_configs entry per base model, as shown in the following configuration:

serveConfigV2: |
  applications:
    - name: maze-vlm
      route_prefix: /
      import_path: ray.serve.llm:build_openai_app
      args:
        llm_configs:
          - model_loading_config:
              model_id: visgym-qwen3vl
              # base SFT model on FSx
              model_source: /shared/models/visgym_sft_mixed_qwen3vl
            lora_config:
              # adapter prefix in S3
              dynamic_lora_loading_path: s3://<your-bucket>/lora-adapters
              max_num_adapters_per_replica: 1
            engine_kwargs:
              enable_lora: true
              # match the training LoRA rank
              max_lora_rank: 32
              max_loras: 1
              max_model_len: 16000
            deployment_config:
              autoscaling_config: { min_replicas: 1, max_replicas: 1 }

In this configuration, model_source points at the base SFT model on the Amazon FSx mount, while dynamic_lora_loading_path points at the S3 prefix you populated in the previous step. Set max_lora_rank to the same LoRA rank you used during training. Finally, make sure the Ray Serve replicas can obtain AWS credentials to read the adapters from your S3 bucket. On Amazon EKS, the recommended mechanism for this is Amazon EKS Pod Identity.

With this configuration in place, a single deployment serves both the base model and every adapter stored under the S3 prefix, and you choose which one to run through the model field on each request:

• Set model to visgym-qwen3vl to run the base SFT model on its own.
• Set model to visgym-qwen3vl:maze-grpo to apply the GRPO adapter on top of that base model. The ID follows the <base-model-id>:<adapter-name> convention, where the adapter name is the subdirectory you created under dynamic_lora_loading_path.

The first time an adapter is requested, Ray Serve downloads it from Amazon S3 onto the replica and caches it, so every later request reuses the loaded weights without downloading again. Because the endpoint speaks the OpenAI API, you can call it from an OpenAI-compatible client. For the complete server configuration options, see the Ray Serve LLM guide.