在 AIStudio Ray 训练任务中将 RLinf GRPO 的生成、推理与训练分配到两个 Worker

本教程展示如何在 AIStudio Ray 训练任务中，把生成、奖励计算和对数概率推理放在 Ray Head，把两路 FSDP actor 训练放在另一个 Ray Worker。两个 Worker 共同完成两个 GRPO step，并将日志、指标和分布式 checkpoint 保存到共享存储。

实战目标

完成本文后，您可以从一次小规模任务中确认：

AIStudio 托管的 Ray 集群包含 2 个 Worker，每个 Worker 可以使用 2 张 A100；
rollout、reward、inference 和 actor 按四个资源 rank 分配到两个节点；
DeepSeek-R1-Distill-Qwen-1.5B 完成两个 GRPO step，TensorBoard 记录奖励、loss、梯度范数和学习率；
两个 FSDP actor rank 写出分布式 checkpoint，任务结束后仍可从共享存储读取。

这次运行用于检查多 Worker 调度，以及生成、推理和 actor 训练能否跨节点衔接，不用于评估模型收敛或训练吞吐。

场景信息

本文使用以下固定范围：

RLinf：commit 7d07a4212ee6858cc333e1d4fab7a37256d1f839，入口为 examples/reasoning/main_grpo.py。
起始配置：RLinf 端到端测试配置 tests/e2e_tests/reasoning/qwen2.5-1.5b-grpo-pipeline-fsdp-sgl.yaml。
上游镜像：rlinf/rlinf:agentic-rlinf0.4-torch2.11.0-sglang0.5.12.post1-vllm0.23.0-megatron0.17.0-te2.17，来自 Docker Hub。
AIStudio 任务镜像：导入自己租户和任务可用区后的同一 tag。
模型：ModelScope 的 deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B@6fc93244f442ee2b5ab5c8000687ef5f7ffe1d03。
数据：inclusionAI/AReaL-boba-Data@1799c00be3f1216ab55a5cae3562d654dbfd7d82 中的 AReaL-boba-106k.jsonl。
资源来源：Spot。
分布式框架：Ray。AIStudio 创建托管 Ray 集群，RLinf Driver 从 Head 连接该集群。
Worker 拓扑：2 个 Worker、每个 Worker 2 张 A100-SXM4-80GB，共 4 张 GPU。
网络：每个 Worker 少于 8 张 GPU，因此不启用 RDMA；任务通过 IP 网络运行，并显式设置 NCCL_IB_DISABLE=1。
共享存储：AICoder 和训练任务都把目标存储卷挂载到 /mnt/rlinf-reproduction，并把 WORK_ROOT 设置为该路径。代码、模型和数据作为固定输入，运行结果写入 runs/two-worker-grpo-pipeline/。每次运行建议至少预留 25 GiB。
运行范围：runner.max_steps=2，第二个 step 结束时保存 FSDP checkpoint。

用小规模运行检查多 Worker 训练链路

本文使用 RLinf 的端到端测试配置 tests/e2e_tests/reasoning/qwen2.5-1.5b-grpo-pipeline-fsdp-sgl.yaml，不是 examples/ 目录中的训练示例。RLinf 使用这份配置检查 SGLang rollout、奖励计算、log probability 和 FSDP actor 优化能否跑通一次小规模训练链路。它调用正式入口 examples/reasoning/main_grpo.py 并实际执行各训练阶段，但使用较小的 batch、较短的序列和一个 step，以缩短检查时间。

本文不改动这份 YAML，而是通过 Hydra 参数把它从单节点四 GPU 调整为两个 AIStudio Worker、每个 Worker 两张 GPU，并把运行范围增加到两个 step。这样既能检查 AIStudio 托管 Ray 的多 Worker 调度、跨节点阶段交接和共享存储输出，也能在第二个 step 后读取分布式 checkpoint。

因此，本教程验证的是一条小规模但真实的多 Worker GRPO 训练链路。完成教程不能证明模型收敛、训练吞吐，也不能视为复现 examples/reasoning/config/math/qwen2.5-1.5b-grpo-fsdp-pipeline.yaml 中的 8 节点、64 GPU 训练方案。

本场景在 AIStudio 中怎样运行

AIStudio 的分布式框架决定 Ray 集群由谁创建，以及启动命令如何运行：

单机：适合把 RLinf 组件放在一个 Worker 内的任务。AIStudio 不创建托管 Ray 集群，由 RLinf 在 Worker 中启动本地 Ray。即使单个 Worker 有足够的 GPU，这种方式也不会检查跨 Worker 调度和阶段交接。
Ray（本教程选择）：AIStudio 创建 Ray Head 和 Ray Worker，并在 Head 中执行启动命令。启动脚本连接已经就绪的集群，不执行 ray start 或 ray job submit。两个 Worker 必须使用相同镜像，并能从相同路径读取共享存储。

本教程需要把 resource rank 0–1 和 2–3 放在不同 Worker 上，因此选择 Ray。RLinf 按以下方式分配四个 resource rank：

language-text

AIStudio Ray 训练任务：2 个 Worker × 2 张 GPU

Worker 0 / Ray Head
├── GPU 0 / resource rank 0：Rollout + Reward
└── GPU 1 / resource rank 1：Inference

Worker 1 / Ray Worker
├── GPU 0 / resource rank 2：FSDP Actor rank 0
└── GPU 1 / resource rank 3：FSDP Actor rank 1

rollout 生成回答，reward 计算数学奖励，inference 重新计算 actor 的 log probability，最后由两个 FSDP actor rank 执行优化阶段。模型、数据和 RLinf checkout 位于同一个共享存储卷；AICoder 和两个训练 Worker 都通过 /mnt/rlinf-reproduction 访问该存储。Head 把 Hydra 文件和主日志写回本次运行目录，actor rank 把 checkpoint shard 写入同一目录。

这个拓扑用于检查 AIStudio 托管 Ray 和跨节点阶段交接。平台托管 Ray 的 Driver 写法和资源规则见提交 Ray Driver 入口命令。

开始前准备

在分配 GPU 前，使用挂载同一共享存储的 AICoder 或开发机准备镜像、固定 checkout、模型、数据和两个辅助文件。以下准备命令和训练任务统一使用 /mnt/rlinf-reproduction。

