在 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_10checkpoint 保留在共享存储中。
场景信息
创建任务时使用以下设置。配置文件名中的 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,且不修改代码。
- 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 start或ray job submit。 - Ray:AIStudio 管理 Ray Head 和 Worker,适用于跨多个 Worker 放置组件的任务。使用这种模式时,还要检查各 Worker 的共享路径、驱动、网络和 rank 放置。
本教程使用以下拓扑:
AIStudio 单机训练任务
└── Worker 0 / GPU 0
└── RLinf 本地 Ray 运行时
├── RLinf Driver
├── RolloutGroup / SGLang
├── RewardGroup / math reward
└── ActorGroup / MegatronWorker 从共享存储读取固定代码、模型和数据,并把日志、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>。
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 或开发机中执行:
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 clone、git fetch 或网络下载。代码、数据和输出目录的一般准备方法见准备代码、数据与输出目录。
在共享存储中下载固定版本的模型和数据
以下命令使用 ModelScope CLI 下载模型,并使用 hf CLI 下载 AReaL-boba 数据。请在挂载共享存储的 AICoder 或开发机中执行。工具准备方法见AICoder 下载模型和使用 HuggingFace 镜像站下载模型和数据集。
模型路径末尾是 ModelScope 仓库的固定 revision。请按原样保留,使后续任务继续读取同一版本的模型文件。
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
)检查模型文件和数据字段
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=128、group_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.chdir | false | 保持 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 或开发机中执行:
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 的导入结果。
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_GRPORLinf 在 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 组件
任务进入运行中后,确认:
- 日志中的完整
IMAGE_TAG、RLinf commit、模型 revision、数据 revision 和TARGET_STEPS=10与任务记录一致。 - 日志只包含一条
NVIDIA_DRIVER_VERSION=<version>;可见 GPU 数量为 1,计算能力为8.0,CUDA 检查和 Transformer Engine 导入均通过。 - RLinf 启动包含 1 个节点和 1 个 accelerator 的本地 Ray,并成功创建 RolloutGroup、RewardGroup 和 ActorGroup。
- 模型、tokenizer、train 和 val 路径均指向共享存储中已校验的固定版本,日志随后进入生成或训练阶段。
Step 7 确认十个 step 和 GRPO 更新完成
任务运行期间,在日志和 TensorBoard 中依次确认:
- 模型和 tokenizer 从固定的本地 revision 目录加载,没有回退到网络;
- AReaL-boba 数据成功解析
prompt和solutions; - 日志已进入生成、数学奖励、GRPO 优势计算和 actor optimizer step;
- TensorBoard 中存在有限的
rollout/奖励类标量和actor/training/损失、梯度或学习率标量; - 日志出现
Step limit given by max_steps=10 reached. Stopping run; - 末尾依次出现
RLINF_PYTHON_EXIT_CODE=0、RLINF_TEE_EXIT_CODE=0、RLINF_VERIFICATION_STATUS=passed和STARTUP_EXIT_CODE=0,任务最终显示运行成功。
以十步完成标志、TensorBoard 事件、最终 checkpoint 和 RLINF_VERIFICATION_STATUS=passed 作为验收结果,不要只依据页面状态或训练进程退出码。
Step 8 Worker 结束后检查共享存储中的输出
任务结束后,在挂载同一共享存储的 AICoder 或开发机中执行。先把示例中的 RUN_ID、TARGET_STEPS 和 <tenant-id> 替换为本次任务的实际值。
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:
- 创建新任务,并使用新的
RUN_ID; - 把
TARGET_STEPS设置为大于10的正整数,例如100; - 保持模型、数据、batch、序列长度和其他训练参数不变;
- 在 Step 8 的检查命令中填写相同的
RUN_ID和TARGET_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.0、FLASH_ATTN_4_METADATA_ACTION=disabled 和 transformer_engine.pytorch import-ok。缺少任一结果时,停止任务,再核对 GPU、完整镜像 tag 和 Step 4 的启动命令。升级 NVIDIA 驱动不能替代本文的镜像兼容处理。
任务显示运行成功,但没有达到目标 step
确认任务使用的是 Step 4 的完整启动命令,并检查日志中的 TARGET_STEPS 与任务设置一致。RLINF_VERIFICATION_STATUS=passed、STARTUP_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 放置。