Skip to content

在 AIStudio 训练任务中使用 RLinf 复现 Qwen 1.5B 单 GPU GRPO

本教程在一个 AIStudio Worker 上运行十步 GRPO。DeepSeek-R1-Distill-Qwen-1.5B 生成数学题回答,RLinf 根据 AReaL-boba 样本中的参考答案计算奖励,再在同一张 GPU 上完成优势计算和 actor 更新。

实战目标

您将创建一个只运行十个 step 的单 GPU GRPO 任务。任务结束后,可以确认:

  • 模型和数据从共享存储加载,RLinf 完成生成、奖励计算和 actor 更新;
  • 十个 step 全部完成,TensorBoard 中可以查看奖励和训练指标;
  • 启动日志、解析后配置、TensorBoard 事件和 global_step_10 checkpoint 保留在共享存储中。

场景信息

创建任务时使用以下设置。配置文件名中的 qwen2.5-1.5b 表示 RLinf 模型适配类型;本场景实际加载的权重是 DeepSeek-R1-Distill-Qwen-1.5B

  • RLinf

    • commit:89b0cd5fce528180559bbd50d055fb2e21a25bb5
    • 配置:examples/reasoning/config/math/qwen2.5-1.5b-single-gpu.yaml
    • 入口:examples/reasoning/main_grpo.py
    • 要求:从共享存储读取指定 commit,且不修改代码。
  • 上游镜像rlinf/rlinf:agentic-rlinf0.4-torch2.11.0-sglang0.5.12.post1-vllm0.23.0-megatron0.17.0-te2.17,来自 Docker Hub。

  • AIStudio 任务镜像cr.infini-ai.com/<tenant-id>/rlinf:agentic-rlinf0.4-torch2.11.0-sglang0.5.12.post1-vllm0.23.0-megatron0.17.0-te2.17。将上游镜像上传或导入自己租户在任务可用区的镜像仓库。

  • 模型:ModelScope 的 deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B@6fc93244f442ee2b5ab5c8000687ef5f7ffe1d03。从共享存储读取固定 revision。

  • 数据inclusionAI/AReaL-boba-Data@1799c00be3f1216ab55a5cae3562d654dbfd7d82 中的 AReaL-boba-106k.jsonl。从共享存储读取,同一文件用于 train 和 val。

  • 资源来源:Spot 资源。中断后保留原输出,并使用新的 RUN_ID 完整重跑。

  • 分布式框架:单机;RLinf 在 Worker 内启动本地 Ray。使用 1 个 Worker 时,创建页只提供“单机”选项,不提供“Ray”。

  • Worker 拓扑:1 个 Worker、1 个节点、1 张 A100-SXM4-80GB。rollout、reward 和 actor 共用 GPU 0

  • 共享存储:挂载到 /mnt/rlinf-reproduction。固定代码、模型和数据,并把日志、TensorBoard 和 checkpoint 写入 runs/

  • 运行范围runner.max_steps=10。保留上游单 GPU 配置的 batch 和序列长度,只限制运行步数。

本场景在 AIStudio 中怎样运行

AIStudio 创建 Worker、挂载共享存储并执行启动命令。分布式框架的选择决定 Ray 由谁管理:

  • 单机(本教程):AIStudio 提供 1 个 Worker,页面只提供“单机”选项;RLinf 在 Worker 内启动本地 Ray。启动命令直接运行 RLinf 入口,不设置 RAY_ADDRESS,也不执行 ray startray job submit
  • Ray:AIStudio 管理 Ray Head 和 Worker,适用于跨多个 Worker 放置组件的任务。使用这种模式时,还要检查各 Worker 的共享路径、驱动、网络和 rank 放置。

本教程使用以下拓扑:

language-text
AIStudio 单机训练任务
└── Worker 0 / GPU 0
    └── RLinf 本地 Ray 运行时
        ├── RLinf Driver
        ├── RolloutGroup / SGLang
        ├── RewardGroup / math reward
        └── ActorGroup / Megatron

Worker 从共享存储读取固定代码、模型和数据,并把日志、TensorBoard 和 checkpoint 写入本次 RUN_ID 的输出目录。SGLang rollout、奖励计算和 Megatron actor 共用一张 GPU。

开始前准备

在创建 GPU 任务前,使用挂载同一共享存储的 AICoder 或开发机准备镜像、RLinf 代码、模型和数据。本文使用宁夏 B;您可以选择其他可用区,但镜像、任务和共享存储需要位于同一可用区。

在目标可用区准备容器镜像

本场景使用 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 代码

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

language-shell
set -euo pipefail

export WORK_ROOT=/mnt/rlinf-reproduction
export RLINF_COMMIT=89b0cd5fce528180559bbd50d055fb2e21a25bb5
export RLINF_ROOT="$WORK_ROOT/code/$RLINF_COMMIT"