在目标可用区准备任务镜像

本场景使用 Docker Hub 镜像 rlinf/rlinf:agentic-rlinf0.4-torch2.11.0-sglang0.5.12.post1-vllm0.23.0-megatron0.17.0-te2.17。

重要

请使用自己的镜像地址

下面的地址属于私有租户，不是公共镜像地址。请先把上游镜像上传或导入自己所选可用区的镜像仓库，再把 te-b905754427352261 替换为自己的 ＜tenant-id＞。

language-text

cr.infini-ai.com/te-b905754427352261/rlinf:agentic-rlinf0.4-torch2.11.0-sglang0.5.12.post1-vllm0.23.0-megatron0.17.0-te2.17

在镜像中心确认该 tag 可用于任务。本文不展开镜像下载和上传步骤；需要准备镜像时，请查看准备并验证容器镜像。

在共享存储中固定 RLinf checkout

以下命令在挂载目标共享存储的 AICoder 或开发机中执行：

language-bash

set -e
set -o pipefail

export WORK_ROOT=/mnt/rlinf-reproduction
export RLINF_COMMIT=7d07a4212ee6858cc333e1d4fab7a37256d1f839
export RLINF_ROOT="$WORK_ROOT/code/checkouts/$RLINF_COMMIT"

test ! -e "$RLINF_ROOT"
mkdir -p "$WORK_ROOT/code/checkouts"
git clone https://github.com/RLinf/RLinf.git "$RLINF_ROOT"
git -C "$RLINF_ROOT" checkout --detach "$RLINF_COMMIT"

test "$(git -C "$RLINF_ROOT" rev-parse HEAD)" = "$RLINF_COMMIT"
test -z "$(git -C "$RLINF_ROOT" status --porcelain)"

目标目录已经存在时，不要重复执行 git clone。直接核对 commit 和干净工作树即可。这里使用独立的 git clone，不要依赖其他工作区的 .git 管理目录。训练任务只读取准备好的 checkout，不在 GPU Worker 中克隆或更新代码。

准备固定版本的模型和数据

本场景与单 GPU GRPO 教程使用同一模型和数据。先按该教程下载并生成 SHA256SUMS，最终路径应为：

language-text

/mnt/rlinf-reproduction/models/modelscope/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B/6fc93244f442ee2b5ab5c8000687ef5f7ffe1d03
/mnt/rlinf-reproduction/datasets/inclusionAI/AReaL-boba-Data/1799c00be3f1216ab55a5cae3562d654dbfd7d82/AReaL-boba-106k.jsonl

模型目录应包含非空的 model.safetensors、config.json、tokenizer_config.json 和 SHA256SUMS；数据目录应包含非空的 JSONL 文件和 SHA256SUMS。

准备运行前后校验脚本

创建 /mnt/rlinf-reproduction/tools/two-worker-grpo-pipeline/verify-grpo-pipeline-v1.py，写入以下内容：

显示代码隐藏代码430 行

python

#!/usr/bin/env python3
"""Validate the two-Worker GRPO tutorial run."""

from __future__ import annotations

import argparse
import json
import math
import os
import re
import sys
from pathlib import Path
from typing import Any


def required_env(name: str) -> str:
    value = os.environ.get(name)
    if not value:
        raise RuntimeError(f"Missing environment variable: {name}")
    return value


def require_file(path: Path) -> None:
    if not path.is_file() or path.stat().st_size <= 0:
        raise RuntimeError(f"Required file is missing or empty: {path}")


def ordered_ray_nodes(ray: Any) -> list[dict[str, Any]]:
    nodes = [node for node in ray.nodes() if node.get("Alive")]
    current_node_id = str(ray.get_runtime_context().get_node_id())
    head = [node for node in nodes if str(node.get("NodeID")) == current_node_id]
    if len(head) != 1:
        raise RuntimeError(
            f"Cannot identify one Ray Head node: current={current_node_id} nodes={nodes}"
        )
    workers = sorted(
        (node for node in nodes if node is not head[0]),
        key=lambda node: str(node.get("NodeManagerAddress", "")),
    )
    return head + workers