findmnt -T "$WORK_ROOT"
df -hT "$WORK_ROOT"
test -w "$WORK_ROOT"

mkdir -p "$WORK_ROOT/code"
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)"

固定 commit 让后续任务使用相同的代码版本。目标目录已存在时,git clone 会停止;请先检查现有目录的 commit 和工作树,再决定是否使用该目录。

训练任务只读取并校验准备好的代码,不会执行 git clonegit fetch 或网络下载。代码、数据和输出目录的一般准备方法见准备代码、数据与输出目录

在共享存储中下载固定版本的模型和数据

以下命令使用 ModelScope CLI 下载模型,并使用 hf CLI 下载 AReaL-boba 数据。请在挂载共享存储的 AICoder 或开发机中执行。工具准备方法见AICoder 下载模型使用 HuggingFace 镜像站下载模型和数据集

模型路径末尾是 ModelScope 仓库的固定 revision。请按原样保留,使后续任务继续读取同一版本的模型文件。

language-shell
set -euo pipefail

export WORK_ROOT=/mnt/rlinf-reproduction
export MODEL_REVISION=6fc93244f442ee2b5ab5c8000687ef5f7ffe1d03
export DATA_REVISION=1799c00be3f1216ab55a5cae3562d654dbfd7d82
export MODEL_WEIGHT_SHA256=58858233513d76b8703e72eed6ce16807b523328188e13329257fb9594462945
export DATA_SHA256=2ee1c91623cbbcefd144e8063d1820576099b799d174ff10fd908798b3dc32b0
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"

test ! -e "$MODEL_ROOT" || {
  echo "Model target already exists: $MODEL_ROOT" >&2
  exit 1
}
test ! -e "$DATA_ROOT" || {
  echo "Dataset target already exists: $DATA_ROOT" >&2
  exit 1
}
mkdir -p "$MODEL_ROOT" "$DATA_ROOT"

modelscope download deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B \
  --revision "$MODEL_REVISION" \
  --local-dir "$MODEL_ROOT" \
  --max-workers 2

hf download inclusionAI/AReaL-boba-Data AReaL-boba-106k.jsonl \
  --repo-type dataset \
  --revision "$DATA_REVISION" \
  --local-dir "$DATA_ROOT"

test "$(stat -c %s "$MODEL_ROOT/model.safetensors")" -eq 3554214621
printf '%s  %s\n' "$MODEL_WEIGHT_SHA256" "$MODEL_ROOT/model.safetensors" \
  | sha256sum -c -
test "$(stat -c %s "$DATA_ROOT/AReaL-boba-106k.jsonl")" -eq 71571518
printf '%s  %s\n' "$DATA_SHA256" "$DATA_ROOT/AReaL-boba-106k.jsonl" \
  | sha256sum -c -

(
  cd "$MODEL_ROOT"
  find . -type f ! -path './.cache/*' ! -name SHA256SUMS -print0 \
    | sort -z \
    | xargs -0 sha256sum > SHA256SUMS
  sha256sum -c SHA256SUMS
)

(
  cd "$DATA_ROOT"
  sha256sum AReaL-boba-106k.jsonl > SHA256SUMS
  sha256sum -c SHA256SUMS
)

检查模型文件和数据字段

language-shell
python - "$MODEL_ROOT" "$DATA_ROOT/AReaL-boba-106k.jsonl" <<'PY'
import json
import sys
from pathlib import Path

model_root = Path(sys.argv[1])
data_file = Path(sys.argv[2])

assert (model_root / "config.json").is_file()
assert (model_root / "tokenizer_config.json").is_file()

index_file = model_root / "model.safetensors.index.json"
if index_file.is_file():
    index = json.loads(index_file.read_text())
    shards = sorted(set(index["weight_map"].values()))
else:
    shards = [path.name for path in model_root.glob("*.safetensors")]

assert shards, "No model weight files found"
for shard in shards:
    path = model_root / shard
    assert path.is_file() and path.stat().st_size > 0, path

count = 0
with data_file.open() as stream:
    for line_number, line in enumerate(stream, 1):
        record = json.loads(line)
        assert isinstance(record.get("prompt"), str) and record["prompt"].strip(), line_number
        assert isinstance(record.get("solutions"), list) and record["solutions"], line_number
        count += 1

assert count > 0
print("model_shards", len(shards))
print("dataset_records", count)
print("model_bytes", sum(path.stat().st_size for path in model_root.rglob("*") if path.is_file()))
print("dataset_bytes", data_file.stat().st_size)
PY

下载、SHA-256 和字段检查都以退出码 0 结束后,继续创建训练任务。

准备训练配置、运行标识和持久化目录