def run_preflight(args: argparse.Namespace) -> None:
    import ray
    from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy

    rlinf_root = Path(required_env("RLINF_ROOT"))
    model_root = Path(required_env("MODEL_ROOT"))
    data_file = Path(required_env("DATA_FILE"))
    expected_commit = required_env("RLINF_COMMIT")

    ray.init(address="auto")
    nodes = ordered_ray_nodes(ray)
    if len(nodes) != args.expected_nodes:
        raise RuntimeError(
            f"Expected {args.expected_nodes} Ray nodes, found {len(nodes)}"
        )
    resources = ray.cluster_resources()
    expected_gpus = args.expected_nodes * args.expected_gpus_per_node
    if int(resources.get("GPU", 0)) != expected_gpus:
        raise RuntimeError(f"Expected {expected_gpus} Ray GPUs, found {resources}")

    @ray.remote(num_cpus=1, num_gpus=1)
    def probe_node(
        expected_rank: int,
        expected_gpus_per_node: int,
        rlinf_path: str,
        expected_rlinf_commit: str,
        model_root_path: str,
        data_file_path: str,
    ) -> dict[str, Any]:
        import hashlib
        import os
        import shutil
        import socket
        import subprocess
        import sys
        from importlib import import_module, metadata
        from pathlib import Path

        import ray
        import torch

        def file_sha256(path: Path) -> str:
            digest = hashlib.sha256()
            with path.open("rb") as stream:
                for chunk in iter(lambda: stream.read(1024 * 1024), b""):
                    digest.update(chunk)
            return digest.hexdigest()

        rows = subprocess.check_output(
            [
                "nvidia-smi",
                "--query-gpu=name,driver_version",
                "--format=csv,noheader,nounits",
            ],
            text=True,
        ).strip().splitlines()
        parsed = [row.rsplit(",", 1) for row in rows if row.strip()]
        gpu_names = [name.strip() for name, _ in parsed]
        drivers = sorted({driver.strip() for _, driver in parsed})
        if len(gpu_names) != expected_gpus_per_node:
            raise RuntimeError(
                f"rank={expected_rank}: expected {expected_gpus_per_node} GPUs, "
                f"found {gpu_names}"
            )
        if not all("A100-SXM4-80GB" in name for name in gpu_names):
            raise RuntimeError(f"rank={expected_rank}: unexpected GPUs={gpu_names}")
        if len(drivers) != 1:
            raise RuntimeError(f"rank={expected_rank}: driver versions={drivers}")
        if not sys.executable.startswith("/opt/venv/reason/bin/"):
            raise RuntimeError(f"rank={expected_rank}: Python={sys.executable}")
        if os.environ.get("NCCL_IB_DISABLE") != "1":
            raise RuntimeError(
                f"rank={expected_rank}: NCCL_IB_DISABLE="
                f"{os.environ.get('NCCL_IB_DISABLE')}"
            )

        configured_rank = os.environ.get("RLINF_NODE_RANK")
        if configured_rank is not None:
            if not configured_rank.isdigit() or int(configured_rank) != expected_rank:
                raise RuntimeError(
                    f"rank={expected_rank}: RLINF_NODE_RANK={configured_rank}"
                )

        commit = subprocess.check_output(
            ["git", "-C", rlinf_path, "rev-parse", "HEAD"], text=True
        ).strip()
        if commit != expected_rlinf_commit:
            raise RuntimeError(f"rank={expected_rank}: RLinf commit={commit}")

        model_root_local = Path(model_root_path)
        data_file_local = Path(data_file_path)
        model_manifest = model_root_local / "SHA256SUMS"
        data_manifest = data_file_local.parent / "SHA256SUMS"
        for path in (
            model_root_local / "config.json",
            model_root_local / "tokenizer_config.json",
            model_root_local / "model.safetensors",
            model_manifest,
            data_file_local,
            data_manifest,
        ):
            if not path.is_file() or path.stat().st_size <= 0:
                raise RuntimeError(f"rank={expected_rank}: missing input={path}")
            with path.open("rb") as stream:
                stream.read(1)

        try:
            distribution = metadata.distribution("flash-attn-4")
        except metadata.PackageNotFoundError:
            disabled_candidates = list(
                Path(sys.prefix).glob("lib/python*/site-packages/*.dist-info.a100-disabled")
            )
            fa4_action = (
                "metadata-already-disabled"
                if any("flash_attn_4" in path.name for path in disabled_candidates)
                else "metadata-absent"
            )
        else:
            dist_info = Path(distribution._path).resolve()
            venv_root = Path(sys.prefix).resolve()
            if venv_root not in dist_info.parents:
                raise RuntimeError(
                    f"rank={expected_rank}: unexpected flash-attn-4 path={dist_info}"
                )
            disabled = Path(f"{dist_info}.a100-disabled")
            if disabled.exists():
                raise RuntimeError(
                    f"rank={expected_rank}: both active and disabled FA4 metadata exist"
                )
            shutil.move(str(dist_info), str(disabled))
            fa4_action = "metadata-disabled"

        try:
            visible_fa4 = metadata.version("flash-attn-4")
        except metadata.PackageNotFoundError:
            pass
        else:
            raise RuntimeError(
                f"rank={expected_rank}: flash-attn-4 remains visible={visible_fa4}"
            )

        if not torch.cuda.is_available():
            raise RuntimeError(f"rank={expected_rank}: CUDA is unavailable")
        capability = ".".join(map(str, torch.cuda.get_device_capability(0)))
        if capability != "8.0":
            raise RuntimeError(
                f"rank={expected_rank}: compute capability={capability}"
            )
        import_module("transformer_engine.pytorch")
        import_module("sglang")

        context = ray.get_runtime_context()
        return {
            "rank": expected_rank,
            "node_id": str(context.get_node_id()),
            "hostname": socket.gethostname(),
            "driver_version": drivers[0],
            "gpu_count": len(gpu_names),
            "gpu_name": gpu_names[0],
            "compute_capability": capability,
            "fa4_action": fa4_action,
            "rlinf_commit": commit,
            "model_manifest": file_sha256(model_manifest),
            "data_manifest": file_sha256(data_manifest),
        }

    refs = []
    for expected_rank, node in enumerate(nodes):
        refs.append(
            probe_node.options(
                scheduling_strategy=NodeAffinitySchedulingStrategy(
                    node_id=node["NodeID"], soft=False
                )
            ).remote(
                expected_rank,
                args.expected_gpus_per_node,
                str(rlinf_root),
                expected_commit,
                str(model_root),
                str(data_file),
            )
        )
    workers = sorted(ray.get(refs), key=lambda item: item["rank"])
    if len({item["node_id"] for item in workers}) != args.expected_nodes:
        raise RuntimeError(f"Probes did not run on distinct Ray nodes: {workers}")
    if len({item["driver_version"] for item in workers}) != 1:
        raise RuntimeError(f"NVIDIA driver versions differ: {workers}")
    if len({item["model_manifest"] for item in workers}) != 1:
        raise RuntimeError(f"Model manifests differ between Workers: {workers}")
    if len({item["data_manifest"] for item in workers}) != 1:
        raise RuntimeError(f"Dataset manifests differ between Workers: {workers}")

    report = {
        "cluster_resources": resources,
        "workers": workers,
    }
    args.report.parent.mkdir(parents=True, exist_ok=True)
    args.report.write_text(
        json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )

    print("RAY_CLUSTER_RESOURCES=" + json.dumps(resources, sort_keys=True))
    for item in workers:
        role = "Ray Head" if item["rank"] == 0 else "Ray Worker"
        print(
            "WORKER_TOPOLOGY="
            f"rank={item['rank']} role={role} host={item['hostname']} "
            f"gpu_count={item['gpu_count']} gpu={item['gpu_name']}"
        )
        print(
            f"NVIDIA_DRIVER_VERSION_NODE_{item['rank']}={item['driver_version']}"
        )
        print(
            "A100_RUNTIME="
            f"rank={item['rank']} "
            f"compute_capability={item['compute_capability']} "
            f"fa4_action={item['fa4_action']}"
        )
    print(f"RLINF_PREFLIGHT_REPORT={args.report}")
    print("RLINF_PREFLIGHT_STATUS=passed")