本场景直接读取固定 checkout 中的 qwen2.5-1.5b-single-gpu.yaml,并在启动命令中覆盖模型、tokenizer、数据和输出路径。保留以下范围和命名规则:

  • TARGET_STEPS=10,启动命令将其传给 runner.max_steps
  • rollout_batch_size=128group_size=2 和 2048 token 序列范围保持不变;
  • 数据、actor 和 Python 哈希 seed 均固定为 1234
  • 每次运行使用新的 RUN_ID,只包含字母、数字、点、下划线或连字符;
  • 输出写入 /mnt/rlinf-reproduction/runs/single-gpu-grpo/<RUN_ID>/,启动命令拒绝复用已有目录。

设置 Hydra 和训练输出目录

Hydra 负责解析配置,RLinf 的 runner 字段负责本场景的训练输出。以下设置避免 Hydra 改变进程工作目录,并把 RLinf 结果写入唯一的共享存储目录。

设置本场景值结果
hydra.job.chdirfalse保持 RLinf 进程工作目录稳定
runner.output_dir/mnt/rlinf-reproduction/runs/single-gpu-grpo设置所有运行共用的持久化根目录
runner.experiment_name<RUN_ID>为每次运行创建独立的日志、TensorBoard 和 checkpoint 目录

确认运行条件

本文使用宁夏 B 的 A100-SXM4-80GB。请在自己的账号中选择满足以下条件的资源:

  • 任务可以把已准备的共享存储挂载到 /mnt/rlinf-reproduction,并能读取代码、模型和数据、写入 runs/
  • 镜像中存在 /opt/venv/reason,并可在 A100 上使用 FlashAttention 2 路径导入 Transformer Engine;
  • Worker 只分配 1 张 A100-SXM4-80GB;
  • 镜像、任务和共享存储位于同一可用区。

运行这个场景

按以下步骤校验输入、创建任务并检查结果。

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

分配 GPU 前,在挂载同一共享存储的 AICoder 或开发机中执行:

language-shell
set -euo pipefail

export WORK_ROOT=/mnt/rlinf-reproduction
export RLINF_COMMIT=89b0cd5fce528180559bbd50d055fb2e21a25bb5
export MODEL_REVISION=6fc93244f442ee2b5ab5c8000687ef5f7ffe1d03
export DATA_REVISION=1799c00be3f1216ab55a5cae3562d654dbfd7d82
export RLINF_ROOT="$WORK_ROOT/code/$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/examples/reasoning/config/math/qwen2.5-1.5b-single-gpu.yaml"
(cd "$MODEL_ROOT" && sha256sum -c SHA256SUMS)
(cd "$DATA_ROOT" && sha256sum -c SHA256SUMS)

所有命令退出码为 0,且校验清单中的文件均返回 OK 后再创建任务。检查失败时,在 AICoder 或开发机中修复固定输入;GPU 启动命令不会克隆代码或下载模型、数据。

Step 2 创建 1 Worker 单 GPU Spot 训练任务

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

AIStudio 字段本场景设置
任务类型训练任务
资源类型Spot 资源
可用区和 GPU本文使用宁夏 B、1 张 A100-SXM4-80GB;请在自己的账号中选择符合要求的资源
分布式框架单机
Worker 数量1
镜像上传或导入到自己租户后的 cr.infini-ai.com/<tenant-id>/rlinf:agentic-rlinf0.4-torch2.11.0-sglang0.5.12.post1-vllm0.23.0-megatron0.17.0-te2.17
自动排队和自动重启不启用

继续前,确认资源摘要显示所选可用区、Spot、1 个 Worker、1 张 A100 和单机框架。本文的启动命令适用于这一资源组合。

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

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

  • 选择已经保存固定输入的共享高性能存储,并把系统挂载路径设置为 /mnt/rlinf-reproduction
  • 添加环境变量 RUN_ID,例如 rlinf-grpo-10step-20260820-01
  • 添加环境变量 TARGET_STEPS,本场景设置为 10
  • 把 TensorBoard 路径设置为 /mnt/rlinf-reproduction/runs/single-gpu-grpo/${RUN_ID}/tensorboard

${RUN_ID} 来自上面添加的任务环境变量,平台会在启动 TensorBoard 前替换它。具体用法见使用环境变量同步 TensorBoard 日志路径。挂载完成后,训练任务中的 /mnt/rlinf-reproduction 必须与 AICoder 中准备代码、模型和数据的目录指向同一个存储位置。每次运行只需设置新的 RUN_ID;如果对应输出目录已经存在,启动命令会停止,避免覆盖上一次结果。

Step 4 设置包含输入、驱动和训练完成检查的启动命令