def run_postflight(args: argparse.Namespace) -> None:
    from omegaconf import OmegaConf
    from tensorboard.backend.event_processing.event_accumulator import EventAccumulator

    output_dir = args.output_dir
    driver_log = output_dir / "driver.log"
    main_log = output_dir / "log" / "main.log"
    preflight_path = output_dir / "preflight.json"
    hydra_config_path = output_dir / "hydra" / ".hydra" / "config.yaml"
    for path in (driver_log, main_log, preflight_path, hydra_config_path):
        require_file(path)

    driver_text = driver_log.read_text(encoding="utf-8", errors="replace")
    main_text = main_log.read_text(encoding="utf-8", errors="replace")
    combined_text = driver_text + "\n" + main_text
    fatal_patterns = {
        "python traceback": r"Traceback \(most recent call last\):",
        "Hydra execution error": r"Error executing job with overrides",
        "Ray task error": r"RayTaskError\(",
        "Ray actor failure": r"(?:ActorDiedError|ActorUnavailableError)",
        "CUDA out of memory": r"CUDA out of memory",
        "preflight failure": r"RLINF_PREFLIGHT_STATUS=failed",
    }
    detected = [
        label
        for label, pattern in fatal_patterns.items()
        if re.search(pattern, combined_text)
    ]
    if detected:
        raise RuntimeError(f"Fatal markers found in logs: {detected}")
    if "RLINF_PYTHON_EXIT_CODE=0" not in driver_text:
        raise RuntimeError("RLinf Python did not report exit code 0")
    completion_line = (
        f"Step limit given by max_steps={args.target_steps} reached. Stopping run"
    )
    if completion_line not in combined_text:
        raise RuntimeError(f"Missing RLinf step-limit marker: {completion_line}")
    preflight = json.loads(preflight_path.read_text(encoding="utf-8"))
    workers = preflight.get("workers", [])

    cfg = OmegaConf.load(hydra_config_path)
    if int(cfg.cluster.num_nodes) != args.expected_nodes:
        raise RuntimeError(f"Resolved cluster.num_nodes={cfg.cluster.num_nodes}")
    placement = OmegaConf.to_container(cfg.cluster.component_placement, resolve=True)
    expected_placement = {
        "rollout": 0,
        "inference": 1,
        "actor": "2-3",
        "reward": 0,
    }
    if placement != expected_placement:
        raise RuntimeError(f"Resolved component placement={placement}")
    if int(cfg.runner.max_steps) != args.target_steps:
        raise RuntimeError(f"Resolved runner.max_steps={cfg.runner.max_steps}")
    if str(cfg.runner.output_dir) != str(output_dir.parent):
        raise RuntimeError(f"Resolved runner.output_dir={cfg.runner.output_dir}")
    if str(cfg.runner.experiment_name) != output_dir.name:
        raise RuntimeError(
            f"Resolved runner.experiment_name={cfg.runner.experiment_name}"
        )

    tensorboard_dir = output_dir / "tensorboard"
    events = sorted(tensorboard_dir.rglob("events.out.tfevents.*"))
    if not events or not all(path.stat().st_size > 0 for path in events):
        raise RuntimeError(f"TensorBoard event file is missing: {tensorboard_dir}")
    required_tags = (
        "rollout/reward_scores",
        "actor/training/actor/final_loss",
        "actor/training/actor/grad_norm",
        "actor/training/actor/lr",
    )
    accumulator = None
    observed_tags = {}
    for event_dir in sorted({path.parent for path in events}):
        candidate = EventAccumulator(str(event_dir), size_guidance={"scalars": 0})
        candidate.Reload()
        scalar_tags = set(candidate.Tags().get("scalars", []))
        observed_tags[str(event_dir)] = sorted(scalar_tags)
        if set(required_tags).issubset(scalar_tags):
            accumulator = candidate
            break
    if accumulator is None:
        raise RuntimeError(
            f"No TensorBoard event directory contains the required scalar tags: "
            f"{observed_tags}"
        )
    metrics = {}
    for tag in required_tags:
        points = accumulator.Scalars(tag)
        if not points:
            raise RuntimeError(f"TensorBoard scalar has no values: {tag}")
        value = float(points[-1].value)
        if not math.isfinite(value):
            raise RuntimeError(f"TensorBoard scalar is not finite: {tag}={value}")
        metrics[tag] = {"step": int(points[-1].step), "value": value}
    if metrics["actor/training/actor/lr"]["value"] <= 0:
        raise RuntimeError(f"Actor learning rate is not positive: {metrics}")

    checkpoint_dir = output_dir / "checkpoints" / f"global_step_{args.target_steps}"
    actor_dcp = checkpoint_dir / "actor" / "dcp_checkpoint"
    data_state = checkpoint_dir / "data" / "data.pt"
    require_file(actor_dcp / ".metadata")
    require_file(data_state)
    shards = sorted(path for path in actor_dcp.glob("*.distcp") if path.stat().st_size > 0)
    if len(shards) < 2:
        raise RuntimeError(f"Expected at least two FSDP checkpoint shards: {shards}")

    report = {
        "checkpoint": {
            "directory": str(checkpoint_dir),
            "fsdp_shards": [
                {"name": path.name, "bytes": path.stat().st_size} for path in shards
            ],
            "data_state_bytes": data_state.stat().st_size,
        },
        "metrics": metrics,
        "workers": workers,
        "resolved_component_placement": placement,
        "target_steps": args.target_steps,
    }
    report_path = output_dir / "verification.json"
    report_path.write_text(
        json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    for tag, point in metrics.items():
        print(f"GRPO_METRIC={tag} step={point['step']} value={point['value']}")
    print(f"ACTOR_FSDP_SHARDS={len(shards)}")
    print(f"CHECKPOINT_DIR={checkpoint_dir}")
    print(f"RLINF_VERIFICATION_REPORT={report_path}")
    print("RLINF_VERIFICATION_STATUS=passed")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="command", required=True)

    preflight = subparsers.add_parser("preflight")
    preflight.add_argument("--expected-nodes", type=int, required=True)
    preflight.add_argument("--expected-gpus-per-node", type=int, required=True)
    preflight.add_argument("--report", type=Path, required=True)

    postflight = subparsers.add_parser("postflight")
    postflight.add_argument("--output-dir", type=Path, required=True)
    postflight.add_argument("--target-steps", type=int, required=True)
    postflight.add_argument("--expected-nodes", type=int, required=True)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.command == "preflight":
        run_preflight(args)
    else:
        run_postflight(args)


if __name__ == "__main__":
    try:
        main()
    except Exception as error:
        print(
            f"RLINF_{'PREFLIGHT' if 'preflight' in sys.argv else 'VERIFICATION'}_STATUS=failed "
            f"error={type(error).__name__}: {error}",
            file=sys.stderr,
        )
        raise

这个脚本的 preflight 子命令会在两个 Ray 节点分别检查 GPU、NVIDIA 驱动、固定输入和 Python 运行时，并为 A100 选择镜像中已有的 FlashAttention 2 路径。postflight 子命令会读取解析后的 Hydra 配置、TensorBoard 指标和 FSDP checkpoint；缺少任一业务信号时返回非 0。

准备 GRPO 训练启动脚本

创建 /mnt/rlinf-reproduction/tools/two-worker-grpo-pipeline/run-grpo-pipeline-v1.sh，写入以下内容：

显示代码隐藏代码132 行

bash

#!/usr/bin/env bash

set -e
set -o pipefail

: "${RUN_ID:?Set RUN_ID in the AIStudio task}"
: "${WORK_ROOT:?Set WORK_ROOT to the shared-storage path mounted in every Worker}"
TARGET_STEPS="${TARGET_STEPS:-2}"

[[ "$RUN_ID" =~ ^[A-Za-z0-9][A-Za-z0-9._-]*$ ]] || {
  echo "RUN_ID contains unsupported characters" >&2
  exit 1
}
[[ "$TARGET_STEPS" =~ ^[1-9][0-9]*$ ]] && (( TARGET_STEPS >= 2 )) || {
  echo "TARGET_STEPS must be an integer greater than or equal to 2" >&2
  exit 1
}

export RLINF_COMMIT=7d07a4212ee6858cc333e1d4fab7a37256d1f839
export MODEL_REVISION=6fc93244f442ee2b5ab5c8000687ef5f7ffe1d03
export DATA_REVISION=1799c00be3f1216ab55a5cae3562d654dbfd7d82
export RLINF_ROOT="$WORK_ROOT/code/checkouts/$RLINF_COMMIT"
export MODEL_ROOT="$WORK_ROOT/models/modelscope/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B/$MODEL_REVISION"
export DATA_FILE="$WORK_ROOT/datasets/inclusionAI/AReaL-boba-Data/$DATA_REVISION/AReaL-boba-106k.jsonl"
export VERIFY_SCRIPT="$WORK_ROOT/tools/two-worker-grpo-pipeline/verify-grpo-pipeline-v1.py"
export OUTPUT_ROOT="$WORK_ROOT/runs/two-worker-grpo-pipeline"
export OUTPUT_DIR="$OUTPUT_ROOT/$RUN_ID"
export REPO_PATH="$RLINF_ROOT"
export PYTHONPATH="$RLINF_ROOT:/opt/Megatron-LM:${PYTHONPATH:-}"
export CUDA_DEVICE_MAX_CONNECTIONS=1
export TOKENIZERS_PARALLELISM=false
export PYTHONUNBUFFERED=1
export HF_HUB_OFFLINE=1
export HF_DATASETS_OFFLINE=1
export TRANSFORMERS_OFFLINE=1
export NCCL_IB_DISABLE=1
export RLINF_CATCH_FAILURE=0
export RLINF_CODE_WORKING_DIR=0

test ! -e "$OUTPUT_DIR" || {
  echo "Refusing to reuse OUTPUT_DIR: $OUTPUT_DIR" >&2
  exit 1
}
mkdir -p "$OUTPUT_DIR"

exec > >(tee -i "$OUTPUT_DIR/driver.log") 2>&1

record_startup_exit() {
  local startup_exit_code="$?"
  trap - EXIT
  printf 'STARTUP_EXIT_CODE=%s\n' "$startup_exit_code"
  sleep 2
  exit "$startup_exit_code"
}
trap record_startup_exit EXIT

printf 'RLINF_COMMIT=%s\n' "$RLINF_COMMIT"
printf 'MODEL_REVISION=%s\n' "$MODEL_REVISION"
printf 'DATA_REVISION=%s\n' "$DATA_REVISION"
printf 'RUN_ID=%s\n' "$RUN_ID"
printf 'TARGET_STEPS=%s\n' "$TARGET_STEPS"
printf 'OUTPUT_DIR=%s\n' "$OUTPUT_DIR"
printf 'NCCL_IB_DISABLE=%s\n' "$NCCL_IB_DISABLE"

echo "=== [Phase] Checking pinned inputs ==="
findmnt -T "$WORK_ROOT"
test -w "$WORK_ROOT"
test "$(git -C "$RLINF_ROOT" rev-parse HEAD)" = "$RLINF_COMMIT"
test -z "$(git -C "$RLINF_ROOT" status --porcelain)"
test -r "$RLINF_ROOT/tests/e2e_tests/reasoning/qwen2.5-1.5b-grpo-pipeline-fsdp-sgl.yaml"
test -r "$RLINF_ROOT/examples/reasoning/main_grpo.py"
test -r "$VERIFY_SCRIPT"
test -s "$MODEL_ROOT/config.json"
test -s "$MODEL_ROOT/tokenizer_config.json"
test -s "$MODEL_ROOT/model.safetensors"
test -s "$DATA_FILE"

echo "=== [Phase] Activating the RLinf environment ==="
source /opt/venv/reason/bin/activate
python --version
ray status

echo "=== [Phase] Checking the two-Worker Ray runtime ==="
python "$VERIFY_SCRIPT" preflight \
  --expected-nodes 2 \
  --expected-gpus-per-node 2 \
  --report "$OUTPUT_DIR/preflight.json"