将以下代码块完整粘贴到启动命令中。命令直接启动 Bash,不依赖共享存储中的脚本文件。把 IMAGE_TAG 中的 te-b905754427352261 替换为自己的 <tenant-id>,并确认替换后的完整 tag 与任务选择的镜像一致。

注意

为什么 A100 需要 FlashAttention 兼容处理

在本文使用的镜像和 A100 上,Transformer Engine 默认选择 FlashAttention 4/CUTLASS,但找不到可用的 kernel。启动命令会临时停用 FlashAttention 4 的包元数据,让 Transformer Engine 使用镜像中已有的 FlashAttention 2。

这个处理不会安装、卸载或永久修改镜像,只适用于本文的镜像和 A100 配置。更换镜像或 GPU 后,请重新检查 Transformer Engine 的导入结果。

language-shell
exec /bin/bash -s <<'RLINF_GRPO'
set -euo pipefail

: "${RUN_ID:?Set RUN_ID in the AIStudio task}"
TARGET_STEPS="${TARGET_STEPS:-10}"

[[ "$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]*$ ]] || {
  echo "TARGET_STEPS must be a positive integer" >&2
  exit 1
}

export WORK_ROOT=/mnt/rlinf-reproduction
# 导入镜像后,把以下 URL 中的租户 ID 替换为自己的 <tenant-id>。
export IMAGE_TAG=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
export RLINF_COMMIT=89b0cd5fce528180559bbd50d055fb2e21a25bb5
export MODEL_SOURCE=modelscope
export MODEL_REVISION=6fc93244f442ee2b5ab5c8000687ef5f7ffe1d03
export DATA_REVISION=1799c00be3f1216ab55a5cae3562d654dbfd7d82
export RLINF_ROOT="$WORK_ROOT/code/$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 OUTPUT_ROOT="$WORK_ROOT/runs/single-gpu-grpo"
export OUTPUT_DIR="$OUTPUT_ROOT/$RUN_ID"
export PYTHONPATH="$RLINF_ROOT:/opt/Megatron-LM:${PYTHONPATH:-}"
export CUDA_DEVICE_MAX_CONNECTIONS=1
export TOKENIZERS_PARALLELISM=false
export RAY_DEDUP_LOGS=0
export HF_HUB_OFFLINE=1
export HF_DATASETS_OFFLINE=1
export TRANSFORMERS_OFFLINE=1
export PYTHONHASHSEED=1234
export DATA_SEED=1234
export ACTOR_SEED=1234

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

record_startup_exit() {
  local startup_exit_code="$1"
  local -a exit_log_status=()

  trap - EXIT
  set +e
  printf 'STARTUP_EXIT_CODE=%s\n' "$startup_exit_code" |
    tee -a "$OUTPUT_DIR/driver.log" >&2
  exit_log_status=("${PIPESTATUS[@]}")
  if (( startup_exit_code == 0 &&
        (exit_log_status[0] != 0 || exit_log_status[1] != 0) )); then
    exit 1
  fi
  exit "$startup_exit_code"
}
trap 'record_startup_exit "$?"' EXIT

if ! NVIDIA_DRIVER_VERSIONS="$(
  nvidia-smi --query-gpu=driver_version --format=csv,noheader,nounits |
    sed '/^[[:space:]]*$/d' |
    sort -u
)"; then
  echo "Unable to detect the NVIDIA driver version" >&2
  exit 1
fi
NVIDIA_DRIVER_VERSION_COUNT="$(
  printf '%s\n' "$NVIDIA_DRIVER_VERSIONS" | wc -l | tr -d '[:space:]'
)"
if [[ -z "$NVIDIA_DRIVER_VERSIONS" || "$NVIDIA_DRIVER_VERSION_COUNT" -ne 1 ]]; then
  printf 'Expected one NVIDIA driver version, detected: %s\n' \
    "${NVIDIA_DRIVER_VERSIONS:-<none>}" >&2
  exit 1
fi
{
  printf 'NVIDIA_DRIVER_VERSION=%s\n' "$NVIDIA_DRIVER_VERSIONS"
  printf 'IMAGE_TAG=%s\n' "$IMAGE_TAG"
  printf 'RLINF_COMMIT=%s\n' "$RLINF_COMMIT"
  printf 'DATA_SEED=%s\n' "$DATA_SEED"
  printf 'ACTOR_SEED=%s\n' "$ACTOR_SEED"
  printf 'MODEL_SOURCE=%s\n' "$MODEL_SOURCE"
  printf 'MODEL_REVISION=%s\n' "$MODEL_REVISION"
  printf 'DATA_REVISION=%s\n' "$DATA_REVISION"
  printf 'TARGET_STEPS=%s\n' "$TARGET_STEPS"
  printf 'OUTPUT_DIR=%s\n' "$OUTPUT_DIR"
} | tee "$OUTPUT_DIR/driver.log"

test "$(git -C "$RLINF_ROOT" rev-parse HEAD)" = "$RLINF_COMMIT"
test -z "$(git -C "$RLINF_ROOT" status --porcelain)"
(cd "$MODEL_ROOT" && sha256sum -c SHA256SUMS)
(cd "$(dirname "$DATA_FILE")" && sha256sum -c SHA256SUMS)

source /opt/venv/reason/bin/activate

GPU_COMPUTE_CAPABILITY="$(
python - <<'PY'
import torch

major, minor = torch.cuda.get_device_capability(0)
print(f"{major}.{minor}")
PY
)"
if [[ "$GPU_COMPUTE_CAPABILITY" != "8.0" ]]; then
  printf 'The A100 compatibility step requires compute capability 8.0; detected %s\n' \
    "$GPU_COMPUTE_CAPABILITY" >&2
  exit 1
fi

FA4_VERSION_BEFORE_WORKAROUND="$(
python - <<'PY'
from importlib import metadata

try:
    print(metadata.version("flash-attn-4"))
except metadata.PackageNotFoundError:
    print("<not-installed>")
PY
)"
FA4_DIST_INFO="$(
python - <<'PY'
from importlib import metadata

try:
    print(metadata.distribution("flash-attn-4")._path)
except metadata.PackageNotFoundError:
    pass
PY
)"

if [[ -n "$FA4_DIST_INFO" ]]; then
  case "$FA4_DIST_INFO" in
    /opt/venv/reason/lib/python*/site-packages/flash_attn_4-*.dist-info) ;;
    *)
      echo "Unexpected flash-attn-4 metadata path: $FA4_DIST_INFO" >&2
      exit 1
      ;;
  esac

  FA4_DISABLED_PATH="${FA4_DIST_INFO}.a100-disabled"
  test ! -e "$FA4_DISABLED_PATH" || {
    echo "Refusing to overwrite disabled metadata: $FA4_DISABLED_PATH" >&2
    exit 1
  }
  mv -- "$FA4_DIST_INFO" "$FA4_DISABLED_PATH"
  FA4_METADATA_ACTION="disabled"
else
  FA4_METADATA_ACTION="already-absent"
fi

python - <<'PY'
from importlib import metadata

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

{
  printf 'GPU_COMPUTE_CAPABILITY=%s\n' "$GPU_COMPUTE_CAPABILITY"
  printf 'FLASH_ATTN_4_VERSION_BEFORE_WORKAROUND=%s\n' \
    "$FA4_VERSION_BEFORE_WORKAROUND"
  printf 'FLASH_ATTN_4_METADATA_ACTION=%s\n' "$FA4_METADATA_ACTION"
} | tee -ia "$OUTPUT_DIR/driver.log"

if [[ -n "${RAY_ADDRESS:-}" ]]; then
  echo "Unexpected RAY_ADDRESS in an AIStudio single-machine task" >&2
  exit 1
fi
if ray status >/dev/null 2>&1; then
  echo "Unexpected pre-existing Ray runtime" >&2
  exit 1
fi

python - <<'PY' | tee -ia "$OUTPUT_DIR/driver.log"
from importlib import import_module, metadata
import math
import torch

assert torch.cuda.is_available(), "CUDA is unavailable"
assert torch.cuda.device_count() == 1, torch.cuda.device_count()
value = torch.tensor([2.0], device="cuda").square().item()
assert math.isfinite(value) and value == 4.0

for package in (
    "torch",
    "ray",
    "sglang",
    "transformers",
    "transformer-engine",
    "flash-attn",
    "flash-attn-4",
    "nvidia-cutlass-dsl",
):
    try:
        version = metadata.version(package)
    except metadata.PackageNotFoundError:
        version = "<not-installed>"
    print(package, version)
print("cuda", torch.version.cuda)
print("gpu", torch.cuda.get_device_name(0))
print("gpu_compute_capability", ".".join(map(str, torch.cuda.get_device_capability(0))))
print("gpu_memory_bytes", torch.cuda.get_device_properties(0).total_memory)

import_module("transformer_engine.pytorch")
print("transformer_engine.pytorch", "import-ok")
PY

python - "$DATA_FILE" <<'PY' | tee -ia "$OUTPUT_DIR/driver.log"
import json
import sys

count = 0
with open(sys.argv[1]) as stream:
    for line_number, line in enumerate(stream, 1):
        record = json.loads(line)
        assert isinstance(record.get("prompt"), str) and record["prompt"].strip(), line_number
        assert isinstance(record.get("solutions"), list) and record["solutions"], line_number
        count += 1
assert count > 0
print("dataset_records", count)
PY