echo "=== [Phase] Launching the RLinf GRPO Driver ==="
OVERRIDES=(
  "hydra.job.chdir=false"
  "hydra.run.dir=$OUTPUT_DIR/hydra"
  "hydra.output_subdir=.hydra"
  "cluster.num_nodes=2"
  "runner.max_epochs=1"
  "runner.max_steps=$TARGET_STEPS"
  "runner.val_check_interval=1"
  "runner.save_interval=$TARGET_STEPS"
  "runner.output_dir=$OUTPUT_ROOT"
  "runner.experiment_name=$RUN_ID"
  "runner.logger.log_path=$OUTPUT_DIR"
  "runner.logger.experiment_name=$RUN_ID"
  "+runner.per_worker_log=true"
  "inference.model.model_path=$MODEL_ROOT"
  "inference.tokenizer.tokenizer_model=$MODEL_ROOT"
  "rollout.model.model_path=$MODEL_ROOT"
  "actor.model.model_path=$MODEL_ROOT"
  "actor.tokenizer.tokenizer_model=$MODEL_ROOT"
  "+actor.fsdp_config.save_full_model_weights=false"
  "data.train_data_paths=[$DATA_FILE]"
  "data.val_data_paths=[$DATA_FILE]"
  "data.seed=1234"
  "actor.seed=1234"
)

RLINF_EXIT_CODE=0
python "$RLINF_ROOT/examples/reasoning/main_grpo.py" \
  --config-path "$RLINF_ROOT/tests/e2e_tests/reasoning" \
  --config-name qwen2.5-1.5b-grpo-pipeline-fsdp-sgl \
  "${OVERRIDES[@]}" || RLINF_EXIT_CODE="$?"

printf 'RLINF_PYTHON_EXIT_CODE=%s\n' "$RLINF_EXIT_CODE"
if (( RLINF_EXIT_CODE != 0 )); then
  exit "$RLINF_EXIT_CODE"
fi

sleep 2
echo "=== [Phase] Checking GRPO metrics and persistent output ==="
python "$VERIFY_SCRIPT" postflight \
  --output-dir "$OUTPUT_DIR" \
  --target-steps "$TARGET_STEPS" \
  --expected-nodes 2

脚本默认运行两个 step。第一个 step 产生可检查的 TensorBoard 指标，第二个 step 保存分布式 checkpoint。脚本会拒绝覆盖已有 RUN_ID 目录，并在 RLinf 进程返回后继续检查业务信号，避免日志中已经出现异常但任务仍被记录为成功。

在创建任务前，检查两个文件的语法并设置执行权限：

language-bash

export TOOL_ROOT=/mnt/rlinf-reproduction/tools/two-worker-grpo-pipeline

cd "$TOOL_ROOT"
python3 -m py_compile verify-grpo-pipeline-v1.py
bash -n run-grpo-pipeline-v1.sh
chmod 0755 verify-grpo-pipeline-v1.py run-grpo-pipeline-v1.sh

两个语法检查都以退出码 0 结束后，再创建训练任务。

设置 Hydra、训练输出和无 RDMA 网络

本场景使用以下路径和运行设置：

设置	本场景值	作用
`hydra.job.chdir`	`false`	保持 Driver 的工作目录稳定
`hydra.run.dir`	`$OUTPUT_DIR/hydra`	把 Hydra 运行文件写入本次共享存储目录
`hydra.output_subdir`	`.hydra`	保留解析后配置和本次 overrides
`runner.output_dir` + `runner.experiment_name`	`$WORK_ROOT/runs/two-worker-grpo-pipeline/＜RUN_ID＞`	统一保存主日志和 checkpoint
`runner.logger.log_path`	`$OUTPUT_DIR`	保存 TensorBoard 和逐 Worker 日志
`NCCL_IB_DISABLE`	`1`	在未启用 RDMA 的任务中使用 IP 网络

每次运行使用新的 RUN_ID。本教程不设置 hydra.output_subdir=null，因为解析后配置用于核对拓扑和输入路径。

运行这个场景

以下步骤先校验共享输入，再创建 Ray 任务并检查训练结果。

Step 1 使用 AICoder 校验共享存储中的输入

在挂载同一共享存储的 AICoder 或开发机中执行：

language-bash

set -e
set -o pipefail

export WORK_ROOT=/mnt/rlinf-reproduction
export RLINF_COMMIT=7d07a4212ee6858cc333e1d4fab7a37256d1f839
export MODEL_REVISION=6fc93244f442ee2b5ab5c8000687ef5f7ffe1d03
export DATA_REVISION=1799c00be3f1216ab55a5cae3562d654dbfd7d82
export RLINF_ROOT="$WORK_ROOT/code/checkouts/$RLINF_COMMIT"
export MODEL_ROOT="$WORK_ROOT/models/modelscope/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B/$MODEL_REVISION"
export DATA_ROOT="$WORK_ROOT/datasets/inclusionAI/AReaL-boba-Data/$DATA_REVISION"

findmnt -T "$WORK_ROOT"
test -w "$WORK_ROOT"
test "$(git -C "$RLINF_ROOT" rev-parse HEAD)" = "$RLINF_COMMIT"
test -z "$(git -C "$RLINF_ROOT" status --porcelain)"
test -r "$RLINF_ROOT/tests/e2e_tests/reasoning/qwen2.5-1.5b-grpo-pipeline-fsdp-sgl.yaml"
test -r "$WORK_ROOT/tools/two-worker-grpo-pipeline/verify-grpo-pipeline-v1.py"
test -r "$WORK_ROOT/tools/two-worker-grpo-pipeline/run-grpo-pipeline-v1.sh"
(cd "$MODEL_ROOT" && sha256sum -c SHA256SUMS)
(cd "$DATA_ROOT" && sha256sum -c SHA256SUMS)

所有命令都以退出码 0 结束后再分配 GPU。训练任务不会修复或下载缺失的输入。

Step 2 创建双 Worker 四 GPU Ray 训练任务

按使用 Spot 资源提交训练任务打开任务创建流程，并填写以下设置。

AIStudio 字段	本场景设置
任务类型	训练任务
资源类型	Spot 资源
可用区和 GPU	本文使用宁夏 B、A100-SXM4-80GB；请在自己的账号中选择符合要求的资源
分布式框架	Ray
Worker 数量	`2`
每个 Worker 的 GPU 数量	`2`
RDMA 配置	不启用
镜像	上传或导入到自己租户后的完整镜像 tag
自动重启	首次验证不启用