set +e
python "$RLINF_ROOT/examples/reasoning/main_grpo.py" \
  --config-path "$RLINF_ROOT/examples/reasoning/config/math" \
  --config-name qwen2.5-1.5b-single-gpu \
  "hydra.job.chdir=false" \
  "runner.max_epochs=$TARGET_STEPS" \
  "runner.max_steps=$TARGET_STEPS" \
  "runner.output_dir=$OUTPUT_ROOT" \
  "runner.experiment_name=$RUN_ID" \
  "rollout.model.model_path=$MODEL_ROOT" \
  "actor.tokenizer.tokenizer_model=$MODEL_ROOT" \
  "actor.seed=$ACTOR_SEED" \
  "data.seed=$DATA_SEED" \
  "data.train_data_paths=['$DATA_FILE']" \
  "data.val_data_paths=['$DATA_FILE']" \
  2>&1 | tee -ia "$OUTPUT_DIR/driver.log"
PIPE_STATUSES=("${PIPESTATUS[@]}")
set -e

{
  printf 'RLINF_PYTHON_EXIT_CODE=%s\n' "${PIPE_STATUSES[0]}"
  printf 'RLINF_TEE_EXIT_CODE=%s\n' "${PIPE_STATUSES[1]}"
} | tee -ia "$OUTPUT_DIR/driver.log"

if (( PIPE_STATUSES[0] != 0 || PIPE_STATUSES[1] != 0 )); then
  printf 'RLinf or log persistence failed: rlinf=%s tee=%s\n' \
    "${PIPE_STATUSES[0]}" "${PIPE_STATUSES[1]}" >&2
  exit 1
fi

MAIN_LOG="$OUTPUT_DIR/log/main.log"
COMPLETION_MARKER="Step limit given by max_steps=$TARGET_STEPS reached. Stopping run"
CHECKPOINT_DIR="$OUTPUT_DIR/checkpoints/global_step_$TARGET_STEPS"

test -s "$MAIN_LOG" || {
  echo "RLinf main log is missing or empty: $MAIN_LOG" >&2
  exit 1
}
grep -Fq "$COMPLETION_MARKER" "$OUTPUT_DIR/driver.log" || {
  echo "RLinf exited without the expected target-step completion marker" >&2
  exit 1
}
test -s "$OUTPUT_DIR/tensorboard/config.yaml" || {
  echo "TensorBoard config is missing or empty" >&2
  exit 1
}
if ! find "$OUTPUT_DIR/tensorboard" -type f \
  -name 'events.out.tfevents.*' -size +0 -print -quit | grep -q .; then
  echo "No non-empty TensorBoard event file was produced" >&2
  exit 1
fi
if ! find "$CHECKPOINT_DIR" -type f -size +0 -print -quit | grep -q .; then
  echo "The final global_step_$TARGET_STEPS checkpoint is missing or empty" >&2
  exit 1
fi

printf 'RLINF_VERIFICATION_STATUS=passed\n' |
  tee -ia "$OUTPUT_DIR/driver.log"
RLINF_GRPO

RLinf 在 Worker 内启动本地 Ray,不依赖 AIStudio 托管 Ray。启动命令同时检查进程退出码、目标 step 完成标志、TensorBoard 事件和最终 checkpoint;任一检查失败,任务都会返回非 0

Step 5 确认配置并创建任务

在最终摘要中确认 Spot、1 个 Worker、1 张 A100、单机框架、自己的租户镜像和 /mnt/rlinf-reproduction 挂载路径。再确认启动命令中的 IMAGE_TAG 与任务镜像一致,TensorBoard 路径引用 ${RUN_ID},并且 TARGET_STEPS=10,然后创建任务。

本次短时运行不启用自动重启。Spot 资源中断或启动检查失败时,保留原输出并使用新的 RUN_ID 重跑。

Step 6 确认 Worker 通过预检并启动 RLinf 组件

任务进入运行中后,确认:

  1. 日志中的完整 IMAGE_TAG、RLinf commit、模型 revision、数据 revision 和 TARGET_STEPS=10 与任务记录一致。
  2. 日志只包含一条 NVIDIA_DRIVER_VERSION=<version>;可见 GPU 数量为 1,计算能力为 8.0,CUDA 检查和 Transformer Engine 导入均通过。
  3. RLinf 启动包含 1 个节点和 1 个 accelerator 的本地 Ray,并成功创建 RolloutGroup、RewardGroup 和 ActorGroup。
  4. 模型、tokenizer、train 和 val 路径均指向共享存储中已校验的固定版本,日志随后进入生成或训练阶段。

Step 7 确认十个 step 和 GRPO 更新完成