继续前，确认资源摘要显示 Ray、2 个 Worker、每个 Worker 2 张 A100。改变 Worker 数量或单 Worker GPU 数量会改变资源 rank 到节点的映射，不能直接复用本文脚本。

Step 3 挂载共享存储并设置运行参数

在任务创建页中完成以下配置：

选择已经保存固定输入的共享高性能存储，把 Worker 内访问地址设置为 /mnt/rlinf-reproduction；
添加环境变量 WORK_ROOT=/mnt/rlinf-reproduction；
添加环境变量 RUN_ID，例如 rlinf-grpo-pipeline-2x2-20260821-01；
添加环境变量 TARGET_STEPS，本场景设置为 2；
添加环境变量 PATH，值为 /opt/venv/reason/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin；
把 TensorBoard 路径设置为 /mnt/rlinf-reproduction/runs/two-worker-grpo-pipeline/${RUN_ID}/tensorboard。

${RUN_ID} 来自上面添加的任务环境变量，平台会在启动 TensorBoard 前替换它。具体用法见使用环境变量同步 TensorBoard 日志路径。AIStudio 会在执行启动命令前调用 ray CLI 来建立集群，因此任务级 PATH 必须包含镜像中的 /opt/venv/reason/bin。确认两个 Worker 都把目标存储卷挂载到 /mnt/rlinf-reproduction，并让 RuntimeEnv、启动命令和 TensorBoard 使用该路径。对应 RUN_ID 的输出目录必须尚不存在。

Step 4 填写 Ray RuntimeEnv JSON 配置

在 RuntimeEnv 配置栏中填写以下 JSON。working_dir 固定 RLinf 代码版本，env_vars 把离线模式、Python 路径和无 RDMA 网络设置传给 Ray Worker。

language-json

{
  "working_dir": "/mnt/rlinf-reproduction/code/checkouts/7d07a4212ee6858cc333e1d4fab7a37256d1f839",
  "excludes": [
    "/.git/",
    "/docs/"
  ],
  "env_vars": {
    "PYTHONPATH": "/mnt/rlinf-reproduction/code/checkouts/7d07a4212ee6858cc333e1d4fab7a37256d1f839:/opt/Megatron-LM",
    "HF_HUB_OFFLINE": "1",
    "HF_DATASETS_OFFLINE": "1",
    "TRANSFORMERS_OFFLINE": "1",
    "TOKENIZERS_PARALLELISM": "false",
    "NCCL_IB_DISABLE": "1",
    "RLINF_CATCH_FAILURE": "0",
    "RLINF_CODE_WORKING_DIR": "0"
  }
}

确认 working_dir 和 PYTHONPATH 都以 /mnt/rlinf-reproduction 开头。不要在 RuntimeEnv 中设置 RAY_ADDRESS、Head 地址或端口；这些值由 AIStudio 管理。镜像已经包含训练依赖，本场景也不在 RuntimeEnv 中安装 pip 包。

Step 5 填写简短的 Driver 启动命令

将以下代码块完整粘贴到启动命令。启动命令只核对并执行共享存储中的版本化 Shell 文件，不会临时生成 Python 或 YAML 文件。

language-bash

set -e
set -o pipefail

export STARTUP_ROOT="$WORK_ROOT/tools/two-worker-grpo-pipeline"
export STARTUP_SCRIPT="$STARTUP_ROOT/run-grpo-pipeline-v1.sh"

test -r "$STARTUP_SCRIPT"
exec /bin/bash "$STARTUP_SCRIPT"

Shell 文件激活 /opt/venv/reason 前只启用 set -e 和 set -o pipefail。镜像的激活脚本可能读取未设置的可选变量，因此这里不启用 set -u。

注意

A100 上的 FlashAttention 兼容处理

本文使用的镜像同时包含 FlashAttention 2 和 FlashAttention 4 相关包。在 A100 上，Transformer Engine 可能选择镜像中不可用的 FlashAttention 4 kernel。预检会在每个 Worker 的临时容器中隐藏 FlashAttention 4 的包元数据，使 Transformer Engine 使用镜像中已有的 FlashAttention 2；共享存储和镜像仓库不会被修改。

Step 6 核对配置并创建任务

在最终摘要中确认 Spot、Ray、2 个 Worker、每个 Worker 2 张 A100、RDMA 未启用、自己的租户镜像，以及共享存储的 Worker 内访问地址为 /mnt/rlinf-reproduction。再核对 WORK_ROOT 和 RuntimeEnv 使用该路径，TensorBoard 路径引用 ${RUN_ID}，RUN_ID 尚未使用，并且 TARGET_STEPS=2，然后创建任务。

Step 7 确认两个 Ray Worker 通过预检

任务开始运行后，在日志中确认：

RAY_CLUSTER_RESOURCES 显示 2 个节点和 4 张 GPU；
日志分别输出 NVIDIA_DRIVER_VERSION_NODE_0 和 NVIDIA_DRIVER_VERSION_NODE_1，且两个版本相同；
两条 WORKER_TOPOLOGY 分别对应 Ray Head 和 Ray Worker，每个节点看到 2 张 A100；
日志出现 RLINF_PREFLIGHT_STATUS=passed。

预检还会检查两个 Worker 使用相同的模型、数据和 RLinf commit，并确认 Transformer Engine 与 SGLang 可以导入。任一检查失败时，启动脚本会在训练开始前退出。

Step 8 确认两个 GRPO step 完成

训练开始后，确认：

日志出现 Step limit given by max_steps=2 reached；
TensorBoard 中可以查看 reward、loss、梯度范数和学习率；
日志末尾出现 RLINF_VERIFICATION_STATUS=passed 和 STARTUP_EXIT_CODE=0。

校验脚本会检查这些指标均为有限值、学习率大于 0，并确认至少写出两个 FSDP checkpoint shard。在这个两步检查任务中，同一组样本得到相同奖励时，梯度范数可能为 0；这不代表校验失败，也不用于判断模型是否开始收敛。页面显示“运行成功”本身不检查这些训练信号。

Step 9 从共享存储回读分布式 checkpoint

任务结束后，在挂载同一共享存储的 AICoder 或开发机中执行。把示例 RUN_ID 替换为本次实际值：

language-bash