任务运行期间,在日志和 TensorBoard 中依次确认:

  • 模型和 tokenizer 从固定的本地 revision 目录加载,没有回退到网络;
  • AReaL-boba 数据成功解析 promptsolutions
  • 日志已进入生成、数学奖励、GRPO 优势计算和 actor optimizer step;
  • TensorBoard 中存在有限的 rollout/ 奖励类标量和 actor/training/ 损失、梯度或学习率标量;
  • 日志出现 Step limit given by max_steps=10 reached. Stopping run
  • 末尾依次出现 RLINF_PYTHON_EXIT_CODE=0RLINF_TEE_EXIT_CODE=0RLINF_VERIFICATION_STATUS=passedSTARTUP_EXIT_CODE=0,任务最终显示运行成功。

以十步完成标志、TensorBoard 事件、最终 checkpoint 和 RLINF_VERIFICATION_STATUS=passed 作为验收结果,不要只依据页面状态或训练进程退出码。

Step 8 Worker 结束后检查共享存储中的输出

任务结束后,在挂载同一共享存储的 AICoder 或开发机中执行。先把示例中的 RUN_IDTARGET_STEPS<tenant-id> 替换为本次任务的实际值。

language-shell
set -euo pipefail

export RUN_ID=rlinf-grpo-10step-20260820-01
export TARGET_STEPS=10
export EXPECTED_IMAGE_TAG='cr.infini-ai.com/<tenant-id>/rlinf:agentic-rlinf0.4-torch2.11.0-sglang0.5.12.post1-vllm0.23.0-megatron0.17.0-te2.17'

: "${RUN_ID:?Set the completed run ID}"
: "${EXPECTED_IMAGE_TAG:?Set the complete image tag used by the task}"
[[ "$TARGET_STEPS" =~ ^[1-9][0-9]*$ ]] || {
  echo "TARGET_STEPS must be a positive integer" >&2
  exit 1
}
OUTPUT_DIR="/mnt/rlinf-reproduction/runs/single-gpu-grpo/$RUN_ID"
MAIN_LOG="$OUTPUT_DIR/log/main.log"
COMPLETION_MARKER="Step limit given by max_steps=$TARGET_STEPS reached. Stopping run"
CHECKPOINT_DIR="$OUTPUT_DIR/checkpoints/global_step_$TARGET_STEPS"
EXPECTED_RLINF_COMMIT=89b0cd5fce528180559bbd50d055fb2e21a25bb5
EXPECTED_MODEL_REVISION=6fc93244f442ee2b5ab5c8000687ef5f7ffe1d03
EXPECTED_DATA_REVISION=1799c00be3f1216ab55a5cae3562d654dbfd7d82

test -s "$OUTPUT_DIR/driver.log"
test -s "$MAIN_LOG"
test -s "$OUTPUT_DIR/tensorboard/config.yaml"
find "$OUTPUT_DIR/tensorboard" -type f -name 'events.out.tfevents.*' -size +0 -print -quit \
  | grep -q .
find "$CHECKPOINT_DIR" -type f -size +0 -print -quit \
  | grep -q .

grep -aFx "IMAGE_TAG=$EXPECTED_IMAGE_TAG" "$OUTPUT_DIR/driver.log"
grep -aFx "RLINF_COMMIT=$EXPECTED_RLINF_COMMIT" "$OUTPUT_DIR/driver.log"
grep -aFx "MODEL_REVISION=$EXPECTED_MODEL_REVISION" "$OUTPUT_DIR/driver.log"
grep -aFx "DATA_REVISION=$EXPECTED_DATA_REVISION" "$OUTPUT_DIR/driver.log"
grep -aFx "TARGET_STEPS=$TARGET_STEPS" "$OUTPUT_DIR/driver.log"
test "$(grep -aEc '^NVIDIA_DRIVER_VERSION=[^[:space:]]+$' "$OUTPUT_DIR/driver.log")" -eq 1
test "$(grep -aFxc 'GPU_COMPUTE_CAPABILITY=8.0' "$OUTPUT_DIR/driver.log")" -eq 1
test "$(grep -aEc '^FLASH_ATTN_4_VERSION_BEFORE_WORKAROUND=.+$' "$OUTPUT_DIR/driver.log")" -eq 1
test "$(grep -aEc '^FLASH_ATTN_4_METADATA_ACTION=(disabled|already-absent)$' "$OUTPUT_DIR/driver.log")" -eq 1
test "$(grep -aFxc 'transformer_engine.pytorch import-ok' "$OUTPUT_DIR/driver.log")" -eq 1
test "$(grep -aFxc 'RLINF_PYTHON_EXIT_CODE=0' "$OUTPUT_DIR/driver.log")" -eq 1
test "$(grep -aFxc 'RLINF_TEE_EXIT_CODE=0' "$OUTPUT_DIR/driver.log")" -eq 1
test "$(grep -aFxc 'RLINF_VERIFICATION_STATUS=passed' "$OUTPUT_DIR/driver.log")" -eq 1
test "$(grep -aFxc 'STARTUP_EXIT_CODE=0' "$OUTPUT_DIR/driver.log")" -eq 1
grep -aFq "$COMPLETION_MARKER" "$OUTPUT_DIR/driver.log"