set -e
set -o pipefail

export RUN_ID=rlinf-grpo-pipeline-2x2-20260821-01
export OUTPUT_DIR=/mnt/rlinf-reproduction/runs/two-worker-grpo-pipeline/$RUN_ID
export CHECKPOINT_DIR="$OUTPUT_DIR/checkpoints/global_step_2"

test -s "$OUTPUT_DIR/verification.json"
test -s "$CHECKPOINT_DIR/actor/dcp_checkpoint/.metadata"
test "$(find "$CHECKPOINT_DIR/actor/dcp_checkpoint" -type f -name '*.distcp' -size +0 | wc -l | tr -d '[:space:]')" -ge 2

find "$CHECKPOINT_DIR/actor/dcp_checkpoint" -maxdepth 1 -type f -printf '%s %f\n' | sort

verification.json 记录两个节点、资源映射、关键 TensorBoard 指标和 checkpoint shard。任务容器回收后仍能读取该报告和 checkpoint，说明输出已经写入共享存储。

运行更多 step

确认两个 step 的链路后，可以创建新任务并提高 TARGET_STEPS。使用新的 RUN_ID，保持镜像、模型、数据、Worker 数量、单 Worker GPU 数量和资源映射不变；同时根据 checkpoint 实际大小规划共享存储空间。

本教程的脚本要求 TARGET_STEPS 至少为 2，这样 TensorBoard 至少会保存第一个 step 的训练指标。更长运行还需要单独确定训练步数、验证集、保存间隔、恢复策略和模型效果评测，不能只提高这个检查任务的 step 数就把它当作完整训练方案。

处理本场景特有故障

Ray 只发现一个节点

检查任务是否选择 Ray、Worker 数量是否为 2，以及两个 Worker 是否都已进入运行状态。启动命令只连接平台托管集群；不要在 Head 中另外执行 ray start。

Worker 使用了不同的输入或驱动

在 preflight.json 中比较两个 Worker 的驱动版本和 RLinf commit。如果模型或数据 manifest 不一致，确认两个 Worker 挂载了同一个共享存储，并且 Worker 内访问地址和 WORK_ROOT 都是 /mnt/rlinf-reproduction。修复后使用新的 RUN_ID 重跑。

Transformer Engine 找不到可用的 FlashAttention kernel

确认任务使用本文指定的 A100 和镜像，并且没有跳过预检。预检会在每个 Worker 中选择镜像已有的 FlashAttention 2 路径；更换镜像或 GPU 后，需要重新检查依赖兼容性。

RLinf 日志报错但页面显示运行成功

检查任务是否使用 Step 5 中的完整启动命令，以及 RUN_ID 是否对应本次输出。运行后检查要求 step 完成标志、TensorBoard 指标、两个 FSDP shard 和 RLINF_VERIFICATION_STATUS=passed 同时存在；任一项缺失都会让启动脚本返回非 0。

确认多 Worker 训练链路已经贯通

当日志出现 RLINF_VERIFICATION_STATUS=passed，并且可以从共享存储读取验证报告和两个 FSDP checkpoint shard 时，这次双 Worker 运行已经完成。生成、奖励计算、log probability、actor 优化和持久化输出均已贯通。

需要继续扩大训练规模时，请从 RLinf 的训练示例重新规划 Worker、GPU、网络、保存策略和效果评测，不要把这次小规模验证直接当作完整训练方案。

简介

AICoder 功能

租户管理

凭证管理

费用管理

算力资源池管理

存储管理

镜像中心

通知与告警

存储

开发机

任务

推理服务

资源管理

简介

集群管理

简介

实例管理

开始使用

模型与能力

开发工具与集成

计费、限制与用量

排查与支持

更新与参考

Megatron-Infinigence 训练框架

典型模型部署

实战目标 ​

场景信息 ​

用小规模运行检查多 Worker 训练链路 ​

本场景在 AIStudio 中怎样运行 ​

开始前准备 ​

在目标可用区准备任务镜像 ​

在共享存储中固定 RLinf checkout ​

准备固定版本的模型和数据 ​

准备运行前后校验脚本 ​

准备 GRPO 训练启动脚本 ​

设置 Hydra、训练输出和无 RDMA 网络 ​

运行这个场景 ​

Step 1 使用 AICoder 校验共享存储中的输入 ​

Step 2 创建双 Worker 四 GPU Ray 训练任务 ​

Step 3 挂载共享存储并设置运行参数 ​

Step 4 填写 Ray RuntimeEnv JSON 配置 ​

Step 5 填写简短的 Driver 启动命令 ​

Step 6 核对配置并创建任务 ​

Step 7 确认两个 Ray Worker 通过预检 ​

Step 8 确认两个 GRPO step 完成 ​

Step 9 从共享存储回读分布式 checkpoint ​

运行更多 step ​

处理本场景特有故障 ​

Ray 只发现一个节点 ​

Worker 使用了不同的输入或驱动 ​

Transformer Engine 找不到可用的 FlashAttention kernel ​

RLinf 日志报错但页面显示运行成功 ​

确认多 Worker 训练链路已经贯通 ​

实战目标

场景信息

用小规模运行检查多 Worker 训练链路

本场景在 AIStudio 中怎样运行

开始前准备

在目标可用区准备任务镜像

在共享存储中固定 RLinf checkout

准备固定版本的模型和数据

准备运行前后校验脚本

准备 GRPO 训练启动脚本

设置 Hydra、训练输出和无 RDMA 网络

运行这个场景

Step 1 使用 AICoder 校验共享存储中的输入

Step 2 创建双 Worker 四 GPU Ray 训练任务

Step 3 挂载共享存储并设置运行参数

Step 4 填写 Ray RuntimeEnv JSON 配置

Step 5 填写简短的 Driver 启动命令

Step 6 核对配置并创建任务

Step 7 确认两个 Ray Worker 通过预检

Step 8 确认两个 GRPO step 完成

Step 9 从共享存储回读分布式 checkpoint

运行更多 step

处理本场景特有故障

Ray 只发现一个节点

Worker 使用了不同的输入或驱动

Transformer Engine 找不到可用的 FlashAttention kernel

RLinf 日志报错但页面显示运行成功

确认多 Worker 训练链路已经贯通