find "$OUTPUT_DIR" -maxdepth 4 -type f -printf '%P\t%s\n' | sort

所有命令退出码为 0 时,说明训练日志、解析后配置、TensorBoard 和目标 step 的 checkpoint 在 Worker 结束后仍可读取。输出的查看方法见获取训练结果

运行更多训练步数

完成十步任务后,可以用相同配置运行更多 step:

  1. 创建新任务,并使用新的 RUN_ID
  2. TARGET_STEPS 设置为大于 10 的正整数,例如 100
  3. 保持模型、数据、batch、序列长度和其他训练参数不变;
  4. 在 Step 8 的检查命令中填写相同的 RUN_IDTARGET_STEPS

启动命令会把 TARGET_STEPS 传给 runner.max_steps,同时提高 runner.max_epochs 上限,确保训练由明确的目标 step 停止。命令随后检查对应的完成标志和 global_step_<TARGET_STEPS> checkpoint。上游配置每 50 step 保存一次 checkpoint,并在目标 step 再保存一次;因此,目标值不是 50 的倍数时也会生成最终 checkpoint。

TARGET_STEPS 只接受正整数,不支持使用 -1 按 epoch 决定结束位置。明确的目标 step 便于预估运行范围,并让启动命令能够核对完成标志和最终 checkpoint。

增加 step 会延长资源占用时间,也会提高 Spot 任务被中断的可能性。本教程不包含 checkpoint 恢复流程;需要恢复训练时,请参见验证 checkpoint 保存与恢复。达到目标 step 只表示本次配置执行结束,模型效果和收敛情况仍需使用独立验证集评估。

停止、清理或重新运行任务

启动命令检测到输入、GPU、Ray、训练完成标志或持久化输出异常时会返回非 0。保留该次输出目录用于排查,再使用新的 RUN_ID 创建任务。通用问题请按训练故障处理检查任务事件、日志和存储挂载。

处理本场景特有故障

以下问题与本场景的共享存储路径、单 GPU 共置或 A100 兼容处理直接相关。

/mnt/rlinf-reproduction 只读或找不到代码

检查训练任务的存储配置,确认共享存储的系统挂载路径是 /mnt/rlinf-reproduction。如果 AICoder 使用不同的容器路径,请确认两个路径指向同一个存储卷和目录,并同时对齐代码、模型、数据、输出和 TensorBoard 路径。

Transformer Engine 报告 FlashAttention 或 CUTLASS 错误

driver.log 中检查 GPU_COMPUTE_CAPABILITY=8.0FLASH_ATTN_4_METADATA_ACTION=disabledtransformer_engine.pytorch import-ok。缺少任一结果时,停止任务,再核对 GPU、完整镜像 tag 和 Step 4 的启动命令。升级 NVIDIA 驱动不能替代本文的镜像兼容处理。

任务显示运行成功,但没有达到目标 step

确认任务使用的是 Step 4 的完整启动命令,并检查日志中的 TARGET_STEPS 与任务设置一致。RLINF_VERIFICATION_STATUS=passedSTARTUP_EXIT_CODE=0、目标 step 完成标志、TensorBoard 事件和最终 checkpoint 应同时存在。

actor、rollout 和 reward 共置时出现 OOM

先确认任务使用 1 张 A100-SXM4-80GB 和本文给出的原始配置。如果仍然出现 OOM,请保留日志并停止本次复现;降低 batch、序列长度或 rollout.gpu_memory_utilization 后,应使用新的 RUN_ID,并把结果作为另一种配置单独检查。

适用范围

  • 同一个 AReaL-boba 文件用于 train 和 val,因此验证奖励用于检查训练链路。评估泛化效果时,请准备独立验证集。
  • 十个 step 是本教程的验收基线。增加 TARGET_STEPS 只会改变停止位置;模型效果、收敛情况和长时间稳定性需要单独训练和评测。
  • FlashAttention 兼容处理适用于本文的镜像和 A100-SXM4-80GB。更换镜像或 GPU 后,请重新检查 Transformer Engine 的导入结果。
  • 本次运行确认 checkpoint 已写入共享存储。恢复训练时,请使用新的 RUN_ID 和显式 runner.resume_dir,并按验证 checkpoint 保存与恢复检查恢复的 global step。
  • 多 Worker 任务需要选择 AIStudio 托管 Ray,并重新配置组件和 rank 放置。