diff --git a/.claude/settings.local.json b/.claude/settings.local.json new file mode 100644 index 0000000000000000000000000000000000000000..9c43fcd75fe327827b78f80ae9078f0ba1fbf48e --- /dev/null +++ b/.claude/settings.local.json @@ -0,0 +1,21 @@ +{ + "permissions": { + "allow": [ + "Bash(env)", + "Bash(curl -s -o /dev/null -w \"HTTP %{http_code} time=%{time_total}s\\\\n\" --max-time 5 http://nexus.sii.shaipower.online/repository/pypi/simple/)", + "Bash(curl -s -o /dev/null -w \"HTTPS %{http_code} time=%{time_total}s\\\\n\" --max-time 5 https://nexus.sii.shaipower.online/repository/pypi/simple/)", + "Bash(curl -s -o /dev/null -w \"HTTP %{http_code} time=%{time_total}s\\\\n\" --max-time 15 http://nexus.sii.shaipower.online/repository/pypi/simple/)", + "Bash(python3 -c \"import sys, json; d=json.loads\\(sys.stdin.read\\(\\)\\); print\\(json.dumps\\(d, indent=2, ensure_ascii=False\\)[:3000]\\)\")", + "Bash(python3 -c ' *)", + "Bash(nvidia-smi)", + "Bash(python3 *)", + "Read(//root/miniconda3/envs/**)", + "Bash(conda env *)", + "Bash(/root/miniconda3/envs/verl/bin/python *)", + "Bash(/root/miniconda3/envs/verl/bin/pip list *)", + "Bash(chmod +x *)", + "Bash(/root/miniconda3/envs/verl/bin/pip show *)", + "Bash(grep -E \"\\\\.\\(py|sh\\)$\")" + ] + } +} diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..e70c9e1952c7d1d87824475240e2a67a0ab28d06 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +verl/core.778099 filter=lfs diff=lfs merge=lfs -text +verl/verl/trainer/ppo/__pycache__/core_algos.cpython-312.pyc filter=lfs diff=lfs merge=lfs -text diff --git a/.sii/ide-port b/.sii/ide-port new file mode 100644 index 0000000000000000000000000000000000000000..5557dd06083122cbb07116debebce0c7589fb302 --- /dev/null +++ b/.sii/ide-port @@ -0,0 +1 @@ +30003 \ No newline at end of file diff --git a/0519_sft_test/data/convert_to_parquet.py b/0519_sft_test/data/convert_to_parquet.py new file mode 100644 index 0000000000000000000000000000000000000000..76b27120863e1c410bebe27491423625a83b5288 --- /dev/null +++ b/0519_sft_test/data/convert_to_parquet.py @@ -0,0 +1,154 @@ +# Convert ind_v4 SFT JSONL into a verl-compatible parquet for MultiTurnSFTDataset. +# +# Source format (per line): +# { +# "messages": [ +# {"role": "user", "content": "\n..."}, +# {"role": "assistant", "content": "", "tool_calls": [{"id": "...", "type": "function", +# "function": {"name": "...", +# "arguments": "{...}"}}]} +# ], +# "images": ["/abs/path/to/img.png"], +# "tools": [{"type": "function", "function": {...}}], +# "metadata": {...} +# } +# +# Target parquet columns (verl MultiTurnSFTDataset): +# - messages: list[dict] +# - images: list[{"bytes": }] (loadable by qwen_vl_utils.fetch_image) +# - tools: list[dict] + +import argparse +import json +import os +import random +from pathlib import Path + +import pandas as pd + + +def load_jsonl(path): + rows = [] + with open(path, "r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if line: + rows.append(json.loads(line)) + return rows + + +def read_image_bytes(path: str) -> bytes: + with open(path, "rb") as f: + return f.read() + + +def normalize_messages(messages): + out = [] + for m in messages: + nm = {"role": m["role"], "content": m.get("content", "") or ""} + if "tool_calls" in m and m["tool_calls"]: + tcs = [] + for tc in m["tool_calls"]: + fn = tc.get("function", {}) + args = fn.get("arguments", "") + if not isinstance(args, str): + args = json.dumps(args, ensure_ascii=False) + tcs.append( + { + "id": tc.get("id", ""), + "type": tc.get("type", "function"), + "function": { + "name": fn.get("name", ""), + "arguments": args, + }, + } + ) + nm["tool_calls"] = tcs + out.append(nm) + return out + + +def convert(rows, image_root: str | None, max_samples: int = -1): + if max_samples > 0: + rows = rows[:max_samples] + out = [] + for row in rows: + messages = normalize_messages(row["messages"]) + + images = [] + for img_path in row.get("images", []) or []: + p = img_path + if image_root and not os.path.isabs(p): + p = os.path.join(image_root, p) + images.append({"bytes": read_image_bytes(p)}) + + tools = row.get("tools", []) or [] + + out.append( + { + "messages": messages, + "images": images, + "tools": tools, + } + ) + return out + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--src", + default="/inspire/qb-ilm2/project/26summer-camp-21/26210880/0518_dataBuilder_liubw" + "/in_domain/output/ind_v4/sft_rl/sft_train.jsonl", + help="Path to source SFT JSONL.", + ) + parser.add_argument( + "--image_root", + default=None, + help="Root for relative image paths. JSONL already has absolute paths, leave None.", + ) + parser.add_argument( + "--out_dir", + default=str(Path(__file__).resolve().parent), + help="Directory to write {train,val}.parquet.", + ) + parser.add_argument("--val_ratio", type=float, default=0.02) + parser.add_argument("--max_samples", type=int, default=-1) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + rows = load_jsonl(args.src) + print(f"loaded {len(rows)} rows from {args.src}") + + rng = random.Random(args.seed) + rng.shuffle(rows) + + if args.max_samples > 0: + rows = rows[: args.max_samples] + print(f"truncated to {len(rows)} samples") + + n_val = max(1, int(len(rows) * args.val_ratio)) if args.val_ratio > 0 else 0 + val_rows = rows[:n_val] + train_rows = rows[n_val:] + print(f"split: train={len(train_rows)}, val={len(val_rows)}") + + train_records = convert(train_rows, args.image_root) + val_records = convert(val_rows, args.image_root) if val_rows else [] + + train_df = pd.DataFrame(train_records) + train_path = out_dir / "train.parquet" + train_df.to_parquet(train_path, index=False) + print(f"wrote {train_path} ({len(train_df)} rows)") + + if val_records: + val_df = pd.DataFrame(val_records) + val_path = out_dir / "val.parquet" + val_df.to_parquet(val_path, index=False) + print(f"wrote {val_path} ({len(val_df)} rows)") + + +if __name__ == "__main__": + main() diff --git a/0519_sft_test/data/train.parquet b/0519_sft_test/data/train.parquet new file mode 100644 index 0000000000000000000000000000000000000000..76aca3c3ea9c8204f9d87e2ae81fcea6133d8c57 --- /dev/null +++ b/0519_sft_test/data/train.parquet @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e61cc170fa7c579bb897cfca7ec2058ead67bfbab7e3c7830b80211958737d4c +size 131141435 diff --git a/0519_sft_test/data/val.parquet b/0519_sft_test/data/val.parquet new file mode 100644 index 0000000000000000000000000000000000000000..9eff1c196300eb9eaa7a9b2afd81c745a5b439ae --- /dev/null +++ b/0519_sft_test/data/val.parquet @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7741860c3e5aadef4589d9306a95495c232b91798e4f1cc3ddcde799b03b48e7 +size 2522091 diff --git a/0519_sft_test/run_sft_qwen3vl_8b.sh b/0519_sft_test/run_sft_qwen3vl_8b.sh new file mode 100644 index 0000000000000000000000000000000000000000..8da77858781fdc5589068e64bc0e7e4c3f14b81b --- /dev/null +++ b/0519_sft_test/run_sft_qwen3vl_8b.sh @@ -0,0 +1,75 @@ +#!/usr/bin/env bash +# SFT for Qwen3-VL-8B-Thinking on the ind_v4 dataset (verl FSDP engine). +# +# Usage: +# bash run_sft_qwen3vl_8b.sh # default 8 GPUs +# NUM_TRAINERS=4 bash run_sft_qwen3vl_8b.sh # override GPU count +# SP_SIZE=2 NUM_TRAINERS=8 bash run_sft_qwen3vl_8b.sh +set -xeuo pipefail + +WORKDIR=$(cd "$(dirname "$0")" && pwd) + +MODEL_PATH=${MODEL_PATH:-/inspire/qb-ilm2/project/26summer-camp-21/26210880/0518_test_liubw/models/Qwen3-VL-8B-Thinking} +TRAIN_FILES=${TRAIN_FILES:-${WORKDIR}/data/train.parquet} +VAL_FILES=${VAL_FILES:-${WORKDIR}/data/val.parquet} + +NUM_TRAINERS=${NUM_TRAINERS:-8} +SP_SIZE=${SP_SIZE:-2} # ulysses sequence parallel +FSDP_SIZE=${FSDP_SIZE:--1} +FSDP_STRATEGY=${FSDP_STRATEGY:-fsdp2} + +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-32} +MICRO_BATCH_SIZE_PER_GPU=${MICRO_BATCH_SIZE_PER_GPU:-1} +MAX_LENGTH=${MAX_LENGTH:-8192} +MAX_TOKEN_LEN_PER_GPU=${MAX_TOKEN_LEN_PER_GPU:-16384} +TOTAL_EPOCHS=${TOTAL_EPOCHS:-3} +LR=${LR:-1e-5} + +PROJECT_NAME=${PROJECT_NAME:-ind_v4_sft} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-qwen3vl-8b-thinking-fsdp${FSDP_STRATEGY}-sp${SP_SIZE}} +SAVE_DIR=${SAVE_DIR:-${WORKDIR}/checkpoints/${PROJECT_NAME}/${EXPERIMENT_NAME}} +LOGGER=${LOGGER:-"['console']"} # set to "['console','wandb']" once wandb is logged in +RESUME_MODE=${RESUME_MODE:-auto} + +mkdir -p "${SAVE_DIR}" + +cd /inspire/qb-ilm2/project/26summer-camp-21/26210880/0519_verl/verl + +torchrun --standalone --nnodes=1 --nproc-per-node="${NUM_TRAINERS}" \ + -m verl.trainer.sft_trainer \ + data.train_files="${TRAIN_FILES}" \ + data.val_files="${VAL_FILES}" \ + data.train_batch_size="${TRAIN_BATCH_SIZE}" \ + data.micro_batch_size_per_gpu="${MICRO_BATCH_SIZE_PER_GPU}" \ + data.max_length="${MAX_LENGTH}" \ + data.max_token_len_per_gpu="${MAX_TOKEN_LEN_PER_GPU}" \ + data.use_dynamic_bsz=True \ + data.pad_mode=no_padding \ + data.truncation=right \ + data.messages_key=messages \ + data.tools_key=tools \ + data.ignore_input_ids_mismatch=True \ + model.path="${MODEL_PATH}" \ + model.use_remove_padding=True \ + engine=fsdp \ + optim=fsdp \ + optim.lr="${LR}" \ + optim.lr_warmup_steps_ratio=0.03 \ + optim.weight_decay=0.1 \ + optim.betas="[0.9,0.95]" \ + optim.clip_grad=1.0 \ + optim.min_lr_ratio=0.1 \ + optim.warmup_style=cosine \ + engine.ulysses_sequence_parallel_size="${SP_SIZE}" \ + engine.strategy="${FSDP_STRATEGY}" \ + engine.fsdp_size="${FSDP_SIZE}" \ + trainer.default_local_dir="${SAVE_DIR}" \ + trainer.project_name="${PROJECT_NAME}" \ + trainer.experiment_name="${EXPERIMENT_NAME}" \ + trainer.total_epochs="${TOTAL_EPOCHS}" \ + trainer.logger="${LOGGER}" \ + trainer.save_freq=500 \ + trainer.test_freq=200 \ + trainer.max_ckpt_to_keep=3 \ + trainer.resume_mode="${RESUME_MODE}" \ + checkpoint.save_contents=[model,optimizer,extra] diff --git a/0520_rej_test/README.md b/0520_rej_test/README.md new file mode 100644 index 0000000000000000000000000000000000000000..443ce289fc9003ebe1bdf734b6f850c1b3dd572b --- /dev/null +++ b/0520_rej_test/README.md @@ -0,0 +1,69 @@ +# 0520_rej_test — Offline RAFT/RFT framework + +verl 主仓库**没有现成的 reward-based 拒绝采样训练循环**： +- `algorithm.filter_groups` 只有 dataclass，没有 trainer 消费 +- `algorithm.rollout_correction` 是 IS-based 拒绝，依据是重要性比值，不是 reward + +这里搭的是经典 **RAFT (Reward-rAnked Fine-Tuning)** —— 完全离线、字面意义上的拒绝采样： +对当前模型采样 N 条 → reward 打分 → 取 top-K → 喂给 SFT。 + +``` +rl_train.jsonl + ↓ build_prompts.py +prompts.parquet + ↓ Stage A: verl.trainer.main_generation_server (vLLM, n=N) +gen.parquet + ↓ Stage B: reward/score.py + reward/reward_fn.py (STUB) +scored.parquet + ↓ Stage C: select/topk_to_sft.py +raft_train.parquet + ↓ Stage D: scripts/run_sft.sh -> 0519_sft_test/run_sft_qwen3vl_8b.sh +checkpoints/ +``` + +## 现状 +- 框架完整可执行； +- **`reward/reward_fn.py` 是 stub**，固定 `compute_reward(prompt, response, metadata, tools, ...)` 签名，永远返回 `0.0`，等后续实现； +- 不动 `0519_sft_test/` 任何文件，Stage D 通过环境变量复用其 SFT 脚本。 + +## 使用 + +整体跑： +```bash +bash run_pipeline.sh +``` + +只跑某一段： +```bash +STAGE=A_pre bash run_pipeline.sh # 只生成 prompts.parquet +STAGE=A bash run_pipeline.sh # 跑 vLLM 采样 +STAGE=B bash run_pipeline.sh # 打分 +STAGE=C bash run_pipeline.sh # 选 top-K + 转 SFT 格式 +STAGE=D bash run_pipeline.sh # SFT +``` + +Smoke 试跑（4 条 prompt × 2 个 sample）： +```bash +MAX_SAMPLES=4 N_SAMPLES=2 RESPONSE_LENGTH=512 STAGE=A_pre bash run_pipeline.sh +MAX_SAMPLES=4 N_SAMPLES=2 RESPONSE_LENGTH=512 STAGE=A bash run_pipeline.sh +STAGE=B bash run_pipeline.sh +TOP_K=1 STAGE=C bash run_pipeline.sh +``` + +## 关键依赖 + +| 用途 | 路径 | +|---|---| +| Stage A 入口 | `verl/trainer/main_generation_server.py` | +| 多模态 chat template | `verl/workers/rollout/schemas.py` | +| SFT 目标 schema | `verl/utils/dataset/multiturn_sft_dataset.py` | +| SFT 启动脚本（被复用） | `0519_sft_test/run_sft_qwen3vl_8b.sh` | +| 模型 | `0518_test_liubw/models/Qwen3-VL-8B-Thinking` | +| RL 原始数据 | `0518_dataBuilder_liubw/.../sft_rl/rl_train.jsonl` | + +## 设计要点 + +- **图像如何进 vLLM？** verl 的 `main_generation_server` 直接把 `prompt` 列里的 OpenAI 格式 messages 转给 vLLM 的 chat completions 接口 —— 我们把每张图编成 `data:image/png;base64,...` URL 塞进 multimodal content block。Stage C 再从 `image_paths` 重新读字节进 SFT parquet。 +- **tools 怎么走？** `main_generation_server` 不读 `tools` 列；`tools` 的 JSON schema 当作文本前缀拼进 user turn。模型是 Qwen3-VL-Thinking，训练时见过 `{...}` 输出格式，会自然产出。 +- **怎么把模型输出还原成结构化 tool_calls？** Stage C 用正则提 `` 块 → `json.loads` → 构造 OpenAI tool_call dict（`arguments` 强制成字符串，与 `0519_sft_test/data/convert_to_parquet.py` 一致）。解析失败的样本直接丢。 +- **SFT schema 完全和 0519_sft_test 一致**，所以可以直接复用其 dataset 加载 + chat template 验证路径。 diff --git a/0520_rej_test/data/build_prompts.py b/0520_rej_test/data/build_prompts.py new file mode 100644 index 0000000000000000000000000000000000000000..dae591abd6eb058171631195f19b71bbc6d56425 --- /dev/null +++ b/0520_rej_test/data/build_prompts.py @@ -0,0 +1,135 @@ +"""Build prompts.parquet for Stage A (verl.trainer.main_generation_server). + +Source: ind_v4 rl_train.jsonl with shape + {messages, images, tools, metadata} + +main_generation_server expects a single `prompt` column whose value is an +OpenAI-style chat-completions `messages` list. It does NOT read `tools` or +`images` columns directly — both must be baked into `prompt`. + +Encoding: + - `images[i]` is read from disk and embedded as a base64 data-URL inside + a multimodal content block on the user turn. + - `tools` (the JSON spec) is rendered as a textual preamble prepended to + the user turn so that even without server-side tool injection the model + sees the tool name + schema. The model is Qwen3-VL-Thinking, which is + trained to emit `{...}` blocks — Stage C will + parse them back out. + +Extra columns kept on the parquet so downstream stages can use them: + - `metadata`: original metadata dict (test_points, expression_numpy, ...) + - `image_paths`: original image paths (so Stage C can re-read bytes + without going through the data-URL). + - `tools_orig`: original tools spec (for Stage C's SFT reconstruction). +""" + +import argparse +import base64 +import json +import mimetypes +import os +import random +from pathlib import Path + +import pandas as pd + + +def load_jsonl(path): + rows = [] + with open(path, "r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if line: + rows.append(json.loads(line)) + return rows + + +def img_to_data_url(path: str) -> str: + mime, _ = mimetypes.guess_type(path) + if mime is None: + mime = "image/png" + with open(path, "rb") as f: + b64 = base64.b64encode(f.read()).decode("ascii") + return f"data:{mime};base64,{b64}" + + +def render_tools_preamble(tools: list[dict]) -> str: + if not tools: + return "" + lines = ["You have access to the following tools. Reply via a single " + "{...} block whose JSON contains " + "`name` and `arguments`."] + for t in tools: + lines.append(json.dumps(t, ensure_ascii=False)) + return "\n".join(lines) + "\n\n" + + +def build_prompt(row: dict) -> list[dict]: + """Construct OpenAI multimodal `messages` from one raw RL sample.""" + user_text = row["messages"][0]["content"] + if user_text.startswith("\n"): + user_text = user_text[len("\n") :] + elif user_text.startswith(""): + user_text = user_text[len("") :] + + preamble = render_tools_preamble(row.get("tools") or []) + text_block = preamble + user_text + + content = [] + for img_path in row.get("images") or []: + content.append( + {"type": "image_url", + "image_url": {"url": img_to_data_url(img_path)}} + ) + content.append({"type": "text", "text": text_block}) + + return [{"role": "user", "content": content}] + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--src", + default="/inspire/qb-ilm2/project/26summer-camp-21/26210880" + "/0518_dataBuilder_liubw/in_domain/output/ind_v4/sft_rl/rl_train.jsonl", + ) + parser.add_argument( + "--out", + default=str(Path(__file__).resolve().parent / "prompts.parquet"), + ) + parser.add_argument("--max_samples", type=int, default=-1) + parser.add_argument("--shuffle", action="store_true") + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + + rows = load_jsonl(args.src) + print(f"loaded {len(rows)} rows from {args.src}") + + if args.shuffle: + rng = random.Random(args.seed) + rng.shuffle(rows) + + if args.max_samples > 0: + rows = rows[: args.max_samples] + print(f"truncated to {len(rows)}") + + records = [] + for row in rows: + records.append( + { + "prompt": build_prompt(row), + "metadata": row.get("metadata", {}), + "image_paths": row.get("images") or [], + "tools_orig": row.get("tools") or [], + } + ) + + out_path = Path(args.out) + out_path.parent.mkdir(parents=True, exist_ok=True) + df = pd.DataFrame(records) + df.to_parquet(out_path, index=False) + print(f"wrote {out_path} ({len(df)} rows)") + + +if __name__ == "__main__": + main() diff --git a/0520_rej_test/reward/__pycache__/reward_fn.cpython-312.pyc b/0520_rej_test/reward/__pycache__/reward_fn.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..adbd9208072678165cf15eac8ac08ecf34ee71f9 Binary files /dev/null and b/0520_rej_test/reward/__pycache__/reward_fn.cpython-312.pyc differ diff --git a/0520_rej_test/reward/reward_fn.py b/0520_rej_test/reward/reward_fn.py new file mode 100644 index 0000000000000000000000000000000000000000..9e41204725c90f417366f2baca5e64d1dcab8b35 --- /dev/null +++ b/0520_rej_test/reward/reward_fn.py @@ -0,0 +1,34 @@ +"""Stage B reward function. + +This is a STUB. The signature is fixed; downstream `score.py` calls +`compute_reward(...)` row-wise by keyword. Implement the body later. + +Inputs (kwargs): + prompt: list[dict] # OpenAI chat-completions messages used at generation + response: str # one model completion (raw text) + metadata: dict # original sample metadata (test_points, expression_numpy, ...) + tools: list[dict] # original `tools` JSON spec + index: int # sample index (for logging) + sample: int # which of the N samples (0..N-1) + +Returns: + float in [0, 1] -- higher is better +""" + +from __future__ import annotations + + +def compute_reward( + *, + prompt: list[dict], + response: str, + metadata: dict, + tools: list[dict], + index: int = 0, + sample: int = 0, +) -> float: + # TODO: parse {"name": "submit_expression", "arguments": "..."} + # eval expression at metadata["test_points"], compute MSE-based reward. + # Return 0.0 if parse fails or eval errors out. + del prompt, response, metadata, tools, index, sample + return 0.0 diff --git a/0520_rej_test/reward/score.py b/0520_rej_test/reward/score.py new file mode 100644 index 0000000000000000000000000000000000000000..fc4893c3ea0481a9f9811eef80a5f9e285231547 --- /dev/null +++ b/0520_rej_test/reward/score.py @@ -0,0 +1,80 @@ +"""Stage B driver: read gen.parquet, score each (prompt, response) pair via +reward_fn.compute_reward, write scored.parquet with a new `scores` column +(parallel to `responses`). + +Reward fn lives in reward/reward_fn.py and can be swapped via --reward_fn. +""" + +import argparse +import importlib.util +import sys +from pathlib import Path + +import pandas as pd + + +def load_reward_fn(path: str, name: str = "compute_reward"): + spec = importlib.util.spec_from_file_location("reward_fn_module", path) + mod = importlib.util.module_from_spec(spec) + sys.modules["reward_fn_module"] = mod + assert spec.loader is not None + spec.loader.exec_module(mod) + fn = getattr(mod, name) + return fn + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--gen", + default=str(Path(__file__).resolve().parent.parent / "data" / "gen.parquet"), + ) + parser.add_argument( + "--out", + default=str(Path(__file__).resolve().parent.parent / "data" / "scored.parquet"), + ) + parser.add_argument( + "--reward_fn", + default=str(Path(__file__).resolve().parent / "reward_fn.py"), + ) + parser.add_argument("--reward_name", default="compute_reward") + args = parser.parse_args() + + df = pd.read_parquet(args.gen) + print(f"loaded {len(df)} rows from {args.gen}; columns: {list(df.columns)}") + + if "responses" not in df.columns: + raise SystemExit("expected `responses` column from Stage A; not found.") + + compute_reward = load_reward_fn(args.reward_fn, args.reward_name) + + all_scores = [] + for idx, row in df.iterrows(): + prompt = row["prompt"] + responses = list(row["responses"]) + metadata = row.get("metadata", {}) or {} + tools = list(row.get("tools_orig", []) or []) + scores = [] + for s_idx, resp in enumerate(responses): + score = compute_reward( + prompt=prompt, + response=resp, + metadata=metadata, + tools=tools, + index=int(idx), + sample=s_idx, + ) + scores.append(float(score)) + all_scores.append(scores) + + df = df.copy() + df["scores"] = all_scores + + out_path = Path(args.out) + out_path.parent.mkdir(parents=True, exist_ok=True) + df.to_parquet(out_path, index=False) + print(f"wrote {out_path} ({len(df)} rows)") + + +if __name__ == "__main__": + main() diff --git a/0520_rej_test/run_pipeline.sh b/0520_rej_test/run_pipeline.sh new file mode 100644 index 0000000000000000000000000000000000000000..8099c7ce02460233a96ce85a997835f9fd9d7f14 --- /dev/null +++ b/0520_rej_test/run_pipeline.sh @@ -0,0 +1,43 @@ +#!/usr/bin/env bash +# End-to-end RAFT pipeline: build prompts -> generate -> score -> top-K -> SFT. +# Each stage is overrideable via env vars and idempotent. +set -xeuo pipefail + +WORKDIR=$(cd "$(dirname "$0")" && pwd) +PY=${PY:-python} + +MAX_SAMPLES=${MAX_SAMPLES:--1} +N_SAMPLES=${N_SAMPLES:-8} +TOP_K=${TOP_K:-1} + +STAGE=${STAGE:-all} # all | A | B | C | D + +if [[ "${STAGE}" == "all" || "${STAGE}" == "A_pre" || "${STAGE}" == "A" ]]; then + "${PY}" "${WORKDIR}/data/build_prompts.py" \ + --out "${WORKDIR}/data/prompts.parquet" \ + --max_samples "${MAX_SAMPLES}" +fi + +if [[ "${STAGE}" == "all" || "${STAGE}" == "A" ]]; then + N_SAMPLES="${N_SAMPLES}" \ + PROMPTS="${WORKDIR}/data/prompts.parquet" \ + OUT="${WORKDIR}/data/gen.parquet" \ + bash "${WORKDIR}/scripts/run_generate.sh" +fi + +if [[ "${STAGE}" == "all" || "${STAGE}" == "B" ]]; then + "${PY}" "${WORKDIR}/reward/score.py" \ + --gen "${WORKDIR}/data/gen.parquet" \ + --out "${WORKDIR}/data/scored.parquet" +fi + +if [[ "${STAGE}" == "all" || "${STAGE}" == "C" ]]; then + "${PY}" "${WORKDIR}/select/topk_to_sft.py" \ + --scored "${WORKDIR}/data/scored.parquet" \ + --out "${WORKDIR}/data/raft_train.parquet" \ + --k "${TOP_K}" +fi + +if [[ "${STAGE}" == "all" || "${STAGE}" == "D" ]]; then + bash "${WORKDIR}/scripts/run_sft.sh" +fi diff --git a/0520_rej_test/scripts/run_generate.sh b/0520_rej_test/scripts/run_generate.sh new file mode 100644 index 0000000000000000000000000000000000000000..f6fdecc91383a52a29fad960340a9bc71eb4b286 --- /dev/null +++ b/0520_rej_test/scripts/run_generate.sh @@ -0,0 +1,48 @@ +#!/usr/bin/env bash +# Stage A: rollout N completions per prompt with vLLM, via verl's +# main_generation_server. +# +# Usage: +# bash run_generate.sh # full prompts.parquet +# PROMPTS=/path/to/sub.parquet OUT=/tmp/gen.parquet \ +# N_SAMPLES=2 bash run_generate.sh # smoke +set -xeuo pipefail + +WORKDIR=$(cd "$(dirname "$0")/.." && pwd) +VERL_ROOT=/inspire/qb-ilm2/project/26summer-camp-21/26210880/0519_verl/verl + +MODEL_PATH=${MODEL_PATH:-/inspire/qb-ilm2/project/26summer-camp-21/26210880/0518_test_liubw/models/Qwen3-VL-8B-Thinking} +PROMPTS=${PROMPTS:-${WORKDIR}/data/prompts.parquet} +OUT=${OUT:-${WORKDIR}/data/gen.parquet} + +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} +ROLLOUT_TP=${ROLLOUT_TP:-2} +ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.85} + +PROMPT_LENGTH=${PROMPT_LENGTH:-4096} +RESPONSE_LENGTH=${RESPONSE_LENGTH:-4096} +N_SAMPLES=${N_SAMPLES:-8} +TEMPERATURE=${TEMPERATURE:-1.0} +TOP_P=${TOP_P:-0.95} + +mkdir -p "$(dirname "${OUT}")" + +cd "${VERL_ROOT}" + +python3 -m verl.trainer.main_generation_server \ + trainer.nnodes="${NNODES}" \ + trainer.n_gpus_per_node="${NGPUS_PER_NODE}" \ + data.train_files="${PROMPTS}" \ + data.prompt_key=prompt \ + +data.output_path="${OUT}" \ + actor_rollout_ref.model.path="${MODEL_PATH}" \ + actor_rollout_ref.model.trust_remote_code=True \ + actor_rollout_ref.rollout.name=vllm \ + actor_rollout_ref.rollout.temperature="${TEMPERATURE}" \ + actor_rollout_ref.rollout.top_p="${TOP_P}" \ + actor_rollout_ref.rollout.prompt_length="${PROMPT_LENGTH}" \ + actor_rollout_ref.rollout.response_length="${RESPONSE_LENGTH}" \ + actor_rollout_ref.rollout.tensor_model_parallel_size="${ROLLOUT_TP}" \ + actor_rollout_ref.rollout.gpu_memory_utilization="${ROLLOUT_GPU_MEM_UTIL}" \ + actor_rollout_ref.rollout.n="${N_SAMPLES}" "$@" diff --git a/0520_rej_test/scripts/run_sft.sh b/0520_rej_test/scripts/run_sft.sh new file mode 100644 index 0000000000000000000000000000000000000000..ca6fabca05b3c170a50ea167bfede182fd7389d3 --- /dev/null +++ b/0520_rej_test/scripts/run_sft.sh @@ -0,0 +1,20 @@ +#!/usr/bin/env bash +# Stage D wrapper: re-use 0519_sft_test/run_sft_qwen3vl_8b.sh against the +# RAFT-mined parquet, without modifying the original script. +set -xeuo pipefail + +WORKDIR=$(cd "$(dirname "$0")/.." && pwd) + +TRAIN_FILES=${TRAIN_FILES:-${WORKDIR}/data/raft_train.parquet} +VAL_FILES=${VAL_FILES:-${WORKDIR}/data/raft_val.parquet} + +PROJECT_NAME=${PROJECT_NAME:-ind_v4_raft} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-qwen3vl-8b-thinking-raft} +SAVE_DIR=${SAVE_DIR:-${WORKDIR}/checkpoints/${PROJECT_NAME}/${EXPERIMENT_NAME}} + +TRAIN_FILES="${TRAIN_FILES}" \ +VAL_FILES="${VAL_FILES}" \ +PROJECT_NAME="${PROJECT_NAME}" \ +EXPERIMENT_NAME="${EXPERIMENT_NAME}" \ +SAVE_DIR="${SAVE_DIR}" \ +bash /inspire/qb-ilm2/project/26summer-camp-21/26210880/0519_verl/0519_sft_test/run_sft_qwen3vl_8b.sh "$@" diff --git a/0520_rej_test/select/topk_to_sft.py b/0520_rej_test/select/topk_to_sft.py new file mode 100644 index 0000000000000000000000000000000000000000..f26af3e7408a24cd1a19aac928d973d833b86849 --- /dev/null +++ b/0520_rej_test/select/topk_to_sft.py @@ -0,0 +1,148 @@ +"""Stage C: scored.parquet -> raft_train.parquet (verl SFT schema). + +For each row pick the top-K responses by score, parse Qwen3-VL-Thinking +`{...}` blocks back into structured tool_calls, and +emit one SFT example per (row, kept response). Output schema matches +`verl/utils/dataset/multiturn_sft_dataset.py`: + - messages: [user_turn, assistant_turn_with_tool_calls] + - images: [{"bytes": }] (re-read from row["image_paths"]) + - tools: original tools spec +""" + +import argparse +import json +import os +import re +from pathlib import Path + +import pandas as pd + +TOOL_CALL_RE = re.compile(r"\s*(.*?)\s*", re.DOTALL) + + +def read_image_bytes(path: str) -> bytes: + with open(path, "rb") as f: + return f.read() + + +def parse_tool_calls(text: str) -> tuple[str, list[dict]]: + """Strip ... blocks out of `text`. + + Returns: + (residual_text_with_blocks_removed, list_of_tool_call_dicts) + + Each tool_call dict is OpenAI-style: + {"id": "...", "type": "function", + "function": {"name": "...", "arguments": ""}} + Failed JSON parses make the whole sample invalid (caller should drop). + """ + tool_calls = [] + for i, m in enumerate(TOOL_CALL_RE.finditer(text)): + raw = m.group(1).strip() + try: + obj = json.loads(raw) + except json.JSONDecodeError: + return text, [] + name = obj.get("name", "") + args = obj.get("arguments", obj.get("parameters", {})) + if not isinstance(args, str): + args = json.dumps(args, ensure_ascii=False) + tool_calls.append( + { + "id": f"call_{i}", + "type": "function", + "function": {"name": name, "arguments": args}, + } + ) + residual = TOOL_CALL_RE.sub("", text).strip() + return residual, tool_calls + + +def reconstruct_user_turn(prompt_messages: list[dict]) -> dict: + """The prompt for generation was multimodal (image_url + text). For SFT + we replace it with the conventional `\\n` Qwen format — + this matches the existing 0519_sft_test/data parquet exactly. The image + bytes are attached separately via the `images` column. + """ + user_msg = prompt_messages[0] + text_parts = [] + for blk in user_msg["content"]: + if isinstance(blk, dict) and blk.get("type") == "text": + text_parts.append(blk["text"]) + text = "\n".join(text_parts) + return {"role": "user", "content": "\n" + text} + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--scored", + default=str(Path(__file__).resolve().parent.parent / "data" / "scored.parquet"), + ) + parser.add_argument( + "--out", + default=str(Path(__file__).resolve().parent.parent / "data" / "raft_train.parquet"), + ) + parser.add_argument("--k", type=int, default=1, help="Top-K responses per prompt.") + parser.add_argument("--min_score", type=float, default=None) + parser.add_argument("--keep_above_mean", action="store_true", + help="Keep responses scoring strictly above the per-row mean.") + args = parser.parse_args() + + df = pd.read_parquet(args.scored) + print(f"loaded {len(df)} rows from {args.scored}") + + out_records = [] + n_dropped_parse = 0 + n_dropped_score = 0 + for _, row in df.iterrows(): + responses = list(row["responses"]) + scores = list(row["scores"]) + assert len(responses) == len(scores) + + ranked = sorted(zip(scores, responses), key=lambda x: x[0], reverse=True) + if args.keep_above_mean: + mean_s = sum(scores) / max(1, len(scores)) + ranked = [(s, r) for s, r in ranked if s > mean_s] + if args.min_score is not None: + ranked = [(s, r) for s, r in ranked if s >= args.min_score] + ranked = ranked[: args.k] + + if not ranked: + n_dropped_score += 1 + continue + + user_turn = reconstruct_user_turn(list(row["prompt"])) + tools = list(row.get("tools_orig", []) or []) + image_paths = list(row.get("image_paths", []) or []) + images = [{"bytes": read_image_bytes(p)} for p in image_paths] + + for score, resp in ranked: + content_text, tool_calls = parse_tool_calls(resp) + if not tool_calls: + n_dropped_parse += 1 + continue + assistant_turn = { + "role": "assistant", + "content": content_text or "", + "tool_calls": tool_calls, + } + out_records.append( + { + "messages": [user_turn, assistant_turn], + "images": images, + "tools": tools, + } + ) + + print(f"kept {len(out_records)} sft examples; " + f"dropped {n_dropped_parse} (parse), {n_dropped_score} (score filter)") + + out_path = Path(args.out) + out_path.parent.mkdir(parents=True, exist_ok=True) + pd.DataFrame(out_records).to_parquet(out_path, index=False) + print(f"wrote {out_path}") + + +if __name__ == "__main__": + main() diff --git a/data.tar.gz b/data.tar.gz new file mode 100644 index 0000000000000000000000000000000000000000..ea3cbf87910074d152f013ae134786a4b2d45b15 --- /dev/null +++ b/data.tar.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:13d6c907cdce09739e63e32f40977e608f6afe6169b85bb58938b9df2616b3f8 +size 391661196 diff --git a/pip_list.txt b/pip_list.txt new file mode 100644 index 0000000000000000000000000000000000000000..eb1954d02190f8f120a2c1aed1a0ab8966fd6341 --- /dev/null +++ b/pip_list.txt @@ -0,0 +1,206 @@ +absl-py 2.4.0 +accelerate 1.12.0 +aiohappyeyeballs 2.6.1 +aiohttp 3.13.3 +aiohttp-cors 0.8.1 +aiosignal 1.4.0 +annotated-doc 0.0.4 +annotated-types 0.7.0 +antlr4-python3-runtime 4.9.3 +anyio 4.12.1 +astor 0.8.1 +async-timeout 5.0.1 +attrs 25.4.0 +av 16.1.0 +blake3 1.0.8 +cachetools 6.2.5 +cbor2 5.8.0 +certifi 2026.1.4 +cffi 2.0.0 +charset-normalizer 3.4.4 +click 8.3.1 +cloudpickle 3.1.2 +codetiming 1.4.0 +colorful 0.5.8 +compressed-tensors 0.11.0 +cupy-cuda12x 13.6.0 +datasets 4.5.0 +depyf 0.19.0 +dill 0.4.0 +diskcache 5.6.3 +distlib 0.4.0 +distro 1.9.0 +dnspython 2.8.0 +einops 0.8.2 +email-validator 2.3.0 +exceptiongroup 1.3.1 +fastapi 0.128.0 +fastapi-cli 0.0.20 +fastapi-cloud-cli 0.11.0 +fastar 0.8.0 +fastrlock 0.8.3 +filelock 3.20.3 +flash_attn 2.8.3+cu12torch28cxx11abitrue +frozendict 2.4.7 +frozenlist 1.8.0 +fsspec 2025.10.0 +gguf 0.17.1 +gitdb 4.0.12 +GitPython 3.1.46 +google-api-core 2.29.0 +google-auth 2.47.0 +googleapis-common-protos 1.72.0 +grpcio 1.76.0 +h11 0.16.0 +hf-xet 1.2.0 +httpcore 1.0.9 +httptools 0.7.1 +httpx 0.28.1 +huggingface-hub 0.36.0 +idna 3.11 +importlib_metadata 8.7.1 +interegular 0.3.3 +Jinja2 3.1.6 +jiter 0.12.0 +jsonschema 4.26.0 +jsonschema-specifications 2025.9.1 +lark 1.2.2 +liger_kernel 0.6.4 +llguidance 0.7.30 +llvmlite 0.44.0 +lm-format-enforcer 0.11.3 +Markdown 3.10.1 +markdown-it-py 4.0.0 +MarkupSafe 3.0.3 +mathruler 0.1.0 +mdurl 0.1.2 +mistral_common 1.8.8 +mpmath 1.3.0 +msgpack 1.1.2 +msgspec 0.20.0 +multidict 6.7.1 +multiprocess 0.70.18 +networkx 3.4.2 +ninja 1.13.0 +numba 0.61.2 +numpy 2.2.6 +nvidia-cublas-cu12 12.6.4.1 +nvidia-cuda-cupti-cu12 12.6.80 +nvidia-cuda-nvrtc-cu12 12.6.77 +nvidia-cuda-runtime-cu12 12.6.77 +nvidia-cudnn-cu12 9.10.2.21 +nvidia-cufft-cu12 11.3.0.4 +nvidia-cufile-cu12 1.11.1.6 +nvidia-curand-cu12 10.3.7.77 +nvidia-cusolver-cu12 11.7.1.2 +nvidia-cusparse-cu12 12.5.4.2 +nvidia-cusparselt-cu12 0.7.1 +nvidia-nccl-cu12 2.27.3 +nvidia-nvjitlink-cu12 12.6.85 +nvidia-nvtx-cu12 12.6.77 +omegaconf 2.3.0 +openai 2.15.0 +openai-harmony 0.0.8 +opencensus 0.11.4 +opencensus-context 0.1.3 +opencv-python-headless 4.13.0.90 +opentelemetry-api 1.39.1 +opentelemetry-exporter-prometheus 0.60b1 +opentelemetry-proto 1.39.1 +opentelemetry-sdk 1.39.1 +opentelemetry-semantic-conventions 0.60b1 +orjson 3.11.5 +outlines_core 0.2.11 +packaging 25.0 +pandas 2.3.3 +partial-json-parser 0.2.1.1.post7 +peft 0.18.1 +pillow 12.1.0 +pip 25.3 +platformdirs 4.5.1 +prometheus_client 0.24.1 +prometheus-fastapi-instrumentator 7.1.0 +propcache 0.4.1 +proto-plus 1.27.0 +protobuf 6.33.4 +psutil 7.2.1 +py-cpuinfo 9.0.0 +py-spy 0.4.1 +pyarrow 23.0.0 +pyasn1 0.6.2 +pyasn1_modules 0.4.2 +pybase64 1.4.3 +pycountry 24.6.1 +pycparser 3.0 +pydantic 2.12.5 +pydantic_core 2.41.5 +pydantic-extra-types 2.11.0 +pydantic-settings 2.12.0 +Pygments 2.19.2 +pylatexenc 2.10 +python-dateutil 2.9.0.post0 +python-dotenv 1.2.1 +python-json-logger 4.0.0 +python-multipart 0.0.22 +pytz 2025.2 +pyvers 0.1.0 +PyYAML 6.0.3 +pyzmq 27.1.0 +qwen-vl-utils 0.0.14 +ray 2.53.0 +referencing 0.37.0 +regex 2026.1.15 +requests 2.32.5 +rich 14.3.1 +rich-toolkit 0.17.1 +rignore 0.7.6 +rpds-py 0.30.0 +rsa 4.9.1 +safetensors 0.7.0 +scipy 1.15.3 +sentencepiece 0.2.1 +sentry-sdk 2.50.0 +setproctitle 1.3.7 +setuptools 80.9.0 +shellingham 1.5.4 +six 1.17.0 +smart_open 7.5.0 +smmap 5.0.2 +sniffio 1.3.1 +soundfile 0.13.1 +soxr 1.0.0 +starlette 0.50.0 +sympy 1.14.0 +tensorboard 2.20.0 +tensorboard-data-server 0.7.2 +tensordict 0.10.0 +tiktoken 0.12.0 +tokenizers 0.22.2 +tomli 2.4.0 +torch 2.8.0+cu126 +torchaudio 2.8.0+cu126 +torchdata 0.11.0 +torchvision 0.23.0+cu126 +tqdm 4.67.1 +transformers 4.57.0 +triton 3.4.0 +typer 0.21.1 +typing_extensions 4.15.0 +typing-inspection 0.4.2 +tzdata 2025.3 +urllib3 2.6.3 +uvicorn 0.40.0 +uvloop 0.22.1 +virtualenv 20.36.1 +vllm 0.11.0 +wandb 0.24.0 +watchfiles 1.1.1 +websockets 16.0 +Werkzeug 3.1.5 +wheel 0.45.1 +wrapt 2.0.1 +xformers 0.0.32.post1 +xgrammar 0.1.25 +xxhash 3.6.0 +yarl 1.22.0 +zipp 3.23.0 \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..62edc1d9ede6ea1c97fd0d39eb66fca25993ea54 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,132 @@ +# Generated from pip_list.txt — exact version pinning +# Note: +# - torch / torchaudio / torchvision carry the +cu124 local tag; +# install with: pip install --index-url https://download.pytorch.org/whl/cu124 ... +# - conda / mamba / libmambapy / menuinst / libarchive-c are normally managed +# by conda, not pip. Remove them if installing into a non-conda env. +# - pip / setuptools / wheel are bootstrap packages; pin only if you intend to +# control their versions explicitly. + +absl-py==2.4.0 +annotated-doc==0.0.4 +anyio==4.13.0 +archspec==0.2.3 +asttokens==2.4.1 +astunparse==1.6.3 +attrs==24.2.0 +beautifulsoup4==4.12.3 +blessed==1.33.0 +boltons==24.0.0 +Brotli==1.1.0 +certifi==2024.8.30 +cffi==1.17.1 +chardet==5.2.0 +charset-normalizer==3.4.0 +click==8.3.3 +colorama==0.4.6 +cycler==0.12.1 +decorator==5.1.1 +distro==1.9.0 +dnspython==2.7.0 +exceptiongroup==1.2.2 +executing==2.1.0 +expecttest==0.2.1 +filelock==3.16.1 +fonttools==4.62.1 +frozendict==2.4.6 +fsspec==2024.10.0 +gpustat==1.1.1 +grpcio==1.80.0 +h11==0.16.0 +h2==4.1.0 +hf-xet==1.5.0 +hpack==4.0.0 +httpcore==1.0.9 +httpx==0.28.1 +huggingface-hub==1.14.0 +hyperframe==6.0.1 +hypothesis==6.115.5 +idna==3.10 +importlib-resources==6.4.5 +ipython==8.29.0 +jedi==0.19.1 +Jinja2==3.1.4 +jsonpatch==1.33 +jsonpointer==3.0.0 +jsonschema==4.23.0 +jsonschema-specifications==2024.10.1 +kiwisolver==1.5.0 +libarchive-c==5.1 +lief==0.14.1 +lintrunner==0.12.5 +Markdown==3.10.2 +markdown-it-py==4.2.0 +MarkupSafe==3.0.2 +matplotlib==3.10.9 +matplotlib-inline==0.1.7 +mdurl==0.1.2 +modelscope==1.35.2 +more-itertools==10.5.0 +mpmath==1.3.0 +networkx==3.4.2 +ninja==1.11.1.1 +numpy<2 +nvitop==1.6.2 +optree==0.13.0 +packaging==24.1 +parso==0.8.4 +pexpect==4.9.0 +pickleshare==0.7.5 +pillow==10.2.0 +pip==24.2 +pkginfo==1.11.2 +pkgutil-resolve-name==1.3.10 +platformdirs==4.3.6 +pluggy==1.5.0 +prompt-toolkit==3.0.48 +protobuf==7.34.1 +psutil==6.1.0 +ptyprocess==0.7.0 +pure-eval==0.2.3 +pyarrow==24.0.0 +pycosat==0.6.6 +pycparser==2.22 +Pygments==2.18.0 +pyparsing==3.3.2 +PySocks==1.7.1 +python-dateutil==2.9.0.post0 +python-etcd==0.4.5 +pytz==2024.2 +PyYAML==6.0.2 +referencing==0.35.1 +requests==2.32.3 +rich==15.0.0 +rpds-py==0.20.0 +ruamel.yaml==0.18.6 +ruamel.yaml.clib==0.2.8 +setuptools==72.1.0 +shellingham==1.5.4 +six==1.16.0 +sortedcontainers==2.4.0 +soupsieve==2.5 +stack-data==0.6.2 +sympy==1.13.1 +tensorboard==2.20.0 +tensorboard-data-server==0.7.2 +torch==2.5.1+cu124 +torchaudio==2.5.1+cu124 +torchelastic==0.2.2 +torchvision==0.20.1+cu124 +tqdm==4.66.5 +traitlets==5.14.3 +triton==3.1.0 +truststore==0.9.2 +typer==0.25.1 +types-dataclasses==0.6.6 +typing-extensions==4.12.2 +urllib3==2.2.3 +wcwidth==0.6.0 +Werkzeug==3.1.8 +wheel==0.44.0 +zipp==3.20.2 +zstandard==0.23.0 diff --git a/verl.tar.gz b/verl.tar.gz new file mode 100644 index 0000000000000000000000000000000000000000..4c27aa3d70de328059d8df7633623b5da53ae33f --- /dev/null +++ b/verl.tar.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ee07dc56371e123a4f5afcd399a6b1b4624ce8af1edac0f8512ef1d98f5bced9 +size 5809181481 diff --git a/verl/.agent/skills/issue/SKILL.md b/verl/.agent/skills/issue/SKILL.md new file mode 100644 index 0000000000000000000000000000000000000000..84ceb8cd11e16e340d1f3bc8d74cfbc2b1db8c6b --- /dev/null +++ b/verl/.agent/skills/issue/SKILL.md @@ -0,0 +1,47 @@ +--- +name: issue +description: Create or update a GitHub issue following verl project conventions. +user_invocable: true +--- + +When the user asks to create or update an issue, follow these steps: + +### 1. Gather Context + +Read the following to understand available issue types and their required fields: + +- [`bug-report.yml`](.github/ISSUE_TEMPLATE/bug-report.yml) +- [`feature-request.yml`](.github/ISSUE_TEMPLATE/feature-request.yml) + +If updating an existing issue, read its current title, body, labels, and comments first. + +### 2. Determine Issue Type + +Based on the user's description, select the appropriate template: + +- **Bug report** ([`bug-report.yml`](.github/ISSUE_TEMPLATE/bug-report.yml)) — something is broken or behaves unexpectedly +- **Feature request** ([`feature-request.yml`](.github/ISSUE_TEMPLATE/feature-request.yml)) — a new capability or enhancement +- **Blank issue** — if neither template fits + +### 3. Compose the Issue + +Fill in the template fields based on information from the user and the codebase. For bug reports, run `python scripts/diagnose.py` to gather system info if possible. + +When updating, ensure the title and body still accurately reflect the current state of the issue. + +### 4. Check for Duplicates + +Search for existing issues before creating: + +``` +gh issue list --repo verl-project/verl --state open --search "" +``` + +If a duplicate exists, inform the user instead of creating a new one. + +### 5. Create or Update the Issue + +- **Create**: add `good first issue` and/or `call for contribution` labels if the issue is straightforward and suitable for new contributors. +- **Update**: update title, body, and labels as needed. + +Return the issue URL when done. diff --git a/verl/.agent/skills/pr/SKILL.md b/verl/.agent/skills/pr/SKILL.md new file mode 100644 index 0000000000000000000000000000000000000000..58055f617776eb5ae24fa220eaa2aa68a15476f4 --- /dev/null +++ b/verl/.agent/skills/pr/SKILL.md @@ -0,0 +1,33 @@ +--- +name: pr +description: Create or update a pull request following verl project conventions. +user_invocable: true +--- + +When the user asks to create or update a PR, follow these steps: + +### 1. Gather Context + +Read the following and understand the current branch's changes compared to main: + +- [`CONTRIBUTING.md`](CONTRIBUTING.md) +- [`PULL_REQUEST_TEMPLATE.md`](.github/PULL_REQUEST_TEMPLATE.md) + +If a PR already exists for this branch, also read its current title, body, and review comments. + +### 2. Compose PR Title and Body + +Follow the PR template strictly for both title format and body sections. Only check checklist boxes for steps that have actually been completed. + +When updating, ensure the title and body still accurately reflect **all** changes on the branch, not just the latest commit. + +### 3. Pre-submit Checks + +Run pre-commit and fix any issues before creating or pushing. + +### 4. Create or Update the PR + +- **Create**: target `main` by default unless the user specifies otherwise. +- **Update**: push new commits and update the title and body if the scope has changed. **Read the current PR title and body first** and incorporate any edits the user may have made directly on GitHub — never overwrite with a version generated from scratch. + +Return the PR URL when done. diff --git a/verl/.gemini/config.yaml b/verl/.gemini/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..66015ad30ed6768e06dceb91f11004be8a74bb04 --- /dev/null +++ b/verl/.gemini/config.yaml @@ -0,0 +1,10 @@ +have_fun: false +code_review: + disable: false + comment_severity_threshold: HIGH + max_review_comments: -1 + pull_request_opened: + help: false + summary: false + code_review: true +ignore_patterns: [] diff --git a/verl/.git-blame-ignore-revs b/verl/.git-blame-ignore-revs new file mode 100644 index 0000000000000000000000000000000000000000..649ba3ca862e8e47a92b932b337fe189fbd14e7c --- /dev/null +++ b/verl/.git-blame-ignore-revs @@ -0,0 +1,13 @@ +# Local uasge: git config blame.ignoreRevsFile .git-blame-ignore-revs + +# [dev] feat: immigrate from yapf & pylint to ruff based on pre-commit +# Changed 268 files, +10k/-9k lines. This is the biggest formatter change. +b00f77d8559b48d57a33c0132a5ba1c81891a536 + +# [ci] refactor: reduce ruff line-length from 300 to 120 +# Changed 238 files, +6k/-1k lines. Global formatting change. +00a10a8ef389556f957a2f36132b2358fd6a109f + +# [Lint] fix: linting errors in all files +# Changed 179 files, +1k/-3k lines. Global lint fix. +8e5ad4688a13de81727c014a3c2e2fb26324bc20 diff --git a/verl/.github/CODEOWNERS b/verl/.github/CODEOWNERS new file mode 100644 index 0000000000000000000000000000000000000000..a369368128a58c36d55cbce78a131876cd67be4b --- /dev/null +++ b/verl/.github/CODEOWNERS @@ -0,0 +1,25 @@ +/docs @eric-haibin-lin @zhaochenyang20 @hongpeng-guo +/docs/amd_tutorial @yushengsu-thu +/docs/slang_multiturn @zhaochenyang20 @SwordFaith +/docs/ascend_tutorial @wucong25 + +/third_party/sglang @zhaochenyang20 @SwordFaith +/third_party/vllm @PeterSH6 @wuxibin89 + +/examples/grpo_trainer @vermouth1992 @PeterSH6 @tardis-key @wucong25 @ji-huazhong + +/verl/single_controller @zw0610 @wuxibin89 @hongpeng-guo +/verl/trainer @eric-haibin-lin @vermouth1992 @tongyx361 @PeterSH6 +/verl/models/mcore @ISEEKYAN @vermouth1992 +/verl/models/transformers @vermouth1992 @PeterSH6 @tardis-key @wucong25 @ji-huazhong +/verl/workers/engine @eric-haibin-lin @vermouth1992 @ZihengJiang @ArronHZG +/verl/workers/roles @eric-haibin-lin @vermouth1992 @ZihengJiang +/verl/workers/engine/fsdp @eric-haibin-lin @vermouth1992 @ZihengJiang +/verl/workers/rollout/vllm_rollout @wuxibin89 @PeterSH6 @chenhaiq @ArronHZG +/verl/workers/rollout/sglang_rollout @zhaochenyang20 @SwordFaith @chenhaiq @ArronHZG +/verl/workers/engine/megatron @ISEEKYAN @vermouth1992 +/verl/experimental @wuxibin89 @ArronHZG + +/tests/single_controller @zw0610 @wuxibin89 +/tests/trainer @eric-haibin-lin @vermouth1992 @tongyx361 @PeterSH6 +/tests/workers/rollout/vllm_rollout @wuxibin89 @PeterSH6 @chenhaiq diff --git a/verl/.github/ISSUE_TEMPLATE/bug-report.yml b/verl/.github/ISSUE_TEMPLATE/bug-report.yml new file mode 100644 index 0000000000000000000000000000000000000000..67341f4139d9a5348f3887242b7c6accf10abc7b --- /dev/null +++ b/verl/.github/ISSUE_TEMPLATE/bug-report.yml @@ -0,0 +1,65 @@ +# modified from https://github.com/huggingface/transformers/blob/main/.github/ISSUE_TEMPLATE/bug-report.yml?plain=1 +name: "\U0001F41B Bug Report" +description: Submit a bug report to help us improve verl +labels: [ "bug" ] +body: + - type: markdown + attributes: + value: | + Thanks for taking the time to fill out this bug report! 🤗 + + - type: textarea + id: system-info + attributes: + label: System Info + description: Please share your system info with us. You can run the command `python scripts/diagnose.py` and copy-paste its output below. + placeholder: verl version, platform, python version, ... + validations: + required: true + + - type: checkboxes + id: information-scripts-examples + attributes: + label: Information + description: 'The problem arises when using:' + options: + - label: "The official example scripts" + - label: "My own modified scripts" + + - type: checkboxes + id: information-tasks + attributes: + label: Tasks + description: "The tasks I am working on are:" + options: + - label: "An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)" + - label: "My own task or dataset (give details below)" + + - type: textarea + id: reproduction + validations: + required: true + attributes: + label: Reproduction + description: | + Please provide a code sample that reproduces the problem you ran into. It can be a Colab link or just a code snippet. + Please include relevant config information with your code. + If you have code snippets, error messages, stack traces please provide them here as well. + Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting + Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code. + + placeholder: | + Steps to reproduce the behavior: + + 1. + 2. + 3. + + + - type: textarea + id: expected-behavior + validations: + required: true + attributes: + label: Expected behavior + description: "A clear and concise description of what you would expect to happen." \ No newline at end of file diff --git a/verl/.github/ISSUE_TEMPLATE/config.yml b/verl/.github/ISSUE_TEMPLATE/config.yml new file mode 100644 index 0000000000000000000000000000000000000000..ac09e8636d7a7577999a55ffdf7095dd0b656e52 --- /dev/null +++ b/verl/.github/ISSUE_TEMPLATE/config.yml @@ -0,0 +1,2 @@ +blank_issues_enabled: true +version: 0.1 diff --git a/verl/.github/ISSUE_TEMPLATE/feature-request.yml b/verl/.github/ISSUE_TEMPLATE/feature-request.yml new file mode 100644 index 0000000000000000000000000000000000000000..79cdb5a0d8e63ef1f0d62fa1160ab9dbc6ec1e28 --- /dev/null +++ b/verl/.github/ISSUE_TEMPLATE/feature-request.yml @@ -0,0 +1,32 @@ +# modified from https://github.com/huggingface/transformers/blob/main/.github/ISSUE_TEMPLATE/feature-request.yml?plain=1 +name: "\U0001F680 Feature request" +description: Submit a proposal/request for a new verl feature +labels: [ "Feature request" ] +body: + - type: textarea + id: feature-request + validations: + required: true + attributes: + label: Feature request + description: | + A clear and concise description of the feature proposal. Please provide a link to the paper and code in case they exist. + + - type: textarea + id: motivation + validations: + required: true + attributes: + label: Motivation + description: | + Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too. + + + - type: textarea + id: contribution + validations: + required: true + attributes: + label: Your contribution + description: | + Is there any way that you could help, e.g. by submitting a PR? Make sure to read the CONTRIBUTING.MD [readme](https://github.com/verl-project/verl/blob/main/CONTRIBUTING.md) \ No newline at end of file diff --git a/verl/.github/PULL_REQUEST_TEMPLATE.md b/verl/.github/PULL_REQUEST_TEMPLATE.md new file mode 100644 index 0000000000000000000000000000000000000000..e8b722b395575a74c8518ecf0d1cc08a7119fd8d --- /dev/null +++ b/verl/.github/PULL_REQUEST_TEMPLATE.md @@ -0,0 +1,41 @@ +### What does this PR do? + +> Add **concise** overview of what this PR aims to achieve or accomplish. Reference related GitHub issues and PRs that help with the review. + +### Checklist Before Starting + +- [ ] Search for similar PRs. Paste at least one query link here: ... +- [ ] Format the PR title as `[{modules}] {type}: {description}` (This will be checked by the CI) + - `{modules}` include `fsdp`, `megatron`, `veomni`, `sglang`, `vllm`, `rollout`, `trainer`, `ci`, `training_utils`, `recipe`, `hardware`, `deployment`, `ray`, `worker`, `single_controller`, `misc`, `perf`, `model`, `algo`, `env`, `tool`, `ckpt`, `doc`, `data`, `cfg`, `reward`, `fully_async`, `one_step_off` + - If this PR involves multiple modules, separate them with `,` like `[megatron, fsdp, doc]` + - `{type}` is in `feat`, `fix`, `refactor`, `chore`, `test` + - If this PR breaks any API (CLI arguments, config, function signature, etc.), add `[BREAKING]` to the beginning of the title. + - Example: `[BREAKING][fsdp, megatron] feat: dynamic batching` + +### Test + +> For changes that can not be tested by CI (e.g., algorithm implementation, new model support), validate by experiment(s) and show results like training curve plots, evaluation results, etc. + +### API and Usage Example + +> Demonstrate how the API changes if any, and provide usage example(s) if possible. + +```python +# Add code snippet or script demonstrating how to use this +``` + +### Design & Code Changes + +> Demonstrate the high-level design if this PR is complex, and list the specific changes. + +### Checklist Before Submitting + +> [!IMPORTANT] +> Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review. + +- [ ] Read the [Contribute Guide](https://github.com/verl-project/verl/blob/main/CONTRIBUTING.md). +- [ ] Apply [pre-commit checks](https://github.com/verl-project/verl/blob/main/CONTRIBUTING.md#code-linting-and-formatting): `pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always` +- [ ] Add / Update [the documentation](https://github.com/verl-project/verl/tree/main/docs). +- [ ] Add unit or end-to-end test(s) to [the CI workflow](https://github.com/verl-project/verl/tree/main/.github/workflows) to cover all the code. If not feasible, explain why: ... +- [ ] Once your PR is ready for CI, send a message in [the `ci-request` channel](https://verl-project.slack.com/archives/C091TCESWB1) in [the `verl` Slack workspace](https://join.slack.com/t/verl-project/shared_invite/zt-3855yhg8g-CTkqXu~hKojPCmo7k_yXTQ). (If not accessible, please try [the Feishu group (飞书群)](https://applink.larkoffice.com/client/chat/chatter/add_by_link?link_token=772jd4f1-cd91-441e-a820-498c6614126a).) +- [ ] If your PR is related to the `recipe` submodule, please also update the reference to the submodule commit via `git submodule update --remote` or `cd recipe && git pull origin main`. diff --git a/verl/.github/dependabot.yml b/verl/.github/dependabot.yml new file mode 100644 index 0000000000000000000000000000000000000000..24a3571c620c9c1e5ef7b4011851124c7a020626 --- /dev/null +++ b/verl/.github/dependabot.yml @@ -0,0 +1,9 @@ +## Enabled the dependabot to check the dependencies of the project +## Dependabot will open pull requests to update dependencies automatically + +version: 2 +updates: + - package-ecosystem: pip + directory: "/" + schedule: + interval: weekly \ No newline at end of file diff --git a/verl/.github/workflows/README.md b/verl/.github/workflows/README.md new file mode 100644 index 0000000000000000000000000000000000000000..d83c87b2e710db4f362171d1a6ee1ccf5db7d829 --- /dev/null +++ b/verl/.github/workflows/README.md @@ -0,0 +1,73 @@ +### Adding a New Workflow + +When adding a new workflow for continuous integration (CI), you have two runner options: a fixed runner or a machine from the vemlp. + +- **Fixed Runner**: To use a fixed runner, specify it in your workflow using the `runs-on` keyword, like `runs-on: [L20x8]`. +- **Vemlp Runner**: Opting for a Vemlp machine allows you to launch tasks elastically. + +Here is a template to assist you. This template is designed for using Vemlp machines. Currently, for each workflow, you need to create a `setup` and a `cleanup` job. When using this template, the main parts you need to modify are the `IMAGE` environment variable and the specific `job steps`. + +```yaml +name: Your Default Workflow + +on: + push: + branches: + - main + - v0.* + pull_request: + branches: + - main + - v0.* + paths: + - "**/*.py" + - ".github/workflows/template.yml" + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +permissions: + contents: read + +env: + IMAGE: "your vemlp image" # e.g. "verl-ci-cn-beijing.cr.volces.com/verlai/verl:sgl059.dev2" + DYNAMIC_RUNNER_URL: "https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner" # public veFaas api + +jobs: + setup: + if: github.repository_owner == 'verl-project' + runs-on: ubuntu-latest + outputs: + runner-label: ${{ steps.create-runner.outputs.runner-label }} + task-id: ${{ steps.create-runner.outputs.task-id }} + steps: + - uses: actions/checkout@v4 + - id: create-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "create" + faas-url: "${{ env.DYNAMIC_RUNNER_URL }}" + image: "${{ env.DEFAULT_IMAGE }}" + + your_job: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'default-runner' }}"] + steps: + xxxx # your jobs + + cleanup: + runs-on: ubuntu-latest + needs: [setup, your_job] + if: always() + steps: + - id: destroy-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "destroy" + faas-url: "${{ env.DYNAMIC_RUNNER_URL }}" + task-id: "${{ needs.setup.outputs.task-id }}" +``` + +### Model and Dataset +To avoid CI relies on network, we pre-download dataset on a NFS on the CI machine. The path for models are \${HOME}/models and the path for dataset is \${HOME}/models/hf_data. \ No newline at end of file diff --git a/verl/.github/workflows/check-pr-title.yml b/verl/.github/workflows/check-pr-title.yml new file mode 100644 index 0000000000000000000000000000000000000000..948ce5e3f01498ce4f569230cdb2dd384fc0cbfd --- /dev/null +++ b/verl/.github/workflows/check-pr-title.yml @@ -0,0 +1,58 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + + +on: + pull_request: + types: [opened, edited, synchronize] + +jobs: + check-title: + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.11' + + - name: Run PR title checker + run: python3 tests/special_sanity/check_pr_title.py + env: + PR_TITLE: ${{ github.event.pull_request.title }} + + - name: Run PR description checker + run: python3 tests/special_sanity/check_pr_description.py + env: + PR_TITLE: ${{ github.event.pull_request.title }} + GITHUB_EVENT_PATH: ${{ github.event_path }} diff --git a/verl/.github/workflows/cpu_unit_tests.yml b/verl/.github/workflows/cpu_unit_tests.yml new file mode 100644 index 0000000000000000000000000000000000000000..776f37ff98b1b468bece84b99152f97225588000 --- /dev/null +++ b/verl/.github/workflows/cpu_unit_tests.yml @@ -0,0 +1,117 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + +name: cpu_unit_tests + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch + push: + branches: + - main + - v0.* + pull_request: + branches: + - main + - v0.* + paths: + - "**/*.py" + - "!tests/special_sanity/**" + - .github/workflows/cpu_unit_tests.yml + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +env: + IMAGE: "verl-ci-cn-beijing.cr.volces.com/verlai/verl:vllm018.dev1" + DYNAMIC_RUNNER_ENDPOINT: "https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner" + +jobs: + setup: + if: github.repository_owner == 'verl-project' + runs-on: ubuntu-latest + outputs: + runner-label: ${{ steps.create-runner.outputs.runner-label }} + mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }} + steps: + - uses: actions/checkout@v4 + - id: create-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "create" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-image: "${{ env.IMAGE }}" + + cpu_unit_tests: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 20 # Increase this timeout value as needed + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + TORCH_COMPILE_DISABLE: 1 + TORCHINDUCTOR_DISABLE: 1 + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + - name: Download datasets + run: | + python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k + python3 examples/data_preprocess/geo3k.py --local_dataset_path ${HOME}/models/hf_data/hiyouga/geometry3k + - name: Running CPU unit tests + run: | + echo '[pytest]' > pytest.ini + echo 'python_files = *_on_cpu.py' >> pytest.ini + pytest -s -x --asyncio-mode=auto tests/ + cleanup: + runs-on: ubuntu-latest + needs: [setup, cpu_unit_tests] + if: always() + steps: + - id: destroy-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "destroy" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-task-id: "${{ needs.setup.outputs.mlp-task-id }}" diff --git a/verl/.github/workflows/doc.yml b/verl/.github/workflows/doc.yml new file mode 100644 index 0000000000000000000000000000000000000000..aa4a713deac84bcc021481f20dc887db5132cf8e --- /dev/null +++ b/verl/.github/workflows/doc.yml @@ -0,0 +1,101 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + + +name: doc_test + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch + push: + branches: + - main + - v0.* + pull_request: + branches: + - main + - v0.* + paths: + - "**/*.py" + - "docs/**" + - .github/workflows/doc.yml + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read # for checkout + pages: write # for deploy-pages + id-token: write # for deploy-pages + +jobs: + doc_test: + runs-on: ubuntu-latest + timeout-minutes: 5 # Increase this timeout value as needed + strategy: + matrix: + python-version: ["3.10"] + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0 + with: + python-version: ${{ matrix.python-version }} + - name: Install the current repository + run: | + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + pip install -r docs/requirements-docs.txt + + - name: Run doc make html + run: | + cd docs + make clean + make html SPHINXOPTS="--keep-going -w _build/sphinx.log" + if grep -q ": ERROR:" _build/sphinx.log; then + echo "🚨 Sphinx doc build contained ERRORs - see _build/sphinx.log" + exit 1 + fi + if grep -q "WARNING: document isn't included in any toctree" _build/sphinx.log; then + echo "🚨 Sphinx doc build contained WARNING. Please include newly added docs in index.rst. See _build/sphinx.log for details" + exit 1 + fi + if grep -q "WARNING: Inline emphasis" _build/sphinx.log; then + echo "🚨 Sphinx doc build contained WARNING. Please check inline emphasis is correct. See _build/sphinx.log for details" + exit 1 + fi + if grep -q "WARNING: Definition list ends without a blank line" _build/sphinx.log; then + echo "🚨 Sphinx doc build contained WARNING. Please check if the indentation is correct. See _build/sphinx.log for details" + exit 1 + fi diff --git a/verl/.github/workflows/docker-build-ascend-a2-qwen3_5.yml b/verl/.github/workflows/docker-build-ascend-a2-qwen3_5.yml new file mode 100644 index 0000000000000000000000000000000000000000..9c0f3eb2d9fa439fb873eb8040c1be2d77b4648f --- /dev/null +++ b/verl/.github/workflows/docker-build-ascend-a2-qwen3_5.yml @@ -0,0 +1,123 @@ +name: docker-build-ascend-a2-qwen3_5 + +on: + workflow_dispatch: + +env: + QUAY_REPO: quay.io/ascend/verl + +jobs: + build-push-digest: + name: build-${{ matrix.tag }} + if: ${{ github.repository_owner == 'verl-project' }} + runs-on: ${{ matrix.runner }} + permissions: + packages: write + contents: read + strategy: + fail-fast: false + matrix: + include: + - arch: linux/amd64 + runner: ubuntu-latest + tag: amd64 + - arch: linux/arm64 + runner: ubuntu-22.04-arm + tag: arm64 + steps: + - name: Remove unnecessary parts in github actions runners to free up disk space + uses: jlumbroso/free-disk-space@v1.3.1 + with: + tool-cache: true + + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Login to Quay.io + uses: docker/login-action@v3 + with: + registry: quay.io + username: ${{ secrets.QUAY_USERNAME }} + password: ${{ secrets.QUAY_PASSWORD }} + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + with: + install: true + driver: docker-container + use: true + + - name: Clean Docker cache before build + run: | + docker system prune -a -f --volumes || true + + - name: Build and push by digest + uses: docker/build-push-action@v6 + id: build + with: + context: . + platforms: ${{ matrix.arch }} + file: ./docker/ascend/Dockerfile.ascend_8.5.2_a2_qwen3-5 + push: true + # 关键：push-by-digest=true 只推送镜像层，不创建 tag + outputs: type=image,name=${{ env.QUAY_REPO }},push-by-digest=true,name-canonical=true,push=true + provenance: false + cache-from: type=gha + cache-to: type=gha,mode=max + build-args: | + BUILDKIT_INLINE_CACHE=1 + + - name: Export digest + run: | + mkdir -p ${{ runner.temp }}/digests + digest="${{ steps.build.outputs.digest }}" + touch "${{ runner.temp }}/digests/${digest#sha256:}" + + - name: Upload digest + uses: actions/upload-artifact@v4 + with: + name: digests-${{ matrix.tag }} + path: ${{ runner.temp }}/digests/* + if-no-files-found: error + retention-days: 1 + + merge-image: + name: merge + runs-on: ubuntu-latest + needs: build-push-digest + steps: + - name: Download digests + uses: actions/download-artifact@v4 + with: + path: ${{ runner.temp }}/digests + pattern: digests-* + merge-multiple: true + + - name: Login to Quay.io + uses: docker/login-action@v3 + with: + registry: quay.io + username: ${{ secrets.QUAY_USERNAME }} + password: ${{ secrets.QUAY_PASSWORD }} + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Merge and push multi-arch image + env: + IMAGE: ${{ env.QUAY_REPO }} + run: | + # 读取所有 digest 文件，构建 IMAGE@sha256:DIGEST 格式 + DIGESTS=$(printf "$IMAGE@sha256:%s " $(ls ${{ runner.temp }}/digests)) + + echo "Digests: $DIGESTS" + + # 创建多架构 manifest + docker buildx imagetools create \ + -t "$IMAGE:verl-8.5.2-910b-ubuntu22.04-py3.11-qwen3-5" \ + $DIGESTS diff --git a/verl/.github/workflows/docker-build-ascend-a2.yml b/verl/.github/workflows/docker-build-ascend-a2.yml new file mode 100644 index 0000000000000000000000000000000000000000..7c99649c762e6afd820629d12a0fa36a431766b4 --- /dev/null +++ b/verl/.github/workflows/docker-build-ascend-a2.yml @@ -0,0 +1,281 @@ +name: docker-build-ascend-a2 + +on: + workflow_dispatch: + push: + branches: ["main"] + paths: + - "docker/ascend/Dockerfile.ascend_8.5.0_a2" + - "docker/ascend/Dockerfile.ascend_8.5.0_a2_v0.7.1" + - ".github/workflows/docker-build-ascend-a2.yml" + + release: + types: [published] + schedule: + - cron: "0 16 * * *" + +env: + QUAY_REPO: quay.io/ascend/verl + +jobs: + # ============================================================ + # 第一阶段：分架构构建并推送 digest（main 镜像） + # ============================================================ + build-push-digest: + name: build-${{ matrix.tag }} + if: ${{ github.event_name != 'pull_request' && github.repository_owner == 'verl-project' }} + concurrency: + group: ${{ github.workflow }}-${{ github.ref }}-${{ matrix.tag }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + runs-on: ${{ matrix.runner }} + permissions: + packages: write + contents: read + strategy: + fail-fast: false + matrix: + include: + - arch: linux/amd64 + runner: ubuntu-latest + tag: amd64 + - arch: linux/arm64 + runner: ubuntu-22.04-arm + tag: arm64 + outputs: + tag: ${{ steps.version.outputs.tag }} + steps: + - name: Remove unnecessary parts in github actions runners to free up disk space + uses: jlumbroso/free-disk-space@v1.3.1 + with: + tool-cache: true + + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Get base image name and tag + id: base_image + run: | + BASE_IMAGE_FULL=$(grep '^FROM' ./docker/ascend/Dockerfile.ascend_8.5.0_a2 | head -1 | cut -d' ' -f2) + echo "Base image full: $BASE_IMAGE_FULL" + BASE_IMAGE_TAG=$(echo "$BASE_IMAGE_FULL" | cut -d':' -f2) + echo "Base image tag: $BASE_IMAGE_TAG" + NEW_IMAGE_NAME="verl-$BASE_IMAGE_TAG" + echo "New image name: $NEW_IMAGE_NAME" + echo "base_image_tag=$BASE_IMAGE_TAG" >> "$GITHUB_OUTPUT" + echo "new_image_name=$NEW_IMAGE_NAME" >> "$GITHUB_OUTPUT" + + - name: Get image tag + id: version + run: | + BRANCH_NAME=$(echo "${{ github.ref }}" | sed 's/refs\/heads\///g' | sed 's/[^a-zA-Z0-9._-]/_/g') + if [ "${{ github.event_name }}" = "release" ]; then + echo "tag=${{ steps.base_image.outputs.new_image_name }}-${{ github.event.release.tag_name }}" >> "$GITHUB_OUTPUT" + elif [ "$BRANCH_NAME" = "main" ]; then + echo "tag=${{ steps.base_image.outputs.new_image_name }}-latest" >> "$GITHUB_OUTPUT" + fi + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + with: + install: true + driver: docker-container + use: true + + - name: Login to Quay.io + uses: docker/login-action@v3 + with: + registry: quay.io + username: ${{ secrets.QUAY_USERNAME }} + password: ${{ secrets.QUAY_PASSWORD }} + + - name: Clean Docker cache before build + run: | + docker system prune -a -f --volumes || true + + - name: Build and push by digest (main image) + uses: docker/build-push-action@v6 + id: build_main + with: + context: . + platforms: ${{ matrix.arch }} + file: ./docker/ascend/Dockerfile.ascend_8.5.0_a2 + push: true + outputs: type=image,name=${{ env.QUAY_REPO }},push-by-digest=true,name-canonical=true,push=true + provenance: false + cache-from: type=gha + cache-to: type=gha,mode=max + build-args: | + BUILDKIT_INLINE_CACHE=1 + + - name: Export digest (main image) + run: | + mkdir -p ${{ runner.temp }}/digests + digest="${{ steps.build_main.outputs.digest }}" + touch "${{ runner.temp }}/digests/main-${digest#sha256:}" + + - name: Upload digest (main image) + uses: actions/upload-artifact@v4 + with: + name: digests-main-${{ matrix.tag }} + path: ${{ runner.temp }}/digests/* + if-no-files-found: error + retention-days: 1 + + # ============================================================ + # 第一阶段（v0.7.1）：分架构构建并推送 digest + # ============================================================ + build-push-digest-v071: + name: build-v0.7.1-${{ matrix.tag }} + if: ${{ github.event_name != 'pull_request' && github.repository_owner == 'verl-project' }} + concurrency: + group: ${{ github.workflow }}-${{ github.ref }}-v0.7.1-${{ matrix.tag }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + runs-on: ${{ matrix.runner }} + permissions: + packages: write + contents: read + strategy: + fail-fast: false + matrix: + include: + - arch: linux/amd64 + runner: ubuntu-latest + tag: amd64 + - arch: linux/arm64 + runner: ubuntu-22.04-arm + tag: arm64 + steps: + - name: Remove unnecessary parts in github actions runners to free up disk space + uses: jlumbroso/free-disk-space@v1.3.1 + with: + tool-cache: true + + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + with: + install: true + driver: docker-container + use: true + + - name: Login to Quay.io + uses: docker/login-action@v3 + with: + registry: quay.io + username: ${{ secrets.QUAY_USERNAME }} + password: ${{ secrets.QUAY_PASSWORD }} + + - name: Clean Docker cache before build + run: | + docker system prune -a -f --volumes || true + + - name: Build and push by digest (v0.7.1 image) + uses: docker/build-push-action@v6 + id: build_v071 + with: + context: . + platforms: ${{ matrix.arch }} + file: ./docker/ascend/Dockerfile.ascend_8.5.0_a2_v0.7.1 + push: true + outputs: type=image,name=${{ env.QUAY_REPO }},push-by-digest=true,name-canonical=true,push=true + provenance: false + cache-from: type=gha + cache-to: type=gha,mode=max + build-args: | + BUILDKIT_INLINE_CACHE=1 + + - name: Export digest (v0.7.1 image) + run: | + mkdir -p ${{ runner.temp }}/digests-v071 + digest="${{ steps.build_v071.outputs.digest }}" + touch "${{ runner.temp }}/digests-v071/v071-${digest#sha256:}" + + - name: Upload digest (v0.7.1 image) + uses: actions/upload-artifact@v4 + with: + name: digests-v071-${{ matrix.tag }} + path: ${{ runner.temp }}/digests-v071/* + if-no-files-found: error + retention-days: 1 + + # ============================================================ + # 第二阶段：合并多架构 manifest（main 镜像） + # ============================================================ + merge-image: + name: merge + runs-on: ubuntu-latest + needs: build-push-digest + steps: + - name: Download digests (main image) + uses: actions/download-artifact@v4 + with: + path: ${{ runner.temp }}/digests-main + pattern: digests-main-* + merge-multiple: true + + - name: Login to Quay.io + uses: docker/login-action@v3 + with: + registry: quay.io + username: ${{ secrets.QUAY_USERNAME }} + password: ${{ secrets.QUAY_PASSWORD }} + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Merge and push multi-arch image (main) + env: + IMAGE: ${{ env.QUAY_REPO }} + run: | + DIGESTS=$(printf "$IMAGE@sha256:%s " $(ls ${{ runner.temp }}/digests-main | sed 's/^main-//')) + echo "Main image digests: $DIGESTS" + docker buildx imagetools create \ + -t "$IMAGE:${{ needs.build-push-digest.outputs.tag }}" \ + $DIGESTS + + # ============================================================ + # 第二阶段（v0.7.1）：合并多架构 manifest + # ============================================================ + merge-image-v071: + name: merge-v0.7.1 + runs-on: ubuntu-latest + needs: build-push-digest-v071 + steps: + - name: Download digests (v0.7.1 image) + uses: actions/download-artifact@v4 + with: + path: ${{ runner.temp }}/digests-v071 + pattern: digests-v071-* + merge-multiple: true + + - name: Login to Quay.io + uses: docker/login-action@v3 + with: + registry: quay.io + username: ${{ secrets.QUAY_USERNAME }} + password: ${{ secrets.QUAY_PASSWORD }} + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Merge and push multi-arch image (v0.7.1) + env: + IMAGE: ${{ env.QUAY_REPO }} + run: | + DIGESTS=$(printf "$IMAGE@sha256:%s " $(ls ${{ runner.temp }}/digests-v071 | sed 's/^v071-//')) + echo "v0.7.1 image digests: $DIGESTS" + docker buildx imagetools create \ + -t "$IMAGE:verl-8.5.0-910b-ubuntu22.04-py3.11-v0.7.1" \ + $DIGESTS \ No newline at end of file diff --git a/verl/.github/workflows/docker-build-ascend-a3-qwen3_5.yml b/verl/.github/workflows/docker-build-ascend-a3-qwen3_5.yml new file mode 100644 index 0000000000000000000000000000000000000000..b32e9b17a77cec7ad1107648d75fd5424d1997d8 --- /dev/null +++ b/verl/.github/workflows/docker-build-ascend-a3-qwen3_5.yml @@ -0,0 +1,123 @@ +name: docker-build-ascend-a3-qwen3_5 + +on: + workflow_dispatch: + +env: + QUAY_REPO: quay.io/ascend/verl + +jobs: + build-push-digest: + name: build-${{ matrix.tag }} + if: ${{ github.repository_owner == 'verl-project' }} + runs-on: ${{ matrix.runner }} + permissions: + packages: write + contents: read + strategy: + fail-fast: false + matrix: + include: + - arch: linux/amd64 + runner: ubuntu-latest + tag: amd64 + - arch: linux/arm64 + runner: ubuntu-22.04-arm + tag: arm64 + steps: + - name: Remove unnecessary parts in github actions runners to free up disk space + uses: jlumbroso/free-disk-space@v1.3.1 + with: + tool-cache: true + + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Login to Quay.io + uses: docker/login-action@v3 + with: + registry: quay.io + username: ${{ secrets.QUAY_USERNAME }} + password: ${{ secrets.QUAY_PASSWORD }} + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + with: + install: true + driver: docker-container + use: true + + - name: Clean Docker cache before build + run: | + docker system prune -a -f --volumes || true + + - name: Build and push by digest + uses: docker/build-push-action@v6 + id: build + with: + context: . + platforms: ${{ matrix.arch }} + file: ./docker/ascend/Dockerfile.ascend_8.5.2_a3_qwen3-5 + push: true + # 关键：push-by-digest=true 只推送镜像层，不创建 tag + outputs: type=image,name=${{ env.QUAY_REPO }},push-by-digest=true,name-canonical=true,push=true + provenance: false + cache-from: type=gha + cache-to: type=gha,mode=max + build-args: | + BUILDKIT_INLINE_CACHE=1 + + - name: Export digest + run: | + mkdir -p ${{ runner.temp }}/digests + digest="${{ steps.build.outputs.digest }}" + touch "${{ runner.temp }}/digests/${digest#sha256:}" + + - name: Upload digest + uses: actions/upload-artifact@v4 + with: + name: digests-${{ matrix.tag }} + path: ${{ runner.temp }}/digests/* + if-no-files-found: error + retention-days: 1 + + merge-image: + name: merge + runs-on: ubuntu-latest + needs: build-push-digest + steps: + - name: Download digests + uses: actions/download-artifact@v4 + with: + path: ${{ runner.temp }}/digests + pattern: digests-* + merge-multiple: true + + - name: Login to Quay.io + uses: docker/login-action@v3 + with: + registry: quay.io + username: ${{ secrets.QUAY_USERNAME }} + password: ${{ secrets.QUAY_PASSWORD }} + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Merge and push multi-arch image + env: + IMAGE: ${{ env.QUAY_REPO }} + run: | + # 读取所有 digest 文件，构建 IMAGE@sha256:DIGEST 格式 + DIGESTS=$(printf "$IMAGE@sha256:%s " $(ls ${{ runner.temp }}/digests)) + + echo "Digests: $DIGESTS" + + # 创建多架构 manifest + docker buildx imagetools create \ + -t "$IMAGE:verl-8.5.2-a3-ubuntu22.04-py3.11-qwen3-5" \ + $DIGESTS diff --git a/verl/.github/workflows/docker-build-ascend-a3.yml b/verl/.github/workflows/docker-build-ascend-a3.yml new file mode 100644 index 0000000000000000000000000000000000000000..2446f65bcbcb94497711253f9c6b198353d72d95 --- /dev/null +++ b/verl/.github/workflows/docker-build-ascend-a3.yml @@ -0,0 +1,281 @@ +name: docker-build-ascend-a3 + +on: + workflow_dispatch: + push: + branches: ["main"] + paths: + - "docker/ascend/Dockerfile.ascend_8.5.0_a3" + - "docker/ascend/Dockerfile.ascend_8.5.0_a3_v0.7.1" + - ".github/workflows/docker-build-ascend-a3.yml" + + release: + types: [published] + schedule: + - cron: "0 19 * * *" + +env: + QUAY_REPO: quay.io/ascend/verl + +jobs: + # ============================================================ + # 第一阶段：分架构构建并推送 digest（main 镜像） + # ============================================================ + build-push-digest: + name: build-${{ matrix.tag }} + if: ${{ github.event_name != 'pull_request' && github.repository_owner == 'verl-project' }} + concurrency: + group: ${{ github.workflow }}-${{ github.ref }}-${{ matrix.tag }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + runs-on: ${{ matrix.runner }} + permissions: + packages: write + contents: read + strategy: + fail-fast: false + matrix: + include: + - arch: linux/amd64 + runner: ubuntu-latest + tag: amd64 + - arch: linux/arm64 + runner: ubuntu-22.04-arm + tag: arm64 + outputs: + tag: ${{ steps.version.outputs.tag }} + steps: + - name: Remove unnecessary parts in github actions runners to free up disk space + uses: jlumbroso/free-disk-space@v1.3.1 + with: + tool-cache: true + + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Get base image name and tag + id: base_image + run: | + BASE_IMAGE_FULL=$(grep '^FROM' ./docker/ascend/Dockerfile.ascend_8.5.0_a3 | head -1 | cut -d' ' -f2) + echo "Base image full: $BASE_IMAGE_FULL" + BASE_IMAGE_TAG=$(echo "$BASE_IMAGE_FULL" | cut -d':' -f2) + echo "Base image tag: $BASE_IMAGE_TAG" + NEW_IMAGE_NAME="verl-$BASE_IMAGE_TAG" + echo "New image name: $NEW_IMAGE_NAME" + echo "base_image_tag=$BASE_IMAGE_TAG" >> "$GITHUB_OUTPUT" + echo "new_image_name=$NEW_IMAGE_NAME" >> "$GITHUB_OUTPUT" + + - name: Get image tag + id: version + run: | + BRANCH_NAME=$(echo "${{ github.ref }}" | sed 's/refs\/heads\///g' | sed 's/[^a-zA-Z0-9._-]/_/g') + if [ "${{ github.event_name }}" = "release" ]; then + echo "tag=${{ steps.base_image.outputs.new_image_name }}-${{ github.event.release.tag_name }}" >> "$GITHUB_OUTPUT" + elif [ "$BRANCH_NAME" = "main" ]; then + echo "tag=${{ steps.base_image.outputs.new_image_name }}-latest" >> "$GITHUB_OUTPUT" + fi + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + with: + install: true + driver: docker-container + use: true + + - name: Login to Quay.io + uses: docker/login-action@v3 + with: + registry: quay.io + username: ${{ secrets.QUAY_USERNAME }} + password: ${{ secrets.QUAY_PASSWORD }} + + - name: Clean Docker cache before build + run: | + docker system prune -a -f --volumes || true + + - name: Build and push by digest (main image) + uses: docker/build-push-action@v6 + id: build_main + with: + context: . + platforms: ${{ matrix.arch }} + file: ./docker/ascend/Dockerfile.ascend_8.5.0_a3 + push: true + outputs: type=image,name=${{ env.QUAY_REPO }},push-by-digest=true,name-canonical=true,push=true + provenance: false + cache-from: type=gha + cache-to: type=gha,mode=max + build-args: | + BUILDKIT_INLINE_CACHE=1 + + - name: Export digest (main image) + run: | + mkdir -p ${{ runner.temp }}/digests + digest="${{ steps.build_main.outputs.digest }}" + touch "${{ runner.temp }}/digests/main-${digest#sha256:}" + + - name: Upload digest (main image) + uses: actions/upload-artifact@v4 + with: + name: digests-main-${{ matrix.tag }} + path: ${{ runner.temp }}/digests/* + if-no-files-found: error + retention-days: 1 + + # ============================================================ + # 第一阶段（v0.7.1）：分架构构建并推送 digest + # ============================================================ + build-push-digest-v071: + name: build-v0.7.1-${{ matrix.tag }} + if: ${{ github.event_name != 'pull_request' && github.repository_owner == 'verl-project' }} + concurrency: + group: ${{ github.workflow }}-${{ github.ref }}-v0.7.1-${{ matrix.tag }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + runs-on: ${{ matrix.runner }} + permissions: + packages: write + contents: read + strategy: + fail-fast: false + matrix: + include: + - arch: linux/amd64 + runner: ubuntu-latest + tag: amd64 + - arch: linux/arm64 + runner: ubuntu-22.04-arm + tag: arm64 + steps: + - name: Remove unnecessary parts in github actions runners to free up disk space + uses: jlumbroso/free-disk-space@v1.3.1 + with: + tool-cache: true + + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + with: + install: true + driver: docker-container + use: true + + - name: Login to Quay.io + uses: docker/login-action@v3 + with: + registry: quay.io + username: ${{ secrets.QUAY_USERNAME }} + password: ${{ secrets.QUAY_PASSWORD }} + + - name: Clean Docker cache before build + run: | + docker system prune -a -f --volumes || true + + - name: Build and push by digest (v0.7.1 image) + uses: docker/build-push-action@v6 + id: build_v071 + with: + context: . + platforms: ${{ matrix.arch }} + file: ./docker/ascend/Dockerfile.ascend_8.5.0_a3_v0.7.1 + push: true + outputs: type=image,name=${{ env.QUAY_REPO }},push-by-digest=true,name-canonical=true,push=true + provenance: false + cache-from: type=gha + cache-to: type=gha,mode=max + build-args: | + BUILDKIT_INLINE_CACHE=1 + + - name: Export digest (v0.7.1 image) + run: | + mkdir -p ${{ runner.temp }}/digests-v071 + digest="${{ steps.build_v071.outputs.digest }}" + touch "${{ runner.temp }}/digests-v071/v071-${digest#sha256:}" + + - name: Upload digest (v0.7.1 image) + uses: actions/upload-artifact@v4 + with: + name: digests-v071-${{ matrix.tag }} + path: ${{ runner.temp }}/digests-v071/* + if-no-files-found: error + retention-days: 1 + + # ============================================================ + # 第二阶段：合并多架构 manifest（main 镜像） + # ============================================================ + merge-image: + name: merge + runs-on: ubuntu-latest + needs: build-push-digest + steps: + - name: Download digests (main image) + uses: actions/download-artifact@v4 + with: + path: ${{ runner.temp }}/digests-main + pattern: digests-main-* + merge-multiple: true + + - name: Login to Quay.io + uses: docker/login-action@v3 + with: + registry: quay.io + username: ${{ secrets.QUAY_USERNAME }} + password: ${{ secrets.QUAY_PASSWORD }} + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Merge and push multi-arch image (main) + env: + IMAGE: ${{ env.QUAY_REPO }} + run: | + DIGESTS=$(printf "$IMAGE@sha256:%s " $(ls ${{ runner.temp }}/digests-main | sed 's/^main-//')) + echo "Main image digests: $DIGESTS" + docker buildx imagetools create \ + -t "$IMAGE:${{ needs.build-push-digest.outputs.tag }}" \ + $DIGESTS + + # ============================================================ + # 第二阶段（v0.7.1）：合并多架构 manifest + # ============================================================ + merge-image-v071: + name: merge-v0.7.1 + runs-on: ubuntu-latest + needs: build-push-digest-v071 + steps: + - name: Download digests (v0.7.1 image) + uses: actions/download-artifact@v4 + with: + path: ${{ runner.temp }}/digests-v071 + pattern: digests-v071-* + merge-multiple: true + + - name: Login to Quay.io + uses: docker/login-action@v3 + with: + registry: quay.io + username: ${{ secrets.QUAY_USERNAME }} + password: ${{ secrets.QUAY_PASSWORD }} + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Merge and push multi-arch image (v0.7.1) + env: + IMAGE: ${{ env.QUAY_REPO }} + run: | + DIGESTS=$(printf "$IMAGE@sha256:%s " $(ls ${{ runner.temp }}/digests-v071 | sed 's/^v071-//')) + echo "v0.7.1 image digests: $DIGESTS" + docker buildx imagetools create \ + -t "$IMAGE:verl-8.5.0-a3-ubuntu22.04-py3.11-v0.7.1" \ + $DIGESTS \ No newline at end of file diff --git a/verl/.github/workflows/docker-build-ascend-sglang-a2.yml b/verl/.github/workflows/docker-build-ascend-sglang-a2.yml new file mode 100644 index 0000000000000000000000000000000000000000..aa2e0ddff62988ef92e12f7b254b6c6893742f42 --- /dev/null +++ b/verl/.github/workflows/docker-build-ascend-sglang-a2.yml @@ -0,0 +1,84 @@ +name: docker-build-ascend-sglang-a2 + +on: + workflow_dispatch: + push: + branches: ["main"] + paths: + - "docker/ascend/Dockerfile.ascend.sglang_8.5.0_a2" + - ".github/workflows/docker-build-ascend-sglang-a2.yml" + release: + types: [published] + schedule: + - cron: "0 16 * * *" + +jobs: + build-ascend-sglang-image-a2: + if: ${{ github.event_name != 'pull_request' && github.repository_owner == 'verl-project' }} + runs-on: ubuntu-latest + concurrency: + group: ${{ github.workflow }}-${{ github.ref }}-build-ascend-sglang-image-a2 + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + steps: + - name: Remove unnecessary parts in github actions runners to free up disk space + uses: jlumbroso/free-disk-space@v1.3.1 + with: + tool-cache: true + + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Get base image name and tag + id: base_image + run: | + BASE_IMAGE_FULL=$(grep '^FROM' ./docker/ascend/Dockerfile.ascend.sglang_8.5.0_a2 | head -1 | cut -d' ' -f2) + echo "Base image full: $BASE_IMAGE_FULL" + BASE_IMAGE_TAG=$(echo "$BASE_IMAGE_FULL" | cut -d':' -f2) + echo "Base image tag: $BASE_IMAGE_TAG" + NEW_IMAGE_NAME="verl-sglang-$BASE_IMAGE_TAG" + echo "New image name: $NEW_IMAGE_NAME" + echo "base_image_tag=$BASE_IMAGE_TAG" >> "$GITHUB_OUTPUT" + echo "new_image_name=$NEW_IMAGE_NAME" >> "$GITHUB_OUTPUT" + + - name: Get image tag + id: version + run: | + BRANCH_NAME=$(echo "${{ github.ref }}" | sed 's/refs\/heads\///g' | sed 's/[^a-zA-Z0-9._-]/_/g') + if [ "${{ github.event_name }}" = "release" ]; then + echo "tag=${{ steps.base_image.outputs.new_image_name }}-${{ github.event.release.tag_name }}" >> "$GITHUB_OUTPUT" + elif [ "$BRANCH_NAME" = "main" ]; then + echo "tag=${{ steps.base_image.outputs.new_image_name }}-latest" >> "$GITHUB_OUTPUT" + fi + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Login to Quay.io + uses: docker/login-action@v3 + with: + registry: quay.io + username: ${{ secrets.QUAY_USERNAME }} + password: ${{ secrets.QUAY_PASSWORD }} + + - name: Clean Docker cache before build + run: | + docker system prune -a -f --volumes || true + + - name: Build and push images Quay + uses: docker/build-push-action@v6 + with: + context: . + platforms: linux/amd64,linux/arm64 + file: ./docker/ascend/Dockerfile.ascend.sglang_8.5.0_a2 + push: true + tags: | + quay.io/ascend/verl:${{ steps.version.outputs.tag }} + cache-from: type=gha + cache-to: type=gha,mode=max + build-args: | + BUILDKIT_INLINE_CACHE=1 diff --git a/verl/.github/workflows/docker-build-ascend-sglang-a3.yml b/verl/.github/workflows/docker-build-ascend-sglang-a3.yml new file mode 100644 index 0000000000000000000000000000000000000000..c4b9a88b3b7a6f97322e2132bad6808ef5c50352 --- /dev/null +++ b/verl/.github/workflows/docker-build-ascend-sglang-a3.yml @@ -0,0 +1,84 @@ +name: docker-build-ascend-sglang-a3 + +on: + workflow_dispatch: + push: + branches: ["main"] + paths: + - "docker/ascend/Dockerfile.ascend.sglang_8.5.0_a3" + - ".github/workflows/docker-build-ascend-sglang-a3.yml" + release: + types: [published] + schedule: + - cron: "0 16 * * *" + +jobs: + build-ascend-sglang-image-a3: + if: ${{ github.event_name != 'pull_request' && github.repository_owner == 'verl-project' }} + runs-on: ubuntu-latest + concurrency: + group: ${{ github.workflow }}-${{ github.ref }}-build-ascend-sglang-image-a3 + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + steps: + - name: Remove unnecessary parts in github actions runners to free up disk space + uses: jlumbroso/free-disk-space@v1.3.1 + with: + tool-cache: true + + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Get base image name and tag + id: base_image + run: | + BASE_IMAGE_FULL=$(grep '^FROM' ./docker/ascend/Dockerfile.ascend.sglang_8.5.0_a3 | head -1 | cut -d' ' -f2) + echo "Base image full: $BASE_IMAGE_FULL" + BASE_IMAGE_TAG=$(echo "$BASE_IMAGE_FULL" | cut -d':' -f2) + echo "Base image tag: $BASE_IMAGE_TAG" + NEW_IMAGE_NAME="verl-sglang-$BASE_IMAGE_TAG" + echo "New image name: $NEW_IMAGE_NAME" + echo "base_image_tag=$BASE_IMAGE_TAG" >> "$GITHUB_OUTPUT" + echo "new_image_name=$NEW_IMAGE_NAME" >> "$GITHUB_OUTPUT" + + - name: Get image tag + id: version + run: | + BRANCH_NAME=$(echo "${{ github.ref }}" | sed 's/refs\/heads\///g' | sed 's/[^a-zA-Z0-9._-]/_/g') + if [ "${{ github.event_name }}" = "release" ]; then + echo "tag=${{ steps.base_image.outputs.new_image_name }}-${{ github.event.release.tag_name }}" >> "$GITHUB_OUTPUT" + elif [ "$BRANCH_NAME" = "main" ]; then + echo "tag=${{ steps.base_image.outputs.new_image_name }}-latest" >> "$GITHUB_OUTPUT" + fi + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Login to Quay.io + uses: docker/login-action@v3 + with: + registry: quay.io + username: ${{ secrets.QUAY_USERNAME }} + password: ${{ secrets.QUAY_PASSWORD }} + + - name: Clean Docker cache before build + run: | + docker system prune -a -f --volumes || true + + - name: Build and push images Quay + uses: docker/build-push-action@v6 + with: + context: . + platforms: linux/amd64,linux/arm64 + file: ./docker/ascend/Dockerfile.ascend.sglang_8.5.0_a3 + push: true + tags: | + quay.io/ascend/verl:${{ steps.version.outputs.tag }} + cache-from: type=gha + cache-to: type=gha,mode=max + build-args: | + BUILDKIT_INLINE_CACHE=1 diff --git a/verl/.github/workflows/e2e_ascend.yml b/verl/.github/workflows/e2e_ascend.yml new file mode 100644 index 0000000000000000000000000000000000000000..e5377e597d55c2544561cbd1aa127bedba11956d --- /dev/null +++ b/verl/.github/workflows/e2e_ascend.yml @@ -0,0 +1,215 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + +name: e2e_ascend + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch + push: + branches: + - main + - v0.* + pull_request: + branches: + - main + paths: + - ".github/workflows/e2e_ascend.yml" + - "examples/data_preprocess/**" + - "examples/grpo_trainer/**" + - "examples/ppo_trainer/**" + - "examples/sft/**" + - "verl/experimental/one_step_off_policy/**" + - "tests/special_npu/**" + - "tests/special_sanity/check_device_api_usage.py" + - "verl/**" + - "pyproject.toml" + - "requirements-npu.txt" + - "setup.py" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +permissions: + contents: read + +jobs: + llm_rl_job: + if: github.repository_owner == 'verl-project' + name: E2E Ascend testing for RL training scenarios of LLM models + runs-on: linux-aarch64-a3-8 + timeout-minutes: 120 + container: + image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-a3-ubuntu22.04-py3.11-latest + options: >- + --shm-size 16g + env: + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - name: Check npu and CANN info + run: | + cat /usr/local/Ascend/ascend-toolkit/latest/"$(uname -i)"-linux/ascend_toolkit_install.info + npu-smi info + - name: Check initial pip list from image + run: | + pip list + - name: Checkout verl-project/verl repo + uses: actions/checkout@v4 + with: + fetch-depth: 0 + clean: true + - name: Install the current repository + run: | + pip install --no-deps -e . + - name: Check final pip list + run: | + pip list + - name: Preprocess gsm8k dataset + run: | + python examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k + - name: Running gsm8k e2e training tests with PPO on ASCEND NPU (FSDP backend) + run: | + ray stop --force + bash tests/special_npu/run_qwen3_06b_ppo.sh + rm -rf $HOME/ckpts + - name: Running gsm8k e2e training tests with GRPO on ASCEND NPU (FSDP backend) + run: | + ray stop --force + bash tests/special_npu/run_qwen2_5_05b_grpo.sh + rm -rf $HOME/ckpts + - name: Running gsm8k e2e training tests with GRPO on ASCEND NPU (MindSpeed backend) + run: | + ray stop --force + USE_DIST_CKPT=True bash tests/special_npu/run_qwen2_5_05b_grpo_mindspeed.sh + rm -rf $HOME/dist_ckpt/qwen2_5_05b_grpo_mindspeed + rm -rf $HOME/ckpts + - name: Running gsm8k e2e training tests with GRPO on ASCEND NPU (MindSpeed backend, MoE Model) + run: | + ray stop --force + USE_DIST_CKPT=True USE_DUMMY_MODEL=True DUMMY_MODEL_CONFIG_PATH=tests/special_e2e/ppo_trainer/expert_parallel/qwen3moe_minimal.json DUMMY_MODEL_PATH=$HOME/dist_ckpt/qwen3_30b_grpo_mindspeed bash tests/special_npu/run_qwen3_30b_grpo_mindspeed.sh + + engine_mindspeed_llm_rl_job: + if: github.repository_owner == 'verl-project' + name: E2E Ascend testing for RL training scenarios of LLM models using MindSpeed_LLM engine + runs-on: linux-aarch64-a3-8 + timeout-minutes: 120 + container: + image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-sglang-8.3.rc1-a3-ubuntu22.04-py3.11-latest + options: >- + --shm-size 16g + env: + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - name: Check npu and CANN info + run: | + cat /usr/local/Ascend/ascend-toolkit/latest/"$(uname -i)"-linux/ascend_toolkit_install.info + npu-smi info + - name: Check initial pip list from image + run: | + pip list + - name: Checkout verl-project/verl repo + uses: actions/checkout@v4 + with: + fetch-depth: 0 + clean: true + - name: Install the current repository + run: | + pip install --no-deps --no-build-isolation -e . + - name: Check final pip list + run: | + pip list + - name: Configure related dependencies + run: | + git clone --depth 1 --branch core_v0.12.1 https://github.com/NVIDIA/Megatron-LM.git /Megatron-LM + rm -rf /MindSpeed + git clone https://gitcode.com/ascend/MindSpeed.git /MindSpeed + git clone https://gitcode.com/ascend/MindSpeed-LLM.git /MindSpeed-LLM + - name: Preprocess gsm8k dataset + run: | + python examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k + - name: Running gsm8k e2e training tests with GRPO on ASCEND NPU (MindSpeedLLM backend) + run: | + ray stop --force + export PYTHONPATH=$PYTHONPATH:/Megatron-LM + export PYTHONPATH=$PYTHONPATH:/MindSpeed + export PYTHONPATH=$PYTHONPATH:/MindSpeed-LLM + bash tests/special_npu/run_qwen3_8b_grpo_mindspeedllm.sh + - name: Running gsm8k e2e training tests with GRPO on ASCEND NPU (MindSpeedLLM backend, MoE Model) + run: | + ray stop --force + export PYTHONPATH=$PYTHONPATH:/Megatron-LM + export PYTHONPATH=$PYTHONPATH:/MindSpeed + export PYTHONPATH=$PYTHONPATH:/MindSpeed-LLM + USE_DIST_CKPT=True USE_DUMMY_MODEL=True DUMMY_MODEL_CONFIG_PATH=tests/special_e2e/ppo_trainer/expert_parallel/qwen3moe_minimal.json DUMMY_MODEL_PATH=$HOME/dist_ckpt/qwen3_30b_grpo_mindspeedllm bash tests/special_npu/run_qwen3_30b_grpo_mindspeedllm.sh + + vlm_rl_job: + if: github.repository_owner == 'verl-project' + name: E2E Ascend testing for RL training scenarios of VLM models + runs-on: linux-aarch64-a3-8 + timeout-minutes: 120 + container: + image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-a3-ubuntu22.04-py3.11-latest + options: >- + --shm-size 16g + env: + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - name: Check npu and CANN info + run: | + cat /usr/local/Ascend/ascend-toolkit/latest/"$(uname -i)"-linux/ascend_toolkit_install.info + npu-smi info + - name: Check initial pip list from image + run: | + pip list + - name: Checkout verl-project/verl repo + uses: actions/checkout@v4 + with: + fetch-depth: 0 + clean: true + - name: Install the current repository + run: | + pip install --no-deps -e . + - name: Check final pip list + run: | + pip list + - name: Preprocess geo3k dataset + run: | + python examples/data_preprocess/geo3k.py --local_dataset_path ${HOME}/.cache/datasets/hiyouga/geometry3k + - name: Running geo3k e2e training tests with GRPO on ASCEND NPU + run: | + ray stop --force + bash tests/special_npu/run_qwen2_5_vl_3b_npu.sh + rm -rf $HOME/ckpts diff --git a/verl/.github/workflows/e2e_fully_async_policy.yml b/verl/.github/workflows/e2e_fully_async_policy.yml new file mode 100644 index 0000000000000000000000000000000000000000..5a60e9c37f9848bd7be057e456f7dd5e8b6dad55 --- /dev/null +++ b/verl/.github/workflows/e2e_fully_async_policy.yml @@ -0,0 +1,169 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + +name: e2e_fully_async_policy + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch + # For push, for now only anti-patterns are specified so it is more conservative + # and achieves higher coverage. + push: + branches: + - main + - v0.* + paths: + - "**/*.py" + - "!**/*.md" + - "!**/*.sh" + # Other entrypoints + - "!examples/*trainer*" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" + - "verl/experimental/fully_async_policy" + pull_request: + branches: + - main + - v0.* + paths: + - "**/*.py" + - "!**/*.md" + - "!**/*.sh" + # Other entrypoints + - "!examples/**" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" + # Home + - "verl/experimental/fully_async_policy" + # Entrypoints + - ".github/workflows/e2e_fully_async_policy.yml" + - "examples/data_preprocess/gsm8k.py" + - "tests/special_e2e/run_fully_async_policy.sh" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +env: + IMAGE: "verl-ci-cn-beijing.cr.volces.com/verlai/verl:vllm018.dev1" + DYNAMIC_RUNNER_ENDPOINT: "https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner" + +jobs: + setup: + if: github.repository_owner == 'verl-project' + runs-on: ubuntu-latest + outputs: + runner-label: ${{ steps.create-runner.outputs.runner-label }} + mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }} + steps: + - uses: actions/checkout@v4 + - id: create-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "create" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-image: "${{ env.IMAGE }}" + + # Test FSDP2 strategy + e2e_fully_async_policy_fsdp2: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 10 # Increase timeout for async training + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + ACTOR_STRATEGY: "fsdp2" + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + pip3 install cupy-cuda12x==13.6.0 + - name: Prepare GSM8K dataset + run: | + python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k + - name: Running the E2E test with fully_async_policy algorithm (FSDP2) + run: | + ray stop --force + bash tests/special_e2e/run_fully_async_policy.sh + + # Test Megatron strategy + e2e_fully_async_policy_megatron: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 10 # Increase timeout for async training + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + ACTOR_STRATEGY: "megatron" + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + pip3 install cupy-cuda12x==13.6.0 + pip3 install git+https://github.com/ISEEKYAN/mbridge.git@main --no-deps --no-build-isolation + - name: Prepare GSM8K dataset + run: | + python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k + - name: Running the E2E test with fully_async_policy algorithm (Megatron) + run: | + ray stop --force + bash tests/special_e2e/run_fully_async_policy.sh + + cleanup: + runs-on: ubuntu-latest + needs: [setup, e2e_fully_async_policy_fsdp2] + if: always() + steps: + - id: destroy-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "destroy" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-task-id: "${{ needs.setup.outputs.mlp-task-id }}" diff --git a/verl/.github/workflows/e2e_fully_async_policy_ascend.yml b/verl/.github/workflows/e2e_fully_async_policy_ascend.yml new file mode 100644 index 0000000000000000000000000000000000000000..b028286c0ee836590f4b621e1d3ca0e8118a55cc --- /dev/null +++ b/verl/.github/workflows/e2e_fully_async_policy_ascend.yml @@ -0,0 +1,169 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + +name: e2e_fully_async_policy_ascend + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch + # For push, for now only anti-patterns are specified so it is more conservative + # and achieves higher coverage. + push: + branches: + - main + - v0.* + paths: + - "**/*.py" + - "!**/*.md" + - "!**/*.sh" + # Other entrypoints + - "!examples/*trainer*" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" + - "verl/experimental/fully_async_policy" + pull_request: + branches: + - main + - v0.* + paths: + - "**/*.py" + - "!**/*.md" + - "!**/*.sh" + # Other entrypoints + - "!examples/**" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" + # Home + - "verl/experimental/fully_async_policy" + # Entrypoints + - ".github/workflows/e2e_fully_async_policy_ascend.yml" + - "examples/data_preprocess/gsm8k.py" + - "tests/special_e2e/run_fully_async_policy.sh" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +jobs: + # Test FSDP2 strategy + e2e_fully_async_policy_fsdp2_ascend: + if: github.repository_owner == 'verl-project' + runs-on: linux-aarch64-a3-8 + timeout-minutes: 60 # Increase this timeout value as needed + container: + image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-a3-ubuntu22.04-py3.11-latest + options: >- + --shm-size 16g + env: + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + ACTOR_STRATEGY: "fsdp2" + device_name: "npu" + steps: + - name: Check npu and CANN info + run: | + cat /usr/local/Ascend/ascend-toolkit/latest/"$(uname -i)"-linux/ascend_toolkit_install.info + npu-smi info + - name: Check initial pip list from image + run: | + pip list + - name: Checkout verl-project/verl repo + uses: actions/checkout@v4 + with: + fetch-depth: 0 + clean: true + - name: Install the current repository + run: | + pip install --no-deps -e . + - name: Check final pip list + run: | + pip list + - name: Prepare weights + run: | + ln -s /root/.cache/models ~/models + - name: Prepare GSM8K dataset + run: | + python examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k + - name: Running the E2E test with fully_async_policy algorithm (FSDP2) + run: | + ray stop --force + bash tests/special_e2e/run_fully_async_policy.sh + + # Test Megatron strategy + e2e_fully_async_policy_megatron_ascend: + if: github.repository_owner == 'verl-project' + runs-on: linux-aarch64-a3-8 + timeout-minutes: 60 # Increase this timeout value as needed + container: + image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-a3-ubuntu22.04-py3.11-latest + options: >- + --shm-size 16g + env: + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + ACTOR_STRATEGY: "megatron" + device_name: "npu" + steps: + - name: Check npu and CANN info + run: | + cat /usr/local/Ascend/ascend-toolkit/latest/"$(uname -i)"-linux/ascend_toolkit_install.info + npu-smi info + - name: Check initial pip list from image + run: | + pip list + - name: Checkout verl-project/verl repo + uses: actions/checkout@v4 + with: + fetch-depth: 0 + clean: true + - name: Install the current repository + run: | + pip install --no-deps -e . + - name: Check final pip list + run: | + pip list + - name: Prepare weights + run: | + ln -s /root/.cache/models ~/models + - name: Prepare GSM8K dataset + run: | + python examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k + - name: Running the E2E test with fully_async_policy algorithm (Megatron) + run: | + ray stop --force + bash tests/special_e2e/run_fully_async_policy.sh diff --git a/verl/.github/workflows/e2e_fully_async_policy_opd.yml b/verl/.github/workflows/e2e_fully_async_policy_opd.yml new file mode 100644 index 0000000000000000000000000000000000000000..c5d40a57879460140f8a50a19eccbd8c4ee486fd --- /dev/null +++ b/verl/.github/workflows/e2e_fully_async_policy_opd.yml @@ -0,0 +1,105 @@ +name: e2e_fully_async_policy_opd + +on: + push: + branches: + - main + - v0.* + paths: + - "**/*.py" + - "!**/*.md" + - "!**/*.sh" + - "!examples/*trainer*" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" + - "verl/experimental/fully_async_policy" + - "verl/trainer/distillation" + pull_request: + branches: + - main + - v0.* + paths: + - "**/*.py" + - "!**/*.md" + - "!**/*.sh" + - "!examples/**" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" + - "verl/experimental/fully_async_policy" + - "verl/trainer/distillation" + - ".github/workflows/e2e_fully_async_policy_opd.yml" + - "examples/data_preprocess/gsm8k.py" + - "examples/data_preprocess/geo3k.py" + - "tests/special_e2e/run_fully_async_policy_opd.sh" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +env: + IMAGE: "verl-ci-cn-beijing.cr.volces.com/verlai/verl:vllm018.dev1" + DYNAMIC_RUNNER_ENDPOINT: "https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner" + +jobs: + setup: + if: github.repository_owner == 'verl-project' + runs-on: ubuntu-latest + outputs: + runner-label: ${{ steps.create-runner.outputs.runner-label }} + mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }} + steps: + - uses: actions/checkout@v4 + - id: create-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "create" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-image: "${{ env.IMAGE }}" + + # Test Multi-Teacher OPD with Megatron strategy + e2e_fully_async_policy_opd_megatron: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 30 + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + pip3 install cupy-cuda12x==13.6.0 + pip3 install --force-reinstall --no-deps --no-build-isolation git+https://github.com/ISEEKYAN/mbridge.git@main + - name: Prepare datasets + run: | + python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k + python3 examples/data_preprocess/geo3k.py --local_dataset_path ${HOME}/models/hf_data/hiyouga/geometry3k/ --local_save_dir ${HOME}/data/geo3k + - name: Running the E2E test with fully_async_policy + Multi-Teacher OPD (Megatron) + run: | + ray stop --force + bash tests/special_e2e/run_fully_async_policy_opd.sh + + cleanup: + runs-on: ubuntu-latest + needs: [setup, e2e_fully_async_policy_opd_megatron] + if: always() + steps: + - id: destroy-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "destroy" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-task-id: "${{ needs.setup.outputs.mlp-task-id }}" diff --git a/verl/.github/workflows/e2e_fully_async_policy_trtllm.yml b/verl/.github/workflows/e2e_fully_async_policy_trtllm.yml new file mode 100644 index 0000000000000000000000000000000000000000..9ccd3a729f76fe7106cc7b677a7bc7d12b9edaab --- /dev/null +++ b/verl/.github/workflows/e2e_fully_async_policy_trtllm.yml @@ -0,0 +1,171 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + +name: e2e_fully_async_policy_trtllm + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch + push: + branches: + - main + - v0.* + paths: + - "**/*.py" + - "!**/*.md" + - "!**/*.sh" + - "!examples/*trainer*" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" + - "verl/workers/rollout/trtllm_rollout/**" + - "verl/experimental/fully_async_policy/**" + - "verl/checkpoint_engine/**" + pull_request: + branches: + - main + - v0.* + paths: + - "verl/workers/rollout/trtllm_rollout/**" + - "verl/experimental/fully_async_policy/**" + - "verl/checkpoint_engine/**" + - ".github/workflows/e2e_fully_async_policy_trtllm.yml" + - "tests/special_e2e/run_fully_async_policy.sh" + - "examples/data_preprocess/gsm8k.py" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +env: + IMAGE: "verl-ci-cn-beijing.cr.volces.com/verlai/verl:trtllm1.3.0rc14" + DYNAMIC_RUNNER_ENDPOINT: "https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner" + +jobs: + setup: + if: github.repository_owner == 'verl-project' + runs-on: ubuntu-latest + outputs: + runner-label: ${{ steps.create-runner.outputs.runner-label }} + mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }} + steps: + - uses: actions/checkout@v4 + - id: create-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "create" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-image: "${{ env.IMAGE }}" + + trtllm_async_unit_tests: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 30 + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" + TORCH_CUDA_ARCH_LIST: "7.5;8.0;8.9;9.0;10.0;12.0+PTX" + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Reinstall FlashInfer for runtime GPU arch + run: | + unset TORCH_CUDA_ARCH_LIST + pip3 install --force-reinstall --no-deps flashinfer-python + - name: Install the current repository + run: | + pip3 install pytest-asyncio + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + - name: Run TRT-LLM async unit tests + run: | + unset TORCH_CUDA_ARCH_LIST + export TRTLLM_ENABLE_PDL=0 + export TRTLLM_TEST_MODEL_PATH_ROOT="${HOME}/models" + ray stop --force + pytest -v -s --durations=20 \ + tests/workers/rollout/rollout_trtllm/test_trtllm_abort.py + + # Megatron multi-replica: 4 rollout GPUs split into 2 replicas × TP=2. + e2e_fully_async_policy_trtllm_megatron_multi_replica: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 20 + env: + HTTP_PROXY: ${{ secrets.PROXY_HTTP }} + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" + ACTOR_STRATEGY: "megatron" + ROLLOUT_NAME: "trtllm" + GEN_TP: "2" # 4 rollout GPUs / TP=2 = 2 replicas + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + pip3 install cupy-cuda12x==14.0.1 + - name: Prepare GSM8K dataset + run: | + python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k + - name: Running TRT-LLM fully-async E2E test with Megatron (multi-replica, 2×TP=2) + run: | + ray stop --force + bash tests/special_e2e/run_fully_async_policy.sh \ + data.max_response_length=128 \ + actor_rollout_ref.rollout.n=2 \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.1 \ + actor_rollout_ref.actor.ppo_mini_batch_size=4 \ + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=1536 + + cleanup: + runs-on: ubuntu-latest + needs: [setup, trtllm_async_unit_tests, e2e_fully_async_policy_trtllm_megatron_multi_replica] + if: always() + steps: + - id: destroy-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "destroy" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-task-id: "${{ needs.setup.outputs.mlp-task-id }}" diff --git a/verl/.github/workflows/e2e_one_step_off_policy.yml b/verl/.github/workflows/e2e_one_step_off_policy.yml new file mode 100644 index 0000000000000000000000000000000000000000..11c38f6a81e1cc40aa61d690e29252effd1369b5 --- /dev/null +++ b/verl/.github/workflows/e2e_one_step_off_policy.yml @@ -0,0 +1,169 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + +name: e2e_one_step_off_policy + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch + # For push, for now only anti-patterns are specified so it is more conservative + # and achieves higher coverage. + push: + branches: + - main + - v0.* + paths: + - "**/*.py" + - "!**/*.md" + - "!**/*.sh" + # Other entrypoints + - "!examples/*trainer*" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" + - "verl/experimental/one_step_off_policy" + pull_request: + branches: + - main + - v0.* + paths: + - "**/*.py" + - "!**/*.md" + - "!**/*.sh" + # Other entrypoints + - "!examples/**" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" + # Home + - "verl/experimental/one_step_off_policy" + # Entrypoints + - ".github/workflows/e2e_one_step_off_policy.yml" + - "examples/data_preprocess/gsm8k.py" + - "tests/special_e2e/run_one_step_off_policy.sh" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +env: + IMAGE: "verl-ci-cn-beijing.cr.volces.com/verlai/verl:vllm018.dev1" + DYNAMIC_RUNNER_ENDPOINT: "https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner" + +jobs: + setup: + if: github.repository_owner == 'verl-project' + runs-on: ubuntu-latest + outputs: + runner-label: ${{ steps.create-runner.outputs.runner-label }} + mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }} + steps: + - uses: actions/checkout@v4 + - id: create-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "create" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-image: "${{ env.IMAGE }}" + + # Test FSDP2 strategy + e2e_one_step_off_policy_fsdp2: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 10 # Increase timeout for async training + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + ACTOR_STRATEGY: "fsdp2" + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + pip3 install cupy-cuda12x==13.6.0 + - name: Prepare GSM8K dataset + run: | + python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k + - name: Running the E2E test with one_step_off_policy algorithm (FSDP2) + run: | + ray stop --force + bash tests/special_e2e/run_one_step_off_policy.sh + + # Test Megatron strategy + e2e_one_step_off_policy_megatron: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 10 # Increase timeout for async training + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + ACTOR_STRATEGY: "megatron" + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + pip3 install cupy-cuda12x==13.6.0 + - name: Prepare GSM8K dataset + run: | + python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k + - name: Running the E2E test with one_step_off_policy algorithm (Megatron) + run: | + ray stop --force + bash tests/special_e2e/run_one_step_off_policy.sh + + cleanup: + runs-on: ubuntu-latest + needs: + [setup, e2e_one_step_off_policy_fsdp2, e2e_one_step_off_policy_megatron] + if: always() + steps: + - id: destroy-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "destroy" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-task-id: "${{ needs.setup.outputs.mlp-task-id }}" diff --git a/verl/.github/workflows/e2e_one_step_off_policy_ascend.yml b/verl/.github/workflows/e2e_one_step_off_policy_ascend.yml new file mode 100644 index 0000000000000000000000000000000000000000..6656ccf10b6324992eab16fbed11e9644351cbb2 --- /dev/null +++ b/verl/.github/workflows/e2e_one_step_off_policy_ascend.yml @@ -0,0 +1,170 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + +name: e2e_one_step_off_policy_ascend + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch + # For push, for now only anti-patterns are specified so it is more conservative + # and achieves higher coverage. + push: + branches: + - main + - v0.* + paths: + - "**/*.py" + - "!**/*.md" + - "!**/*.sh" + # Other entrypoints + - "!examples/*trainer*" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" + - "verl/experimental/one_step_off_policy" + pull_request: + branches: + - main + - v0.* + paths: + - "**/*.py" + - "!**/*.md" + - "!**/*.sh" + # Other entrypoints + - "!examples/**" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" + # Home + - "verl/experimental/one_step_off_policy" + # Entrypoints + - ".github/workflows/e2e_one_step_off_policy_ascend.yml" + - "examples/data_preprocess/gsm8k.py" + - "tests/special_e2e/run_one_step_off_policy.sh" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +jobs: + # Test FSDP2 strategy + e2e_one_step_off_policy_fsdp2_ascend: + if: github.repository_owner == 'verl-project' + runs-on: linux-aarch64-a3-8 + timeout-minutes: 60 # Increase this timeout value as needed + container: + image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-a3-ubuntu22.04-py3.11-latest + options: >- + --shm-size 16g + env: + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + ACTOR_STRATEGY: "fsdp2" + device_name: "npu" + steps: + - name: Check npu and CANN info + run: | + cat /usr/local/Ascend/ascend-toolkit/latest/"$(uname -i)"-linux/ascend_toolkit_install.info + npu-smi info + - name: Check initial pip list from image + run: | + pip list + - name: Checkout verl-project/verl repo + uses: actions/checkout@v4 + with: + fetch-depth: 0 + clean: true + - name: Install the current repository + run: | + pip install --no-deps -e . + - name: Check final pip list + run: | + pip list + - name: Prepare weights + run: | + ln -s /root/.cache/models ~/models + - name: Prepare GSM8K dataset + run: | + python examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k + - name: Running the E2E test with one_step_off_policy algorithm (FSDP2) + run: | + ray stop --force + bash tests/special_e2e/run_one_step_off_policy.sh + + # Test Megatron strategy + e2e_one_step_off_policy_megatron_ascend: + if: github.repository_owner == 'verl-project' + runs-on: linux-aarch64-a3-8 + timeout-minutes: 60 # Increase this timeout value as needed + container: + image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-a3-ubuntu22.04-py3.11-latest + options: >- + --shm-size 16g + env: + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + ACTOR_STRATEGY: "megatron" + device_name: "npu" + steps: + - name: Check npu and CANN info + run: | + cat /usr/local/Ascend/ascend-toolkit/latest/"$(uname -i)"-linux/ascend_toolkit_install.info + npu-smi info + - name: Check initial pip list from image + run: | + pip list + - name: Checkout verl-project/verl repo + uses: actions/checkout@v4 + with: + fetch-depth: 0 + clean: true + - name: Install the current repository + run: | + pip install --no-deps -e . + - name: Check final pip list + run: | + pip list + - name: Prepare weights + run: | + ln -s /root/.cache/models ~/models + - name: Prepare GSM8K dataset + run: | + python examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k + - name: Running the E2E test with one_step_off_policy algorithm (Megatron) + run: | + ray stop --force + export PYTHONPATH=$PYTHONPATH:/Megatron-LM + bash tests/special_e2e/run_one_step_off_policy.sh diff --git a/verl/.github/workflows/e2e_ppo_grpo_trainer_trtllm.yml b/verl/.github/workflows/e2e_ppo_grpo_trainer_trtllm.yml new file mode 100644 index 0000000000000000000000000000000000000000..403d9d5d910c90e7dadb78c26fd9e6b1e645096a --- /dev/null +++ b/verl/.github/workflows/e2e_ppo_grpo_trainer_trtllm.yml @@ -0,0 +1,289 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + +name: e2e_ppo_trainer_megatron_trtllm + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch. + # For push, for now only anti-patterns are specified so it is more conservative + # and achieves higher coverage. + push: + branches: + - main + - v0.* + paths: + - "**/*.py" + # Other entrypoints + - "!verl/trainer/fsdp_sft_trainer.py" + # Recipes + - "!recipe/**" + # FSDP + - "!verl/workers/**/*dp_*.py" + pull_request: + branches: + - main + - v0.* + paths: + - "**/*.py" + # Other entrypoints + - "!docker/**" + # Docs + - "!**/*.md" + - "!docs/**" + - "!examples/**" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" + # Recipes + - "!recipe/**" + # FSDP + - "!verl/workers/**/*dp_*.py" + # Entrypoints + - "verl/workers/rollout/trtllm_rollout/**" + - "tests/workers/rollout/rollout_trtllm/**" + - ".github/workflows/e2e_ppo_grpo_trainer_trtllm.yml" + - "examples/data_preprocess/gsm8k.py" + - "examples/data_preprocess/geo3k.py" + - "examples/data_preprocess/dapo_multiturn_w_tool.py" + - "examples/data_preprocess/aime2024_multiturn_w_tool.py" + - "examples/grpo_trainer/run_qwen3_8b_megatron.sh" + - "examples/grpo_trainer/run_qwen3_30b_a3b_megatron.sh" + # add back when ppo flow is ready + # - "tests/special_e2e/run_ppo_trainer_megatron.sh" + # - "verl/trainer/main_ppo.py" + # - "verl/trainer/config/ppo_megatron_trainer.yaml" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +env: + IMAGE: "verl-ci-cn-beijing.cr.volces.com/verlai/verl:trtllm1.3.0rc14" + DYNAMIC_RUNNER_ENDPOINT: "https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner" + +jobs: + setup: + if: github.repository_owner == 'verl-project' + runs-on: ubuntu-latest + outputs: + runner-label: ${{ steps.create-runner.outputs.runner-label }} + mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }} + steps: + - uses: actions/checkout@v4 + - id: create-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "create" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-image: "${{ env.IMAGE }}" + + trtllm_unit_tests: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 30 # Increase this timeout value as needed + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + TORCH_CUDA_ARCH_LIST: "7.5;8.0;8.9;9.0;10.0;12.0+PTX" + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Reinstall FlashInfer for runtime GPU arch + run: | + unset TORCH_CUDA_ARCH_LIST + pip3 install --force-reinstall --no-deps flashinfer-python + - name: Install the current repository + run: | + pip3 install pytest-asyncio + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + - name: Run TRTLLM unit tests + run: | + unset TORCH_CUDA_ARCH_LIST # Let TRT-LLM CUTLASS DSL auto-detect runtime GPU arch + export TRTLLM_ENABLE_PDL=0 # Disable Programmatic Dependent Launch (uses SM 9.0 griddepcontrol) + export TRTLLM_TEST_MODEL_PATH_ROOT="${HOME}/models" + ray stop --force + pytest -v -s --durations=20 \ + tests/workers/rollout/rollout_trtllm/test_adapter.py \ + tests/workers/rollout/rollout_trtllm/test_async_server.py \ + tests/workers/rollout/rollout_trtllm/test_inter_node_rollout.py \ + tests/workers/rollout/rollout_trtllm/test_trtllm_rollout_utils.py + + e2e_grpo_trainer_megatron-qwen2: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 30 # Increase this timeout value as needed + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + TORCH_CUDA_ARCH_LIST: "7.5;8.0;8.9;9.0;10.0;12.0+PTX" + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Reinstall FlashInfer for runtime GPU arch + run: | + unset TORCH_CUDA_ARCH_LIST + pip3 install --force-reinstall --no-deps flashinfer-python + - name: Install the current repository + run: | + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + - name: Prepare GSM8K dataset + run: | + python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k --local_save_dir ${PWD}/data/gsm8k + - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with Megatron (Qwen) + run: | + unset TORCH_CUDA_ARCH_LIST # Let TRT-LLM CUTLASS DSL auto-detect runtime GPU arch + export TRTLLM_ENABLE_PDL=0 # Disable Programmatic Dependent Launch (uses SM 9.0 griddepcontrol) + ray stop --force + ACTOR_TP=2 \ + ACTOR_PP=1 \ + ROLLOUT_TP=2 \ + INFER_BACKEND=trtllm \ + MODEL_PATH="${HOME}/models/Qwen/Qwen2.5-0.5B-Instruct" \ + bash examples/grpo_trainer/run_qwen3_8b_megatron.sh \ + trainer.total_training_steps=1 \ + data.train_batch_size=32 \ + data.max_prompt_length=128 \ + data.max_response_length=64 \ + actor_rollout_ref.rollout.n=1 \ + actor_rollout_ref.actor.ppo_mini_batch_size=32 \ + actor_rollout_ref.rollout.max_num_seqs=32 \ + actor_rollout_ref.rollout.max_num_batched_tokens=1024 \ + actor_rollout_ref.rollout.max_model_len=256 \ + data.train_files="['${PWD}/data/gsm8k/train.parquet']" \ + data.val_files="['${PWD}/data/gsm8k/test.parquet']" \ + trainer.logger='["console"]' + - name: clean up + run: | + rm -rf checkpoints + e2e_grpo_trainer_megatron-vlm: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 30 # Increase this timeout value as needed + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + TORCH_CUDA_ARCH_LIST: "7.5;8.0;8.9;9.0;10.0;12.0+PTX" + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Reinstall FlashInfer for runtime GPU arch + run: | + unset TORCH_CUDA_ARCH_LIST + pip3 install --force-reinstall --no-deps flashinfer-python + - name: Install the current repository + run: | + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + pip3 install qwen_vl_utils + pip3 install mathruler + - name: Prepare GEO3K dataset + run: | + python3 examples/data_preprocess/geo3k.py --local_dataset_path ${HOME}/models/hf_data/hiyouga/geometry3k --local_save_dir ${PWD}/data/geo3k + - name: Running GEO3K E2E training tests with FSDP on 8 L20 GPUs (VLM) + run: | + ray stop --force + ROLLOUT_TP=2 \ + INFER_BACKEND=trtllm \ + MODEL_PATH="${HOME}/models/Qwen/Qwen3-VL-2B-Instruct" \ + bash examples/grpo_trainer/run_qwen2_5_vl_7b_fsdp.sh \ + trainer.total_training_steps=1 \ + data.train_batch_size=8 \ + data.max_response_length=64 \ + actor_rollout_ref.actor.ppo_mini_batch_size=8 \ + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=2048 \ + actor_rollout_ref.rollout.n=2 \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.35 \ + actor_rollout_ref.rollout.max_model_len=2048 \ + actor_rollout_ref.rollout.max_num_seqs=16 \ + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=2048 \ + data.train_files="['${PWD}/data/geo3k/train.parquet']" \ + data.val_files="['${PWD}/data/geo3k/test.parquet']" \ + trainer.logger='["console"]' + - name: clean up + run: | + rm -rf checkpoints + - name: Prepare GSM8K dataset (data_preprocess) + run: | + python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k --local_save_dir ${PWD}/data/gsm8k + - name: Running GRPO E2E with FP8 TRT-LLM rollout (Qwen3-0.6B, gsm8k) + run: | + unset TORCH_CUDA_ARCH_LIST # Let TRT-LLM CUTLASS DSL auto-detect runtime GPU arch + export TRTLLM_ENABLE_PDL=0 # Disable Programmatic Dependent Launch (uses SM 9.0 griddepcontrol) + ray stop --force + export INFER_TP=2 ACTOR_TP=2 ACTOR_PP=2 ACTOR_VPP=2 ACTOR_EP=1 ACTOR_CP=2 REF_TP=2 REF_PP=2 REF_VPP=2 REF_EP=1 REF_CP=2 GEN_MOE_TP=null GEN_MOE_EP=null + export NNODES=1 GPUS_PER_NODE=8 TRTLLM_MOE_BACKEND=TRITON + export DATA_DIR=${PWD} DAPO_MATH_TRAIN=${PWD}/data/gsm8k/train.parquet AIME_VAL=${PWD}/data/gsm8k/test.parquet MODEL_PATH=${HOME}/models/Qwen/Qwen3-0.6B + INFER_BACKEND=trtllm ROLLOUT_QUANTIZATION=fp8 bash examples/grpo_trainer/run_qwen3_30b_a3b_megatron.sh \ + reward_model.reward_kwargs.overlong_buffer_cfg.len=258 \ + reward_model.reward_kwargs.max_resp_len=512 \ + data.max_prompt_length=128 \ + data.max_response_length=64 \ + data.train_batch_size=32 \ + actor_rollout_ref.rollout.n=1 \ + actor_rollout_ref.rollout.max_num_seqs=16 \ + actor_rollout_ref.rollout.max_num_batched_tokens=1024 \ + actor_rollout_ref.rollout.max_model_len=1024 \ + actor_rollout_ref.actor.megatron.override_transformer_config.moe_grouped_gemm=False \ + actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=False \ + trainer.total_training_steps=1 \ + trainer.logger='["console"]' + - name: clean up + run: | + rm -rf checkpoints + + cleanup: + runs-on: ubuntu-latest + needs: [setup, trtllm_unit_tests, e2e_grpo_trainer_megatron-qwen2, e2e_grpo_trainer_megatron-vlm] + if: always() + steps: + - id: destroy-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "destroy" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-task-id: "${{ needs.setup.outputs.mlp-task-id }}" diff --git a/verl/.github/workflows/e2e_ppo_trainer.yml b/verl/.github/workflows/e2e_ppo_trainer.yml new file mode 100644 index 0000000000000000000000000000000000000000..357f0aa6bb6d819ae3098f73aa945dd1bb46123f --- /dev/null +++ b/verl/.github/workflows/e2e_ppo_trainer.yml @@ -0,0 +1,78 @@ +name: e2e_ppo_trainer + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch + # For push, for now only anti-patterns are specified so it is more conservative + # and achieves higher coverage. + push: + branches: + - main + - v0.* + paths: + - "**/*.py" + # Other entrypoints + - "!verl/trainer/fsdp_sft_trainer.py" + + # Megatron + - "!verl/workers/**/megatron_*.py" + + pull_request: + branches: + - main + - v0.* + paths: + - "**/*.py" + # Other entrypoints + - "!**/*.md" + - "!docker/**" + - "!examples/**" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" + # Docs + - "!docs/**" + + # Megatron + - "!verl/workers/**/megatron_*.py" + # Entrypoints + - ".github/workflows/e2e_ppo_trainer.yml" + - "examples/data_preprocess/gsm8k.py" + - "examples/data_preprocess/geo3k.py" + - "tests/special_e2e/ppo_trainer" + - "verl/trainer/main_ppo.py" + - "verl/trainer/config/ppo_trainer.yaml" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +jobs: + pre_commit_for_ppo: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.12"] + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0 + with: + python-version: ${{ matrix.python-version }} + - name: Install the current repository + run: | + pip install pre-commit hydra-core + pip3 install --no-deps -e . + - name: Set ruff --output-format=github + run: | + sed -i 's/--output-format=full/--output-format=github/' .pre-commit-config.yaml + git add .pre-commit-config.yaml + - uses: pre-commit/action@v3.0.1 + with: + extra_args: "" # Overriding default "--all-files" + diff --git a/verl/.github/workflows/e2e_ppo_trainer_megatron_sglang.yml b/verl/.github/workflows/e2e_ppo_trainer_megatron_sglang.yml new file mode 100644 index 0000000000000000000000000000000000000000..93bf17324a5700806ceb4765381202f1229658b7 --- /dev/null +++ b/verl/.github/workflows/e2e_ppo_trainer_megatron_sglang.yml @@ -0,0 +1,202 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + +name: e2e_ppo_trainer_megatron_sglang + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch. + # For push, for now only anti-patterns are specified so it is more conservative + # and achieves higher coverage. + push: + branches: + - main + - v0.* + paths: + - "**/*.py" + # Other entrypoints + - "!verl/trainer/fsdp_sft_trainer.py" # FSDP + - "!verl/workers/**/*dp_*.py" + - "!verl/utils/fsdp_utils.py" + - "!verl/utils/checkpoint/fsdp_checkpoint_manager.py" + - "!verl/model_merger/fsdp_model_merger.py" + pull_request: + branches: + - main + - v0.* + paths: + - "**/*.py" + # Other entrypoints + - "!docker/**" + # Docs + - "!**/*.md" + - "!docs/**" + - "!examples/**" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" # FSDP + - "!verl/workers/**/*dp_*.py" + - "!verl/utils/fsdp_utils.py" + - "!verl/utils/checkpoint/fsdp_checkpoint_manager.py" + - "!verl/model_merger/fsdp_model_merger.py" + # Entrypoints + - "verl/worksers/rollout/sglang_rollout/*" + - ".github/workflows/e2e_ppo_trainer_megatron_sglang.yml" + - "examples/data_preprocess/gsm8k.py" + - "examples/data_preprocess/geo3k.py" + - "tests/special_e2e/run_ppo_trainer_megatron.sh" + - "verl/trainer/main_ppo.py" + - "verl/trainer/config/ppo_megatron_trainer.yaml" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +env: + IMAGE: "verl-ci-cn-beijing.cr.volces.com/verlai/verl:sgl059.dev2" + DYNAMIC_RUNNER_ENDPOINT: "https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner" + +jobs: + setup: + if: github.repository_owner == 'verl-project' + runs-on: ubuntu-latest + outputs: + runner-label: ${{ steps.create-runner.outputs.runner-label }} + mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }} + steps: + - uses: actions/checkout@v4 + - id: create-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "create" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-image: "${{ env.IMAGE }}" + + e2e_ppo_trainer_megatron-deepseek: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 60 # Increase this timeout value as needed + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + ENGINE: sglang + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip3 install -r requirements-test.txt + pip3 install megatron-bridge --no-deps + pip3 install git+https://github.com/ISEEKYAN/mbridge.git@main --no-deps --no-build-isolation + pip3 install --no-deps -e . + - name: Prepare GSM8K dataset + run: | + python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k + - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with Megatron (DeepSeek) + run: | + ray stop --force + OPTIM_MEMORY_EFFICIENT=True ENGINE=sglang SAVE_FREQ=1 MODEL_ID=deepseek-ai/deepseek-coder-1.3b-instruct bash tests/special_e2e/run_ppo_trainer_megatron.sh + - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with Megatron (DeepSeek) + run: | + ray stop --force + export VLLM_USE_V1=1 + ray start --head + ENGINE=sglang MODE=async RESUME_MODE=auto MODEL_ID=deepseek-ai/deepseek-coder-1.3b-instruct TOTAL_TRAIN_STEPS=2 bash tests/special_e2e/run_ppo_trainer_megatron.sh + - name: Profiling GRPO GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with Megatron (Deepseek) + run: | + ray stop --force + PROFILE_ENABLE=True ENGINE=sglang ADV_ESTIMATOR=grpo USE_DYNAMIC_BSZ=False MODEL_ID=deepseek-ai/deepseek-coder-1.3b-instruct bash tests/special_e2e/run_ppo_trainer_megatron.sh + if [ -z "$( ls -A '/tmp/ray/session_latest/logs/nsight/' )" ]; then + echo "[ERROR] not found any profiling files" + exit 1 + else + echo "[SUCCESS] profile success" + fi + - name: clean up + run: | + rm -rf checkpoints + + # Qwen3-0.6B: dense, tie_word_embeddings=True + e2e_ppo_trainer_megatron-qwen3: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 60 # Increase this timeout value as needed + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + ENGINE: sglang + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + pip3 install megatron-bridge --no-deps + - name: Prepare GSM8K dataset + run: | + python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k + - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with Megatron (Qwen3) testing learning rate scheduler + run: | + ray stop --force + ALL_OFFLOAD=True VAL_BEFORE_TRAIN=True TEST_FREQ=1 SAVE_FREQ=1 LR_WARMUP_STEPS=1 TOTAL_TRAIN_STEPS=2 MODEL_ID=Qwen/Qwen3-0.6B bash tests/special_e2e/run_ppo_trainer_megatron.sh + - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with FP8 rollout + run: | + ray stop --force + export VLLM_USE_V1=1 + ROLLOUT_QUANTIZATION=fp8 TOTAL_TRAIN_STEPS=2 MODEL_ID=Qwen/Qwen3-0.6B bash tests/special_e2e/run_ppo_trainer_megatron.sh + - name: clean up + run: | + rm -rf checkpoints + + cleanup: + runs-on: ubuntu-latest + needs: + [setup, e2e_ppo_trainer_megatron-deepseek, e2e_ppo_trainer_megatron-qwen3] + if: always() + steps: + - id: destroy-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "destroy" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-task-id: "${{ needs.setup.outputs.mlp-task-id }}" diff --git a/verl/.github/workflows/e2e_ppo_trainer_megatron_sglang_2.yml b/verl/.github/workflows/e2e_ppo_trainer_megatron_sglang_2.yml new file mode 100644 index 0000000000000000000000000000000000000000..ec90add17594b2ecc3f887d529362f9a7a49d065 --- /dev/null +++ b/verl/.github/workflows/e2e_ppo_trainer_megatron_sglang_2.yml @@ -0,0 +1,199 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + +name: e2e_ppo_trainer_megatron_sglang_2 + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch. + # For push, for now only anti-patterns are specified so it is more conservative + # and achieves higher coverage. + push: + branches: + - main + - v0.* + paths: + - "**/*.py" + # Other entrypoints + - "!verl/trainer/fsdp_sft_trainer.py" # FSDP + - "!verl/workers/**/*dp_*.py" + - "!verl/utils/fsdp_utils.py" + - "!verl/utils/checkpoint/fsdp_checkpoint_manager.py" + - "!verl/model_merger/fsdp_model_merger.py" + pull_request: + branches: + - main + - v0.* + paths: + - "**/*.py" + # Other entrypoints + - "!docker/**" + # Docs + - "!**/*.md" + - "!docs/**" + - "!examples/**" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" # FSDP + - "!verl/workers/**/*dp_*.py" + - "!verl/utils/fsdp_utils.py" + - "!verl/utils/checkpoint/fsdp_checkpoint_manager.py" + - "!verl/model_merger/fsdp_model_merger.py" + # Entrypoints + - "verl/worksers/rollout/sglang_rollout/*" + - ".github/workflows/e2e_ppo_trainer_megatron_sglang.yml" + - "examples/data_preprocess/gsm8k.py" + - "examples/data_preprocess/geo3k.py" + - "tests/special_e2e/run_ppo_trainer_megatron.sh" + - "verl/trainer/main_ppo.py" + - "verl/trainer/config/ppo_megatron_trainer.yaml" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +env: + IMAGE: "verl-ci-cn-beijing.cr.volces.com/verlai/verl:sgl059.dev2" + DYNAMIC_RUNNER_ENDPOINT: "https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner" + +jobs: + setup: + if: github.repository_owner == 'verl-project' + runs-on: ubuntu-latest + outputs: + runner-label: ${{ steps.create-runner.outputs.runner-label }} + mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }} + steps: + - uses: actions/checkout@v4 + - id: create-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "create" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-image: "${{ env.IMAGE }}" + + e2e_ppo_trainer_fsdp_sglang: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 40 # Increase this timeout value as needed + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + - name: Prepare gsm8k dataset + run: | + ray stop --force + python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k + - name: Running GSM8K E2E training tests on 8 L20 GPUs with rmpad using function rm and save ckpt + run: | + ray stop --force + ENGINE=sglang bash tests/special_e2e/ppo_trainer/run_function_reward.sh + + e2e_ppo_trainer_fsdp-qwen2_5vl-3b: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 60 # Increase this timeout value as needed + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + # Geo3k + - name: Prepare GEO3K dataset + run: | + ray stop --force + python3 examples/data_preprocess/geo3k.py --local_dataset_path ${HOME}/models/hf_data/hiyouga/geometry3k/ + - name: Running GEO3K VLM E2E training tests on 8 L20 GPUs with rmpad using function rm + run: | + ray stop --force + TRAIN_FILES=$HOME/data/geo3k/train.parquet VAL_FILES=$HOME/data/geo3k/test.parquet \ + MAX_PROMPT_LEN=1536 MAX_RESPONSE_LEN=1536 \ + MODEL_ID=Qwen/Qwen2.5-VL-3B-Instruct \ + ADV_ESTIMATOR=grpo RM_PAD=True USE_KL=True ENABLE_CHUNKED_PREFILL=False \ + ENGINE=sglang ROLLOUT_MODE=async GPU_MEMORY_UTILIZATION=0.6 ACTOR_FSDP_PARAM_OFFLOAD=True \ + ACTOR_FSDP_OPTIMIZER_OFFLOAD=True REF_FSDP_PARAM_OFFLOAD=True \ + bash tests/special_e2e/ppo_trainer/run_function_reward.sh + - name: Running GEO3K VLM E2E with rmpad using torch fused kernel (Qwen2.5-VL) + run: | + ray stop --force + FUSED_KERNELS=True TRAIN_FILES=$HOME/data/geo3k/train.parquet VAL_FILES=$HOME/data/geo3k/test.parquet \ + MAX_PROMPT_LEN=1536 MAX_RESPONSE_LEN=1536 \ + MODEL_ID=Qwen/Qwen2.5-VL-3B-Instruct \ + ADV_ESTIMATOR=grpo RM_PAD=True USE_KL=True ENABLE_CHUNKED_PREFILL=False \ + ENGINE=sglang ROLLOUT_MODE=async GPU_MEMORY_UTILIZATION=0.6 ACTOR_FSDP_PARAM_OFFLOAD=True \ + ACTOR_FSDP_OPTIMIZER_OFFLOAD=True REF_FSDP_PARAM_OFFLOAD=True \ + bash tests/special_e2e/ppo_trainer/run_function_reward.sh + - name: Running GEO3K VLM E2E with rmpad using triton fused kernel (Qwen2.5-VL) + run: | + ray stop --force + FUSED_KERNELS=True FUSED_KERNEL_BACKEND=triton \ + TRAIN_FILES=$HOME/data/geo3k/train.parquet VAL_FILES=$HOME/data/geo3k/test.parquet \ + MAX_PROMPT_LEN=1536 MAX_RESPONSE_LEN=1536 \ + MODEL_ID=Qwen/Qwen2.5-VL-3B-Instruct \ + ADV_ESTIMATOR=grpo RM_PAD=True USE_KL=True ENABLE_CHUNKED_PREFILL=False \ + ENGINE=sglang ROLLOUT_MODE=async GPU_MEMORY_UTILIZATION=0.6 ACTOR_FSDP_PARAM_OFFLOAD=True \ + ACTOR_FSDP_OPTIMIZER_OFFLOAD=True REF_FSDP_PARAM_OFFLOAD=True \ + bash tests/special_e2e/ppo_trainer/run_function_reward.sh + + cleanup: + runs-on: ubuntu-latest + needs: + [setup, e2e_ppo_trainer_fsdp-qwen2_5vl-3b, e2e_ppo_trainer_fsdp_sglang] + if: always() + steps: + - id: destroy-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "destroy" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-task-id: "${{ needs.setup.outputs.mlp-task-id }}" diff --git a/verl/.github/workflows/e2e_ppo_trainer_megatron_sglang_2_ascend.yml b/verl/.github/workflows/e2e_ppo_trainer_megatron_sglang_2_ascend.yml new file mode 100644 index 0000000000000000000000000000000000000000..d12fb088c44c00469c24e3410b9a059221022287 --- /dev/null +++ b/verl/.github/workflows/e2e_ppo_trainer_megatron_sglang_2_ascend.yml @@ -0,0 +1,134 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + +name: e2e_ppo_trainer_megatron_sglang_2_ascend + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch. + # For push, for now only anti-patterns are specified so it is more conservative + # and achieves higher coverage. + push: + branches: + - main + - v0.* + paths: + - "**/*.py" + # Other entrypoints + - "!verl/trainer/fsdp_sft_trainer.py" # FSDP + - "!verl/workers/**/*dp_*.py" + - "!verl/utils/fsdp_utils.py" + - "!verl/utils/checkpoint/fsdp_checkpoint_manager.py" + - "!verl/model_merger/fsdp_model_merger.py" + pull_request: + branches: + - main + - v0.* + paths: + - "**/*.py" + # Other entrypoints + - "!docker/**" + # Docs + - "!**/*.md" + - "!docs/**" + - "!examples/**" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" # FSDP + - "!verl/workers/**/*dp_*.py" + - "!verl/utils/fsdp_utils.py" + - "!verl/utils/checkpoint/fsdp_checkpoint_manager.py" + - "!verl/model_merger/fsdp_model_merger.py" + # Entrypoints + - "verl/worksers/rollout/sglang_rollout/*" + - ".github/workflows/e2e_ppo_trainer_megatron_sglang_ascend.yml" + - "examples/data_preprocess/gsm8k.py" + - "examples/data_preprocess/geo3k.py" + - "tests/special_e2e/run_ppo_trainer_megatron.sh" + - "verl/trainer/main_ppo.py" + - "verl/trainer/config/ppo_megatron_trainer.yaml" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +jobs: + e2e_ppo_trainer_fsdp_sglang_ascend: + if: github.repository_owner == 'verl-project' + runs-on: linux-aarch64-a3-8 + timeout-minutes: 90 # Increase this timeout value as needed + container: + image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-sglang-8.5.0-a3-ubuntu22.04-py3.11-latest + options: >- + --shm-size 16g + env: + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES: "1" + ASCEND_RT_VISIBLE_DEVICES: "0,1,2,3,4,5,6,7" + HCCL_HOST_SOCKET_PORT_RANGE: "60000-60050" + HCCL_NPU_SOCKET_PORT_RANGE: "61000-61050" + steps: + - name: Check npu and CANN info + run: | + cat /usr/local/Ascend/ascend-toolkit/latest/"$(uname -i)"-linux/ascend_toolkit_install.info + npu-smi info + - name: Check initial pip list from image + run: | + pip list + - name: Checkout verl-project/verl repo + uses: actions/checkout@v4 + with: + fetch-depth: 0 + clean: true + - name: Install the current repository + run: | + pip3 install pytest pytest-asyncio + pip3 install --no-deps --no-build-isolation -e . + - name: Check final pip list + run: | + pip list + - name: Prepare weights + run: | + ln -s /root/.cache/models ~/models + - name: Prepare GSM8K dataset + run: | + ray stop --force + python examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k + - name: Running GSM8K E2E training tests on 8 NPUs with rmpad using function rm and save ckpt + run: | + ray stop --force + ENGINE=sglang bash tests/special_e2e/ppo_trainer/run_function_reward.sh diff --git a/verl/.github/workflows/e2e_ppo_trainer_megatron_sglang_ascend.yml b/verl/.github/workflows/e2e_ppo_trainer_megatron_sglang_ascend.yml new file mode 100644 index 0000000000000000000000000000000000000000..a19383e262350da9f05517fef2ba8ad02bb2377f --- /dev/null +++ b/verl/.github/workflows/e2e_ppo_trainer_megatron_sglang_ascend.yml @@ -0,0 +1,146 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + +name: e2e_ppo_trainer_megatron_sglang_ascend + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch. + # For push, for now only anti-patterns are specified so it is more conservative + # and achieves higher coverage. + push: + branches: + - main + - v0.* + paths: + - "**/*.py" + # Other entrypoints + - "!verl/trainer/fsdp_sft_trainer.py" # FSDP + - "!verl/workers/**/*dp_*.py" + - "!verl/utils/fsdp_utils.py" + - "!verl/utils/checkpoint/fsdp_checkpoint_manager.py" + - "!verl/model_merger/fsdp_model_merger.py" + pull_request: + branches: + - main + - v0.* + paths: + - "**/*.py" + # Other entrypoints + - "!docker/**" + # Docs + - "!**/*.md" + - "!docs/**" + - "!examples/**" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" # FSDP + - "!verl/workers/**/*dp_*.py" + - "!verl/utils/fsdp_utils.py" + - "!verl/utils/checkpoint/fsdp_checkpoint_manager.py" + - "!verl/model_merger/fsdp_model_merger.py" + # Entrypoints + - "verl/worksers/rollout/sglang_rollout/*" + - ".github/workflows/e2e_ppo_trainer_megatron_sglang_ascend.yml" + - "examples/data_preprocess/gsm8k.py" + - "examples/data_preprocess/geo3k.py" + - "tests/special_e2e/run_ppo_trainer_megatron.sh" + - "verl/trainer/main_ppo.py" + - "verl/trainer/config/ppo_megatron_trainer.yaml" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +jobs: + e2e_ppo_trainer_megatron-deepseek_ascend: + if: github.repository_owner == 'verl-project' + runs-on: linux-aarch64-a3-8 + timeout-minutes: 90 # Increase this timeout value as needed + container: + image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-sglang-8.5.0-a3-ubuntu22.04-py3.11-latest + options: >- + --shm-size 16g + env: + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES: "1" + ASCEND_RT_VISIBLE_DEVICES: "0,1,2,3,4,5,6,7" + HCCL_HOST_SOCKET_PORT_RANGE: "60000-60050" + HCCL_NPU_SOCKET_PORT_RANGE: "61000-61050" + ENGINE: sglang + steps: + - name: Check npu and CANN info + run: | + cat /usr/local/Ascend/ascend-toolkit/latest/"$(uname -i)"-linux/ascend_toolkit_install.info + npu-smi info + - name: Check initial pip list from image + run: | + pip list + - name: Checkout verl-project/verl repo + uses: actions/checkout@v4 + with: + fetch-depth: 0 + clean: true + - name: Install the current repository + run: | + pip3 install pytest pytest-asyncio + pip3 install megatron-bridge --no-deps + pip3 install git+https://github.com/ISEEKYAN/mbridge.git@main --no-deps --no-build-isolation + pip3 install --no-deps --no-build-isolation -e . + - name: Check final pip list + run: | + pip list + - name: Prepare weights + run: | + ln -s /root/.cache/models ~/models + - name: Prepare GSM8K dataset + run: | + python examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k + - name: Running GSM8K E2E training tests with 3D parallelism on 8 NPUs with Megatron (DeepSeek) + run: | + ray stop --force + OPTIM_MEMORY_EFFICIENT=True ENGINE=sglang COMMON_VPP=NULL USE_REMOVE_PADDING=True SAVE_FREQ=1 MODEL_ID=deepseek-ai/deepseek-coder-1.3b-instruct bash tests/special_e2e/run_ppo_trainer_megatron.sh + - name: Running GSM8K E2E training tests with 3D parallelism on 8 NPUs with Megatron (DeepSeek) + run: | + ray stop --force + export VLLM_USE_V1=1 + ray start --head + ENGINE=sglang MODE=async RESUME_MODE=auto COMMON_VPP=NULL USE_REMOVE_PADDING=True MODEL_ID=deepseek-ai/deepseek-coder-1.3b-instruct TOTAL_TRAIN_STEPS=2 bash tests/special_e2e/run_ppo_trainer_megatron.sh + - name: clean up + run: | + rm -rf checkpoints + rm -rf "${HOME}/profiling" diff --git a/verl/.github/workflows/e2e_ppo_trainer_megatron_vllm.yml b/verl/.github/workflows/e2e_ppo_trainer_megatron_vllm.yml new file mode 100644 index 0000000000000000000000000000000000000000..184ef19a6cca351732c5efd30be01962ced287f8 --- /dev/null +++ b/verl/.github/workflows/e2e_ppo_trainer_megatron_vllm.yml @@ -0,0 +1,227 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + +name: e2e_ppo_trainer_megatron_vllm + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch. + # For push, for now only anti-patterns are specified so it is more conservative + # and achieves higher coverage. + push: + branches: + - main + - v0.* + paths: + - "**/*.py" + # Other entrypoints + - "!verl/trainer/fsdp_sft_trainer.py" + # FSDP + - "!verl/workers/**/*dp_*.py" + - "!verl/utils/fsdp_utils.py" + - "!verl/utils/checkpoint/fsdp_checkpoint_manager.py" + - "!verl/model_merger/fsdp_model_merger.py" + pull_request: + branches: + - main + - v0.* + paths: + - "**/*.py" + # Other entrypoints + - "!docker/**" + # Docs + - "!**/*.md" + - "!docs/**" + - "!examples/**" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" + # FSDP + - "!verl/workers/**/*dp_*.py" + - "!verl/utils/fsdp_utils.py" + - "!verl/utils/checkpoint/fsdp_checkpoint_manager.py" + - "!verl/model_merger/fsdp_model_merger.py" + # Entrypoints + - ".github/workflows/e2e_ppo_trainer_megatron_vllm.yml" + - "examples/data_preprocess/gsm8k.py" + - "examples/data_preprocess/geo3k.py" + - "tests/special_e2e/run_ppo_trainer_megatron.sh" + - "verl/trainer/main_ppo.py" + - "verl/trainer/config/ppo_megatron_trainer.yaml" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +env: + IMAGE: "verl-ci-cn-beijing.cr.volces.com/verlai/verl:vllm018.dev1" + DYNAMIC_RUNNER_ENDPOINT: "https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner" + +jobs: + setup: + if: github.repository_owner == 'verl-project' + runs-on: ubuntu-latest + outputs: + runner-label: ${{ steps.create-runner.outputs.runner-label }} + mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }} + steps: + - uses: actions/checkout@v4 + - id: create-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "create" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-image: "${{ env.IMAGE }}" + + # deepseek-ai/deepseek-coder-1.3b-instruct: dense, tie_word_embeddings=False + e2e_ppo_trainer_megatron-deepseek: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 60 # Increase this timeout value as needed + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip3 install -r requirements-test.txt + pip3 install --no-deps --force-reinstall . + pip3 install megatron-bridge --no-deps + pip3 install git+https://github.com/ISEEKYAN/mbridge.git@main --no-deps --no-build-isolation + pip3 install math-verify + - name: Prepare GSM8K dataset + run: | + python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k + # Full training save&load + - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with Megatron, use mbridge e2e to pre-load and save (Deepseek) + run: | + ray stop --force + ALL_OFFLOAD=True SAVE_FREQ=1 MODEL_ID=deepseek-ai/deepseek-coder-1.3b-instruct COMMON_PP=4 COMMON_VPP=null COMMON_CP=1 USE_MBRIDGE=True USE_DIST_CKPT=False \ + bash tests/special_e2e/run_ppo_trainer_megatron.sh + - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with Megatron, use mbridge e2e to pre-load and save (Deepseek) + run: | + ray stop --force + RESUME_MODE=auto MODEL_ID=deepseek-ai/deepseek-coder-1.3b-instruct TOTAL_TRAIN_STEPS=2 SAVE_FREQ=1 COMMON_PP=4 COMMON_VPP=null COMMON_CP=1 USE_MBRIDGE=True USE_DIST_CKPT=False \ + bash tests/special_e2e/run_ppo_trainer_megatron.sh + # LoRA training save&load + - name: clean up and install Megatron-Bridge + run: | + rm -rf checkpoints + pip3 install git+https://github.com/NVIDIA-NeMo/Megatron-Bridge.git@6259ae8 --no-deps --no-build-isolation + pip3 install git+https://github.com/NVIDIA/Megatron-LM.git@7ca9dc5 --no-deps --no-build-isolation + pip3 install "nvidia-modelopt[torch]>=0.37.0" + - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with Megatron, use Megatron-Bridge LoRA e2e to pre-load and save (Deepseek) + run: | + ray stop --force + ALL_OFFLOAD=True SAVE_FREQ=1 MODEL_ID=deepseek-ai/deepseek-coder-1.3b-instruct COMMON_PP=4 LORA_RANK=8 COMMON_VPP=null COMMON_CP=1 USE_MBRIDGE=True VANILLA_MBRIDGE=False VALUE_VANILLA_MBRIDGE=False USE_DIST_CKPT=False \ + bash tests/special_e2e/run_ppo_trainer_megatron.sh + - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with Megatron, use Megatron-Bridge LoRA e2e to pre-load and save (Deepseek) + run: | + ray stop --force + RESUME_MODE=auto MODEL_ID=deepseek-ai/deepseek-coder-1.3b-instruct TOTAL_TRAIN_STEPS=2 SAVE_FREQ=1 COMMON_PP=4 LORA_RANK=8 COMMON_VPP=null COMMON_CP=1 USE_MBRIDGE=True VANILLA_MBRIDGE=False VALUE_VANILLA_MBRIDGE=False USE_DIST_CKPT=False \ + bash tests/special_e2e/run_ppo_trainer_megatron.sh + - name: clean up + run: | + rm -rf checkpoints + + # Qwen3-0.6B: dense, tie_word_embeddings=True + e2e_ppo_trainer_megatron-qwen3: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 60 # Increase this timeout value as needed + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + pip3 install megatron-bridge --no-deps + pip3 install math-verify + - name: Prepare GSM8K dataset + run: | + python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k + - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with Megatron (Qwen3) testing learning rate scheduler + run: | + ray stop --force + ALL_OFFLOAD=True VAL_BEFORE_TRAIN=True TEST_FREQ=1 SAVE_FREQ=1 LR_WARMUP_STEPS=1 TOTAL_TRAIN_STEPS=2 MODEL_ID=Qwen/Qwen3-0.6B bash tests/special_e2e/run_ppo_trainer_megatron.sh + - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with FP8 rollout + run: | + ray stop --force + export VLLM_USE_V1=1 + ROLLOUT_QUANTIZATION=fp8 TOTAL_TRAIN_STEPS=2 MODEL_ID=Qwen/Qwen3-0.6B bash tests/special_e2e/run_ppo_trainer_megatron.sh + - name: Install Megatron-LM and Megatron-Bridge for Megatron-FSDP + run: | + pip3 install --no-deps --no-build-isolation git+https://github.com/NVIDIA/Megatron-LM.git@d4cacef87 + pip3 install --no-deps --no-build-isolation git+https://github.com/NVIDIA-NeMo/Megatron-Bridge.git@6fea5bb + pip3 install "nvidia-modelopt[torch]>=0.37.0" + - name: Running GSM8K E2E PPO training tests on 8 L20 GPUs with Megatron-FSDP (Qwen3) + run: | + ray stop --force + ALL_OFFLOAD=False USE_MBRIDGE=True VANILLA_MBRIDGE=False USE_MEGATRON_FSDP=True \ + TOTAL_TRAIN_STEPS=2 MODEL_ID=Qwen/Qwen3-0.6B \ + COMMON_PP=1 COMMON_VPP=null COMMON_CP=1 COMMON_TP=1 INFER_TP=1 \ + bash tests/special_e2e/run_ppo_trainer_megatron.sh \ + ++actor_rollout_ref.actor.megatron.override_transformer_config.gradient_accumulation_fusion=False \ + ++actor_rollout_ref.ref.megatron.override_transformer_config.gradient_accumulation_fusion=False \ + ++critic.megatron.override_transformer_config.gradient_accumulation_fusion=False + - name: clean up + run: | + rm -rf checkpoints + + cleanup: + runs-on: ubuntu-latest + needs: + [setup, e2e_ppo_trainer_megatron-deepseek, e2e_ppo_trainer_megatron-qwen3] + if: always() + steps: + - id: destroy-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "destroy" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-task-id: "${{ needs.setup.outputs.mlp-task-id }}" diff --git a/verl/.github/workflows/e2e_ppo_trainer_megatron_vllm_2.yml b/verl/.github/workflows/e2e_ppo_trainer_megatron_vllm_2.yml new file mode 100644 index 0000000000000000000000000000000000000000..a4f26d2beb75b2f23d3437768c149103d1502fec --- /dev/null +++ b/verl/.github/workflows/e2e_ppo_trainer_megatron_vllm_2.yml @@ -0,0 +1,315 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + +name: e2e_ppo_trainer_megatron_vllm_2 + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch. + # For push, for now only anti-patterns are specified so it is more conservative + # and achieves higher coverage. + push: + branches: + - main + - v0.* + paths: + - "**/*.py" + # Other entrypoints + - "!verl/trainer/fsdp_sft_trainer.py" + # FSDP + - "!verl/workers/**/*dp_*.py" + - "!verl/utils/fsdp_utils.py" + - "!verl/utils/checkpoint/fsdp_checkpoint_manager.py" + - "!verl/model_merger/fsdp_model_merger.py" + pull_request: + branches: + - main + - v0.* + paths: + - "**/*.py" + # Other entrypoints + - "!docker/**" + # Docs + - "!**/*.md" + - "!docs/**" + - "!examples/**" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" + # FSDP + - "!verl/workers/**/*dp_*.py" + - "!verl/utils/fsdp_utils.py" + - "!verl/utils/checkpoint/fsdp_checkpoint_manager.py" + - "!verl/model_merger/fsdp_model_merger.py" + # Entrypoints + - ".github/workflows/e2e_ppo_trainer_megatron_vllm_2.yml" + - "examples/data_preprocess/gsm8k.py" + - "examples/data_preprocess/geo3k.py" + - "tests/special_e2e/run_ppo_trainer_megatron.sh" + - "verl/trainer/main_ppo.py" + - "verl/trainer/config/ppo_megatron_trainer.yaml" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +env: + IMAGE: "verl-ci-cn-beijing.cr.volces.com/verlai/verl:vllm018.dev1" + DYNAMIC_RUNNER_ENDPOINT: "https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner" + +jobs: + setup: + if: github.repository_owner == 'verl-project' + runs-on: ubuntu-latest + outputs: + runner-label: ${{ steps.create-runner.outputs.runner-label }} + mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }} + steps: + - uses: actions/checkout@v4 + - id: create-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "create" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-image: "${{ env.IMAGE }}" + + e2e_ppo_trainer_megatron-moe-expert-parallel: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 60 # Increase this timeout value as needed + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip3 install -r requirements-test.txt + pip3 install --no-deps --force-reinstall . + pip3 install git+https://github.com/NVIDIA-NeMo/Megatron-Bridge.git@6259ae8 --no-deps --no-build-isolation + pip3 install git+https://github.com/NVIDIA/Megatron-LM.git@7ca9dc5 --no-deps --no-build-isolation + pip3 install "nvidia-modelopt[torch]>=0.37.0" + - name: Prepare GSM8K dataset + run: | + python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k + - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with Megatron-Bridge (Qwen3-30B-A3B-Instruct-2507) + run: | + ray stop --force + ADV_ESTIMATOR=grpo USE_DUMMY_MODEL=True DUMMY_MODEL_CONFIG_PATH=tests/special_e2e/ppo_trainer/expert_parallel/qwen2moe_minimal.json \ + PPO_MAX_TOKEN_LEN=1024 FWD_MAX_TOKEN_LEN=1024 \ + MAX_PROMPT_LENGTH=512 MAX_RESPONSE_LENGTH=512 \ + MODEL_ID=Qwen/Qwen3-30B-A3B-Instruct-2507 USE_MBRIDGE=True VANILLA_MBRIDGE=False VALUE_VANILLA_MBRIDGE=False \ + COMMON_PP=2 COMMON_VPP=null COMMON_CP=1 COMMON_TP=4 COMMON_EP=4 COMMON_ETP=1 INFER_TP=8 \ + USE_DIST_CKPT=False ALL_OFFLOAD=True SKIP_SAVE_HF_MODEL=1 bash tests/special_e2e/run_ppo_trainer_megatron.sh + - name: Running GSM8K E2E training tests with 3D parallelism with FP8 rollout on 8 L20 GPUs with Megatron-Bridge (Qwen3-30B-A3B-Instruct-2507) + run: | + ray stop --force + ADV_ESTIMATOR=grpo USE_DUMMY_MODEL=True DUMMY_MODEL_CONFIG_PATH=tests/special_e2e/ppo_trainer/expert_parallel/qwen2moe_minimal.json \ + PPO_MAX_TOKEN_LEN=1024 FWD_MAX_TOKEN_LEN=1024 \ + MAX_PROMPT_LENGTH=512 MAX_RESPONSE_LENGTH=512 \ + MODEL_ID=Qwen/Qwen3-30B-A3B-Instruct-2507 USE_MBRIDGE=True VANILLA_MBRIDGE=False VALUE_VANILLA_MBRIDGE=False \ + COMMON_PP=2 COMMON_VPP=null COMMON_CP=1 COMMON_TP=4 COMMON_EP=4 COMMON_ETP=1 INFER_TP=2 \ + USE_DIST_CKPT=False ALL_OFFLOAD=True SKIP_SAVE_HF_MODEL=1 ROLLOUT_QUANTIZATION=fp8 bash tests/special_e2e/run_ppo_trainer_megatron.sh + - name: clean up + run: | + rm -rf checkpoints + - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with Megatron-Bridge LoRA (Qwen3-30B-A3B-Instruct-2507) + run: | + ray stop --force + ADV_ESTIMATOR=grpo USE_DUMMY_MODEL=True DUMMY_MODEL_CONFIG_PATH=tests/special_e2e/ppo_trainer/expert_parallel/qwen2moe_minimal.json \ + PPO_MAX_TOKEN_LEN=1024 FWD_MAX_TOKEN_LEN=1024 \ + MAX_PROMPT_LENGTH=512 MAX_RESPONSE_LENGTH=512 LORA_RANK=8 CRITIC_LORA_RANK=8 \ + MODEL_ID=Qwen/Qwen3-30B-A3B-Instruct-2507 USE_MBRIDGE=True VANILLA_MBRIDGE=False VALUE_VANILLA_MBRIDGE=False \ + COMMON_PP=2 COMMON_VPP=null COMMON_CP=1 COMMON_TP=4 COMMON_EP=2 COMMON_ETP=1 INFER_TP=8 \ + USE_DIST_CKPT=False LORA_MERGE=True ALL_OFFLOAD=True SKIP_SAVE_HF_MODEL=1 bash tests/special_e2e/run_ppo_trainer_megatron.sh + - name: clean up + run: | + rm -rf checkpoints + + e2e_ppo_trainer_fsdp_vllm: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 60 # Increase this timeout value as needed + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + - name: Prepare GSM8K dataset + run: | + ray stop --force + python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k + # Function RM + - name: Running GSM8K E2E training tests on 8 L20 GPUs with rmpad using function rm with validation and saving (FSDP_SIZE=8) + run: | + ray stop --force + VAL_BEFORE_TRAIN=True TEST_FREQ=1 SAVE_FREQ=1 SAVE_HF_MODEL=True VERL_EXP_NAME="qwen2.5-0.5b-function-reward-minimal-fsdp-size8" bash tests/special_e2e/ppo_trainer/run_function_reward.sh + - name: Running GSM8K E2E training tests on 8 L20 GPUs with rmpad using function rm after resuming + run: | + ray stop --force + RESUME_MODE=auto VERL_EXP_NAME="qwen2.5-0.5b-function-reward-minimal-fsdp-size8" bash tests/special_e2e/ppo_trainer/run_function_reward.sh + - name: Test merging FSDP checkpoints (Qwen Actor) + run: | + exp_name="qwen2.5-0.5b-function-reward-minimal-fsdp-size8" + python -m verl.model_merger test --backend fsdp --local_dir checkpoints/verl-test/${exp_name}/global_step_1/actor --test_hf_dir checkpoints/verl-test/${exp_name}/global_step_1/actor/huggingface + - name: Running GSM8K E2E training tests on 8 L20 GPUs with rmpad using function rm with validation and saving (DDP_SIZE=2, FSDP_SIZE=4) + run: | + ray stop --force + VAL_BEFORE_TRAIN=True TEST_FREQ=1 SAVE_FREQ=1 SAVE_HF_MODEL=True FSDP_SIZE=4 USE_KL=True VERL_EXP_NAME="qwen2.5-0.5b-function-reward-minimal-ddp-size2-fsdp-size4" bash tests/special_e2e/ppo_trainer/run_function_reward.sh + - name: Test merging DDP+FSDP checkpoints (Qwen Actor) + run: | + exp_name="qwen2.5-0.5b-function-reward-minimal-ddp-size2-fsdp-size4" + python -m verl.model_merger test --backend fsdp --local_dir checkpoints/verl-test/${exp_name}/global_step_1/actor --test_hf_dir checkpoints/verl-test/${exp_name}/global_step_1/actor/huggingface + - name: Running GSM8K E2E training tests on 8 L20 GPUs with rmpad using function rm with validation and saving (FSDP2) + run: | + ray stop --force + VAL_BEFORE_TRAIN=True TEST_FREQ=1 SAVE_FREQ=1 SAVE_HF_MODEL=True VERL_EXP_NAME="qwen2.5-0.5b-function-reward-minimal-fsdp2-size8" STRATEGY=fsdp2 bash tests/special_e2e/ppo_trainer/run_function_reward.sh + - name: Test merging FSDP2 checkpoints (Qwen Actor) + run: | + exp_name="qwen2.5-0.5b-function-reward-minimal-fsdp2-size8" + python -m verl.model_merger test --backend fsdp --local_dir checkpoints/verl-test/${exp_name}/global_step_1/actor --test_hf_dir checkpoints/verl-test/${exp_name}/global_step_1/actor/huggingface + - name: Running GSM8K E2E without rmpad using function rm + run: | + ray stop --force + RM_PAD=False bash tests/special_e2e/ppo_trainer/run_function_reward.sh + - name: Running GSM8K E2E training tests on 8 L20 GPUs with rmpad using function rm (GRPO) + run: | + ray stop --force + CUSTOM_REWARD_FN=True ADV_ESTIMATOR=grpo USE_KL=True bash tests/special_e2e/ppo_trainer/run_function_reward.sh + # - name: Running GSM8K E2E training tests on 8 L20 GPUs with rmpad using function rm (ReMax) + # run: | + # ray stop --force + # ADV_ESTIMATOR=remax USE_KL=True bash tests/special_e2e/ppo_trainer/run_function_reward.sh + # LoRA tests + - name: Running GSM8K E2E training tests on 8 L20 GPUs with grpo lora using function rm with use_shm + run: | + ray stop --force + ADV_ESTIMATOR=grpo USE_SHM=True LORA_RANK=32 LOAD_FORMAT=safetensors bash tests/special_e2e/ppo_trainer/run_function_reward.sh + - name: Running GSM8K E2E training tests on 8 L20 GPUs with grpo lora using function rm with use_shm and layered_summon + run: | + ray stop --force + ADV_ESTIMATOR=grpo USE_SHM=True LORA_RANK=32 LOAD_FORMAT=safetensors LAYERED_SUMMON=True TOTAL_TRAIN_STEPS=1 SAVE_FREQ=1 FSDP_SIZE=4 VERL_EXP_NAME="qwen2.5-0.5b-function-reward-minimal" bash tests/special_e2e/ppo_trainer/run_function_reward.sh + - name: Test GRPO LoRA checkpoints merging function + run: | + export EXP_NAME="qwen2.5-0.5b-function-reward-minimal" + ls checkpoints/verl-test/${EXP_NAME}/global_step_1/actor + cat checkpoints/verl-test/${EXP_NAME}/global_step_1/actor/huggingface/config.json + python3 -m verl.model_merger merge --backend fsdp --local_dir checkpoints/verl-test/${EXP_NAME}/global_step_1/actor/ --target_dir checkpoints/verl-test/${EXP_NAME}/global_step_1/actor/huggingface + - name: Running GSM8K E2E training tests on 8 L20 GPUs with grpo lora using function rm with use_shm and layered_summon with fsdp2 + run: | + ray stop --force + ADV_ESTIMATOR=grpo USE_SHM=True LORA_RANK=32 LOAD_FORMAT=safetensors LAYERED_SUMMON=True STRATEGY=fsdp2 bash tests/special_e2e/ppo_trainer/run_function_reward.sh + + e2e_ppo_trainer_fsdp-qwen2_5vl-3b: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 40 # Increase this timeout value as needed + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + # Geo3k + - name: Prepare GEO3K dataset + run: | + python3 examples/data_preprocess/geo3k.py --local_dataset_path ${HOME}/models/hf_data/hiyouga/geometry3k/ + - name: Running GEO3K VLM GRPO E2E training tests on 8 L20 GPUs with rmpad using function rm + run: | + ray stop --force + TRAIN_FILES=$HOME/data/geo3k/train.parquet VAL_FILES=$HOME/data/geo3k/test.parquet \ + MAX_PROMPT_LEN=1536 MAX_RESPONSE_LEN=1536 \ + MODEL_ID=Qwen/Qwen2.5-VL-3B-Instruct \ + ADV_ESTIMATOR=grpo RM_PAD=True USE_KL=True ENABLE_CHUNKED_PREFILL=False \ + SP_SIZE=2 \ + bash tests/special_e2e/ppo_trainer/run_function_reward.sh + + - name: Running GEO3K VLM PPO E2E training tests on 8 L20 GPUs with rmpad using function rm + run: | + ray stop --force + TRAIN_FILES=$HOME/data/geo3k/train.parquet VAL_FILES=$HOME/data/geo3k/test.parquet \ + MAX_PROMPT_LEN=1536 MAX_RESPONSE_LEN=1536 \ + MODEL_ID=Qwen/Qwen2.5-VL-3B-Instruct \ + ADV_ESTIMATOR=gae RM_PAD=True USE_KL=True ENABLE_CHUNKED_PREFILL=False \ + SP_SIZE=2 \ + bash tests/special_e2e/ppo_trainer/run_function_reward.sh + - name: Running GEO3K VLM GRPO E2E lora training tests on 8 L20 GPUs with rmpad using function rm + run: | + ray stop --force + TRAIN_FILES=$HOME/data/geo3k/train.parquet VAL_FILES=$HOME/data/geo3k/test.parquet \ + MAX_PROMPT_LEN=1536 MAX_RESPONSE_LEN=1536 \ + MODEL_ID=Qwen/Qwen2.5-VL-3B-Instruct \ + ADV_ESTIMATOR=grpo RM_PAD=True USE_KL=True ENABLE_CHUNKED_PREFILL=False \ + SP_SIZE=2 \ + LORA_RANK=32 LORA_EXCLUDE=".*visual.*" \ + bash tests/special_e2e/ppo_trainer/run_function_reward.sh + + cleanup: + runs-on: ubuntu-latest + needs: + [ + setup, + e2e_ppo_trainer_megatron-moe-expert-parallel, + e2e_ppo_trainer_fsdp-qwen2_5vl-3b, + e2e_ppo_trainer_fsdp_vllm, + ] + if: always() + steps: + - id: destroy-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "destroy" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-task-id: "${{ needs.setup.outputs.mlp-task-id }}" diff --git a/verl/.github/workflows/e2e_ppo_trainer_megatron_vllm_2_ascend.yml b/verl/.github/workflows/e2e_ppo_trainer_megatron_vllm_2_ascend.yml new file mode 100644 index 0000000000000000000000000000000000000000..878f087651cfc96b90c9aa7b670705db1139032e --- /dev/null +++ b/verl/.github/workflows/e2e_ppo_trainer_megatron_vllm_2_ascend.yml @@ -0,0 +1,231 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + +name: e2e_ppo_trainer_megatron_vllm_2_ascend + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch. + # For push, for now only anti-patterns are specified so it is more conservative + # and achieves higher coverage. + push: + branches: + - main + - v0.* + paths: + - "**/*.py" + # Other entrypoints + - "!verl/trainer/fsdp_sft_trainer.py" + # FSDP + - "!verl/workers/**/*dp_*.py" + - "!verl/utils/fsdp_utils.py" + - "!verl/utils/checkpoint/fsdp_checkpoint_manager.py" + - "!verl/model_merger/fsdp_model_merger.py" + pull_request: + branches: + - main + - v0.* + paths: + - "**/*.py" + # Other entrypoints + - "!docker/**" + # Docs + - "!**/*.md" + - "!docs/**" + - "!examples/**" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" + # FSDP + - "!verl/workers/**/*dp_*.py" + - "!verl/utils/fsdp_utils.py" + - "!verl/utils/checkpoint/fsdp_checkpoint_manager.py" + - "!verl/model_merger/fsdp_model_merger.py" + # Entrypoints + - ".github/workflows/e2e_ppo_trainer_megatron_vllm_2_ascend.yml" + - "examples/data_preprocess/gsm8k.py" + - "examples/data_preprocess/geo3k.py" + - "tests/special_e2e/run_ppo_trainer_megatron.sh" + - "verl/trainer/main_ppo.py" + - "verl/trainer/config/ppo_megatron_trainer.yaml" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +jobs: + e2e_ppo_trainer_fsdp_vllm_ascend: + if: github.repository_owner == 'verl-project' + runs-on: linux-aarch64-a2b3-8 + timeout-minutes: 90 # Increase this timeout value as needed + container: + image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest + options: >- + --shm-size 16g + env: + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - name: Check npu and CANN info + run: | + cat /usr/local/Ascend/ascend-toolkit/latest/"$(uname -i)"-linux/ascend_toolkit_install.info + npu-smi info + - name: Check initial pip list from image + run: | + pip list + - name: Checkout verl-project/verl repo + uses: actions/checkout@v4 + with: + fetch-depth: 0 + clean: true + - name: Install the current repository + run: | + pip install --no-deps -e . + - name: Check final pip list + run: | + pip list + - name: Prepare weights + run: | + ln -s /root/.cache/models ~/models + - name: Prepare GSM8K dataset + run: | + python examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k + # Function RM + - name: Running GSM8K E2E training tests on 8 L20 GPUs with rmpad using function rm with validation and saving (DDP_SIZE=2, FSDP_SIZE=4) + run: | + ray stop --force + VAL_BEFORE_TRAIN=True TEST_FREQ=1 SAVE_FREQ=1 SAVE_HF_MODEL=True FSDP_SIZE=4 USE_KL=True VERL_EXP_NAME="qwen2.5-0.5b-function-reward-minimal-ddp-size2-fsdp-size4" bash tests/special_e2e/ppo_trainer/run_function_reward.sh + - name: Test merging DDP+FSDP checkpoints (Qwen Actor) + run: | + exp_name="qwen2.5-0.5b-function-reward-minimal-ddp-size2-fsdp-size4" + python -m verl.model_merger test --backend fsdp --local_dir checkpoints/verl-test/${exp_name}/global_step_1/actor --test_hf_dir checkpoints/verl-test/${exp_name}/global_step_1/actor/huggingface + - name: Running GSM8K E2E training tests on 8 L20 GPUs with rmpad using function rm with validation and saving (FSDP2) + run: | + ray stop --force + VAL_BEFORE_TRAIN=True TEST_FREQ=1 SAVE_FREQ=1 SAVE_HF_MODEL=True VERL_EXP_NAME="qwen2.5-0.5b-function-reward-minimal-fsdp2-size8" STRATEGY=fsdp2 bash tests/special_e2e/ppo_trainer/run_function_reward.sh + - name: Test merging FSDP2 checkpoints (Qwen Actor) + run: | + exp_name="qwen2.5-0.5b-function-reward-minimal-fsdp2-size8" + python -m verl.model_merger test --backend fsdp --local_dir checkpoints/verl-test/${exp_name}/global_step_1/actor --test_hf_dir checkpoints/verl-test/${exp_name}/global_step_1/actor/huggingface + - name: Running GSM8K E2E without rmpad using function rm + run: | + ray stop --force + RM_PAD=False bash tests/special_e2e/ppo_trainer/run_function_reward.sh + - name: Running GSM8K E2E training tests on 8 L20 GPUs with rmpad using function rm (GRPO) + run: | + ray stop --force + CUSTOM_REWARD_FN=True ADV_ESTIMATOR=grpo USE_KL=True bash tests/special_e2e/ppo_trainer/run_function_reward.sh + - name: Running GSM8K E2E training tests on 8 L20 GPUs with grpo lora using function rm with use_shm and layered_summon + run: | + ray stop --force + ADV_ESTIMATOR=grpo USE_SHM=True LORA_RANK=32 LOAD_FORMAT=safetensors LAYERED_SUMMON=True TOTAL_TRAIN_STEPS=1 SAVE_FREQ=1 FSDP_SIZE=4 VERL_EXP_NAME="qwen2.5-0.5b-function-reward-minimal" bash tests/special_e2e/ppo_trainer/run_function_reward.sh + - name: Test GRPO LoRA checkpoints merging function + run: | + export EXP_NAME="qwen2.5-0.5b-function-reward-minimal" + ls checkpoints/verl-test/${EXP_NAME}/global_step_1/actor + cat checkpoints/verl-test/${EXP_NAME}/global_step_1/actor/huggingface/config.json + python3 -m verl.model_merger merge --backend fsdp --local_dir checkpoints/verl-test/${EXP_NAME}/global_step_1/actor/ --target_dir checkpoints/verl-test/${EXP_NAME}/global_step_1/actor/huggingface + - name: Running GSM8K E2E training tests on 8 L20 GPUs with grpo lora using function rm with use_shm and layered_summon with fsdp2 + run: | + ray stop --force + ADV_ESTIMATOR=grpo USE_SHM=True LORA_RANK=32 LOAD_FORMAT=safetensors LAYERED_SUMMON=True STRATEGY=fsdp2 bash tests/special_e2e/ppo_trainer/run_function_reward.sh + + e2e_ppo_trainer_fsdp-qwen2_5vl-3b_ascend: + if: github.repository_owner == 'verl-project' + runs-on: linux-aarch64-a2b3-8 + timeout-minutes: 60 # Increase this timeout value as needed + container: + image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest + options: >- + --shm-size 16g + env: + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - name: Check npu and CANN info + run: | + cat /usr/local/Ascend/ascend-toolkit/latest/"$(uname -i)"-linux/ascend_toolkit_install.info + npu-smi info + - name: Check initial pip list from image + run: | + pip list + - name: Checkout verl-project/verl repo + uses: actions/checkout@v4 + with: + fetch-depth: 0 + clean: true + - name: Install the current repository + run: | + pip install --no-deps -e . + pip install trl==0.26.0 + - name: Check final pip list + run: | + pip list + - name: Prepare weights + run: | + ln -s /root/.cache/models ~/models + # Geo3k + - name: Prepare GEO3K dataset + run: | + python examples/data_preprocess/geo3k.py --local_dataset_path ${HOME}/.cache/datasets/hiyouga/geometry3k + - name: Running GEO3K VLM GRPO E2E training tests on 8 L20 GPUs with rmpad using function rm + run: | + ray stop --force + TRAIN_FILES=$HOME/data/geo3k/train.parquet VAL_FILES=$HOME/data/geo3k/test.parquet \ + MAX_PROMPT_LEN=1536 MAX_RESPONSE_LEN=1536 \ + MODEL_ID=Qwen/Qwen2.5-VL-3B-Instruct \ + ADV_ESTIMATOR=grpo RM_PAD=True USE_KL=True ENABLE_CHUNKED_PREFILL=False \ + SP_SIZE=2 \ + bash tests/special_e2e/ppo_trainer/run_function_reward.sh + - name: Running GEO3K VLM PPO E2E training tests on 8 L20 GPUs with rmpad using function rm + run: | + ray stop --force + TRAIN_FILES=$HOME/data/geo3k/train.parquet VAL_FILES=$HOME/data/geo3k/test.parquet \ + MAX_PROMPT_LEN=1536 MAX_RESPONSE_LEN=1536 \ + MODEL_ID=Qwen/Qwen2.5-VL-3B-Instruct \ + ADV_ESTIMATOR=gae RM_PAD=True USE_KL=True ENABLE_CHUNKED_PREFILL=False \ + SP_SIZE=2 \ + bash tests/special_e2e/ppo_trainer/run_function_reward.sh + - name: Running GEO3K VLM GRPO E2E lora training tests on 8 L20 GPUs with rmpad using function rm + run: | + ray stop --force + TRAIN_FILES=$HOME/data/geo3k/train.parquet VAL_FILES=$HOME/data/geo3k/test.parquet \ + MAX_PROMPT_LEN=1536 MAX_RESPONSE_LEN=1536 \ + MODEL_ID=Qwen/Qwen2.5-VL-3B-Instruct \ + ADV_ESTIMATOR=grpo RM_PAD=True USE_KL=True ENABLE_CHUNKED_PREFILL=False \ + SP_SIZE=2 \ + LORA_RANK=32 LORA_EXCLUDE=".*visual.*" \ + bash tests/special_e2e/ppo_trainer/run_function_reward.sh diff --git a/verl/.github/workflows/e2e_ppo_trainer_veomni_vllm.yml b/verl/.github/workflows/e2e_ppo_trainer_veomni_vllm.yml new file mode 100644 index 0000000000000000000000000000000000000000..39cc6dca7de319e6091569355966d3427f8f630f --- /dev/null +++ b/verl/.github/workflows/e2e_ppo_trainer_veomni_vllm.yml @@ -0,0 +1,153 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + +name: e2e_ppo_trainer_veomni_vllm + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch. + # For push, for now only anti-patterns are specified so it is more conservative + # and achieves higher coverage. + push: + branches: + - main + - v0.* + paths: + - "**/*.py" + # Other entrypoints + - "!verl/trainer/fsdp_sft_trainer.py" + # Megatron + - "!verl/workers/**/megatron_*.py" + pull_request: + branches: + - main + - v0.* + paths: + - "**/*.py" + # Other entrypoints + - "!docker/**" + # Docs + - "!**/*.md" + - "!docs/**" + - "!examples/**" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" + # Megatron + - "!verl/workers/**/megatron_*.py" + # Entrypoints + - ".github/workflows/e2e_ppo_trainer_veomni_vllm.yml" + - "examples/data_preprocess/gsm8k.py" + - "examples/data_preprocess/geo3k.py" + - "tests/special_e2e/run_ppo_trainer_veomni.sh" + - "verl/trainer/main_ppo.py" + - "verl/trainer/config/ppo_trainer.yaml" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +env: + IMAGE: "verl-ci-cn-beijing.cr.volces.com/verlai/verl:vllm018.dev1" + DYNAMIC_RUNNER_ENDPOINT: "https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner" + +jobs: + setup: + if: github.repository_owner == 'verl-project' + runs-on: ubuntu-latest + outputs: + runner-label: ${{ steps.create-runner.outputs.runner-label }} + mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }} + steps: + - uses: actions/checkout@v4 + - id: create-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "create" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-image: "${{ env.IMAGE }}" + + e2e_ppo_trainer_veomni_vllm: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 60 # Increase this timeout value as needed + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + pip3 install veomni==0.1.9a5 --ignore-requires-python --no-deps --index-url https://pypi.org/simple/ + pip3 install transformers==4.57.3 + - name: Prepare GSM8K dataset + run: | + ray stop --force + python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k + - name: Prepare GEO3K dataset + run: | + ray stop --force + python3 examples/data_preprocess/geo3k.py --local_dataset_path ${HOME}/models/hf_data/hiyouga/geometry3k/ + - name: Running GSM8K E2E training tests on 8 L20 GPUs with veomni engine (FSDP_SIZE=4, USP=2) + run: | + ray stop --force + FSDP_SIZE=4 SP_SIZE=2 bash tests/special_e2e/run_ppo_trainer_veomni.sh + - name: Running GEO3K E2E training tests on 8 L20 GPUs with veomni engine (FSDP_SIZE=8, USP=1) + run: | + ray stop --force + MODEL_ID=Qwen/Qwen3-VL-2B-Instruct TRAIN_FILES=${HOME}/data/geo3k/train.parquet VAL_FILES=${HOME}/data/gsm8k/test.parquet FSDP_SIZE=8 SP_SIZE=1 bash tests/special_e2e/run_ppo_trainer_veomni.sh + + cleanup: + runs-on: ubuntu-latest + needs: + [ + setup, + e2e_ppo_trainer_veomni_vllm, + ] + if: always() + steps: + - id: destroy-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "destroy" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-task-id: "${{ needs.setup.outputs.mlp-task-id }}" diff --git a/verl/.github/workflows/e2e_ppo_trainer_veomni_vllm_ascend.yml b/verl/.github/workflows/e2e_ppo_trainer_veomni_vllm_ascend.yml new file mode 100644 index 0000000000000000000000000000000000000000..917649d4bc205c1cfcf907236c3ac6956a9fcfd7 --- /dev/null +++ b/verl/.github/workflows/e2e_ppo_trainer_veomni_vllm_ascend.yml @@ -0,0 +1,129 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + +name: e2e_ppo_trainer_veomni_vllm_ascend + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch + # For push, for now only anti-patterns are specified so it is more conservative + # and achieves higher coverage. + push: + branches: + - main + - v0.* + paths: + - "**/*.py" + - "!**/*.md" + - "!**/*.sh" + # Other entrypoints + - "!examples/*trainer*" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" + # Megatron + - "!verl/workers/**/megatron_*.py" + pull_request: + branches: + - main + - v0.* + paths: + - "**/*.py" + - "!**/*.md" + - "!**/*.sh" + # Other entrypoints + - "!examples/**" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" + # Entrypoints + - ".github/workflows/e2e_ppo_trainer_veomni_vllm_ascend.yml" + - "examples/data_preprocess/gsm8k.py" + - "tests/special_e2e/run_ppo_trainer_veomni.sh" + - "verl/trainer/main_ppo.py" + - "verl/trainer/config/ppo_trainer.yaml" + # Megatron + - "!verl/workers/**/megatron_*.py" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +jobs: + e2e_ppo_trainer_veomni_vllm_ascend: + if: github.repository_owner == 'verl-project' + runs-on: linux-aarch64-a2b3-8 + timeout-minutes: 60 # Increase this timeout value as needed + container: + image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest + options: >- + --shm-size 16g + env: + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + device_name: "npu" + steps: + - name: Check npu and CANN info + run: | + cat /usr/local/Ascend/ascend-toolkit/latest/"$(uname -i)"-linux/ascend_toolkit_install.info + npu-smi info + - name: Check initial pip list from image + run: | + pip list + - name: Checkout verl-project/verl repo + uses: actions/checkout@v4 + with: + fetch-depth: 0 + clean: true + - name: Install the current repository + run: | + pip install -r requirements-npu.txt + pip install --no-deps -e . + pip install veomni==0.1.9a5 --ignore-requires-python --no-deps --index-url https://pypi.org/simple/ + pip install transformers==4.57.3 + - name: Check final pip list + run: | + pip list + - name: Prepare weights + run: | + ln -s /root/.cache/models ~/models + - name: Prepare GSM8K dataset + run: | + python examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k + - name: Running the E2E test with veomni backend + run: | + ray stop --force + bash tests/special_e2e/run_ppo_trainer_veomni.sh diff --git a/verl/.github/workflows/e2e_sft_llm.yml b/verl/.github/workflows/e2e_sft_llm.yml new file mode 100644 index 0000000000000000000000000000000000000000..fc8a3adc7b9c179cc80a3a04e490c17d0ecea93d --- /dev/null +++ b/verl/.github/workflows/e2e_sft_llm.yml @@ -0,0 +1,153 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + +name: e2e_sft_llm + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch + push: + branches: + - main + - v0.* + pull_request: + branches: + - main + - v0.* + paths: + - "**/*.py" + # Other entrypoints + - "!examples/**" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" + + # Megatron + - "!verl/workers/**/megatron_*.py" + # Entrypoints + - ".github/workflows/e2e_sft_llm.yml" + - "examples/data_preprocess/gsm8k.py" + - "tests/special_e2e/sft" + - "verl/trainer/fsdp_sft_trainer.py" + - "verl/trainer/config/sft_trainer.yaml" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +env: + IMAGE: "verl-ci-cn-beijing.cr.volces.com/verlai/verl:sgl059.dev2" + DYNAMIC_RUNNER_ENDPOINT: "https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner" + +jobs: + setup: + if: github.repository_owner == 'verl-project' + runs-on: ubuntu-latest + outputs: + runner-label: ${{ steps.create-runner.outputs.runner-label }} + mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }} + steps: + - uses: actions/checkout@v4 + - id: create-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "create" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-image: "${{ env.IMAGE }}" + e2e_sft_llm: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 30 # Increase this timeout value as needed + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip3 install peft + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + pip3 install veomni==0.1.9a5 --ignore-requires-python --no-deps --index-url https://pypi.org/simple/ + pip3 install transformers==4.57.3 + - name: Prepare gsm8k dataset + run: | + ray stop --force + python3 examples/data_preprocess/gsm8k_multiturn_sft.py --local_dataset_path ${HOME}/models/hf_data/gsm8k + - name: Running GSM8K E2E training tests on 8 L20 GPUs with rmpad using function rm + run: | + ray stop --force + bash tests/special_e2e/sft/run_sft.sh + - name: Running GSM8K E2E training tests on 8 L20 GPUs w/o rmpad using function rm + run: | + ray stop --force + RM_PAD=False bash tests/special_e2e/sft/run_sft.sh + - name: Running GSM8K E2E training tests on 8 L20 GPUs with sequence parallism + run: | + ray stop --force + SP_SIZE=2 bash tests/special_e2e/sft/run_sft.sh + - name: Running GSM8K E2E training tests on 8 L20 GPUs with sequence parallism and liger + run: | + ray stop --force + SP_SIZE=2 LIGER=True bash tests/special_e2e/sft/run_sft.sh + - name: Running GSM8K E2E training tests with LoRA + run: | + ray stop --force + LORA_RANK=32 bash tests/special_e2e/sft/run_sft.sh + - name: Run GSM8K E2E training and resume tests resuming from the checkpoint manager + run: | + ray stop --force + LORA_RANK=32 RESUME_MODE=auto TOTAL_TRAIN_STEP=2 bash tests/special_e2e/sft/run_sft.sh + # TODO: multiturn + - name: Running GSM8K E2E training tests with multiturn and various configs and compare results + run: | + bash tests/special_e2e/sft/test_sft_engine_all.sh + + cleanup: + runs-on: ubuntu-latest + needs: [setup, e2e_sft_llm] + if: always() + steps: + - id: destroy-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "destroy" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-task-id: "${{ needs.setup.outputs.mlp-task-id }}" diff --git a/verl/.github/workflows/e2e_sft_llm_ascend.yml b/verl/.github/workflows/e2e_sft_llm_ascend.yml new file mode 100644 index 0000000000000000000000000000000000000000..b9d2e65465ed1cafe24fef2368438855023b6bab --- /dev/null +++ b/verl/.github/workflows/e2e_sft_llm_ascend.yml @@ -0,0 +1,160 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + +name: e2e_sft_llm_ascend + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch + push: + branches: + - main + - v0.* + pull_request: + branches: + - main + - v0.* + paths: + - "**/*.py" + # Other entrypoints + - "!examples/**" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" + + # Megatron + - "!verl/workers/**/megatron_*.py" + # Entrypoints + - ".github/workflows/e2e_sft_llm_ascend.yml" + - "examples/data_preprocess/gsm8k.py" + - "tests/special_e2e/sft" + - "verl/trainer/fsdp_sft_trainer.py" + - "verl/trainer/config/sft_trainer.yaml" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +jobs: + e2e_sft_llm_ascend: + if: github.repository_owner == 'verl-project' + runs-on: linux-aarch64-a2b3-8 + timeout-minutes: 90 # Increase this timeout value as needed + container: + image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest + options: >- + --shm-size 16g + env: + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - name: Check npu and CANN info + run: | + cat /usr/local/Ascend/ascend-toolkit/latest/"$(uname -i)"-linux/ascend_toolkit_install.info + npu-smi info + - name: Check initial pip list from image + run: | + pip list + - name: Checkout verl-project/verl repo + uses: actions/checkout@v4 + with: + fetch-depth: 0 + clean: true + - name: Install the current repository + run: | + pip install --no-deps -e . + pip install veomni==0.1.9a5 --ignore-requires-python --no-deps --index-url https://pypi.org/simple/ + pip install transformers==4.57.3 + pip install pandas==2.3.3 + pip uninstall -y mbridge + pip install git+https://github.com/ISEEKYAN/mbridge.git@89eb10 + - name: Check final pip list + run: | + pip list + - name: Prepare weights + run: | + ln -s /root/.cache/models ~/models + - name: Prepare gsm8k dataset + run: | + python3 examples/data_preprocess/gsm8k_multiturn_sft.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k + - name: Running GSM8K E2E training tests on 8 NPUs with rmpad using function rm + run: | + ray stop --force + bash tests/special_e2e/sft/run_sft.sh + - name: Running GSM8K E2E training tests on 8 NPUs w/o rmpad using function rm + run: | + ray stop --force + RM_PAD=False bash tests/special_e2e/sft/run_sft.sh + - name: Running GSM8K E2E training tests on 8 NPUs with sequence parallism + run: | + ray stop --force + SP_SIZE=2 bash tests/special_e2e/sft/run_sft.sh + - name: Running GSM8K E2E training tests with LoRA + run: | + ray stop --force + LORA_RANK=32 bash tests/special_e2e/sft/run_sft.sh + - name: Run GSM8K E2E training and resume tests resuming from the checkpoint manager + run: | + ray stop --force + LORA_RANK=32 RESUME_MODE=auto TOTAL_TRAIN_STEP=2 bash tests/special_e2e/sft/run_sft.sh + - name: Running GSM8K E2E training tests with multiturn and various configs and compare results + run: | + ray stop --force + rm -rf ~/verl/test/log + mkdir -p ~/verl/test/log + export VERL_FILE_LOGGER_ROOT=~/verl/test/log + # test with single gpu as golden + echo "run with single gpu as golden" + BACKEND=fsdp SP_SIZE=1 FSDP_SIZE=1 NUM_GPUS=1 FSDP_STRATEGY=fsdp VERL_FILE_LOGGER_PATH=~/verl/test/log/golden.jsonl bash tests/special_e2e/sft/run_sft_engine.sh + # test with fsdp 1 + echo "run with sp2 fsdp_size2 num_gpus8 fsdp_strategy fsdp pad_mode no_padding" + BACKEND=fsdp SP_SIZE=2 FSDP_SIZE=2 NUM_GPUS=8 FSDP_STRATEGY=fsdp PAD_MODE=no_padding bash tests/special_e2e/sft/run_sft_engine.sh + # test with fsdp 1 use_remove_padding and pad_mode no_padding + echo "run with sp4 fsdp_size4 num_gpus8 fsdp_strategy fsdp pad_mode no_padding use_remove_padding False" + BACKEND=fsdp SP_SIZE=1 FSDP_SIZE=-1 NUM_GPUS=8 FSDP_STRATEGY=fsdp PAD_MODE=no_padding USE_REMOVE_PADDING=False bash tests/special_e2e/sft/run_sft_engine.sh + # test with fsdp 2 + echo "run with sp2 fsdp_size2 num_gpus8 fsdp_strategy fsdp2" + BACKEND=fsdp SP_SIZE=2 FSDP_SIZE=2 NUM_GPUS=8 FSDP_STRATEGY=fsdp2 bash tests/special_e2e/sft/run_sft_engine.sh + # test with veomni + echo "run with sp2 fsdp_size4 num_gpus8 fsdp_strategy fsdp2" + BACKEND=veomni SP_SIZE=2 FSDP_SIZE=4 NUM_GPUS=8 FSDP_STRATEGY=fsdp2 bash tests/special_e2e/sft/run_sft_engine.sh + # test with megatron + echo "run with tp2 pp2 vpp2 cp2 num_gpus8" + BACKEND=megatron TP_SIZE=2 PP_SIZE=2 VPP_SIZE=NULL CP_SIZE=2 NUM_GPUS=8 bash tests/special_e2e/sft/run_sft_engine.sh + # test with cp in ray + echo "run with tp2 pp2 vpp2 cp2 num_gpus8 mode=ray" + BACKEND=megatron TP_SIZE=2 PP_SIZE=2 VPP_SIZE=NULL CP_SIZE=2 NUM_GPUS=8 mode=ray bash tests/special_e2e/sft/run_sft_engine.sh + rm -rf ~/verl/test/log diff --git a/verl/.github/workflows/e2e_sft_vlm.yml b/verl/.github/workflows/e2e_sft_vlm.yml new file mode 100644 index 0000000000000000000000000000000000000000..cc8ffd0fb3a78d0f02fcc2a6672eee0c711c78ed --- /dev/null +++ b/verl/.github/workflows/e2e_sft_vlm.yml @@ -0,0 +1,128 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + +name: e2e_sft_vlm + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch + push: + branches: + - main + - v0.* + pull_request: + branches: + - main + - v0.* + paths: + - "**/*.py" + # Other entrypoints + - "!examples/**" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" + + # Megatron + - "!verl/workers/**/megatron_*.py" + # Entrypoints + - ".github/workflows/e2e_sft_vlm.yml" + - "examples/data_preprocess/gsm8k.py" + - "tests/special_e2e/sft" + - "verl/trainer/fsdp_sft_trainer.py" + - "verl/trainer/config/sft_trainer.yaml" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +env: + IMAGE: "verl-ci-cn-beijing.cr.volces.com/verlai/verl:sgl059.dev2" + DYNAMIC_RUNNER_ENDPOINT: "https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner" + +jobs: + setup: + if: github.repository_owner == 'verl-project' + runs-on: ubuntu-latest + outputs: + runner-label: ${{ steps.create-runner.outputs.runner-label }} + mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }} + steps: + - uses: actions/checkout@v4 + - id: create-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "create" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-image: "${{ env.IMAGE }}" + e2e_sft_vlm: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 30 # Increase this timeout value as needed + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip3 install peft + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + pip3 install veomni==0.1.9a5 --ignore-requires-python --no-deps --index-url https://pypi.org/simple/ + pip3 install transformers==4.57.3 + - name: Prepare pokemon-gpt4o-captions dataset + run: | + ray stop --force + python3 examples/data_preprocess/pokemon.py --local_dataset_path ${HOME}/models/hf_data/pokemon-gpt4o-captions + - name: Running Pokemon E2E training tests with multiturn and various configs and compare results + run: | + MODEL_ID=Qwen/Qwen3-VL-2B-Instruct DATASET_DIR=~/data/pokemon-gpt4o-captions VPP_SIZE=null bash tests/special_e2e/sft/test_sft_engine_all.sh + + cleanup: + runs-on: ubuntu-latest + needs: [setup, e2e_sft_vlm] + if: always() + steps: + - id: destroy-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "destroy" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-task-id: "${{ needs.setup.outputs.mlp-task-id }}" diff --git a/verl/.github/workflows/gpu_unit_tests.yml b/verl/.github/workflows/gpu_unit_tests.yml new file mode 100644 index 0000000000000000000000000000000000000000..24be4e2ba874c0fa05f32ba9b502236135f967c5 --- /dev/null +++ b/verl/.github/workflows/gpu_unit_tests.yml @@ -0,0 +1,134 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + +name: GPU unit tests + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch + push: + branches: + - main + - v0.4.x + paths: + - "**/*.py" + - .github/workflows/gpu_unit_tests.yml + pull_request: + branches: + - main + - v0.4.x + paths: + # The order that you define paths patterns matters: + # A matching negative pattern (prefixed with !) after a positive match will exclude the path. + # A matching positive pattern after a negative match will include the path again. + - "**/*.py" + # Other entrypoints + - "!examples/**" + - "!tests/special_sanity/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" + # Entrypoints + - .github/workflows/gpu_unit_tests.yml + - "tests/**test_*.py" + # Ignore CPU tests + - "!tests/*_on_cpu.py" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +env: + IMAGE: "verl-ci-cn-beijing.cr.volces.com/verlai/verl:sgl059.dev2" + DYNAMIC_RUNNER_ENDPOINT: "https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner" + +jobs: + setup: + if: github.repository_owner == 'verl-project' + runs-on: ubuntu-latest + outputs: + runner-label: ${{ steps.create-runner.outputs.runner-label }} + mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }} + steps: + - uses: actions/checkout@v4 + - id: create-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "create" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-image: "${{ env.IMAGE }}" + + gpu_unit_tests: + if: github.repository_owner == 'verl-project' + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 60 # Increase this timeout value as needed + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1" + HF_HUB_ENABLE_HF_TRANSFER: 1 + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip3 install hf_transfer + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + pip3 install cupy-cuda12x==13.6.0 pytest-asyncio + pip3 install --ignore-installed blinker + pip3 install --ignore-installed mlflow "numpy<2.0" + - name: Run all GPU unit tests + run: | + pytest -s -x --ignore-glob="*on_npu.py" --ignore-glob="*test_special_*.py" --ignore-glob='*on_cpu.py' --ignore-glob="*test_vllm*" --ignore-glob="*_sglang*" --ignore-glob="*_hf_rollout*" --ignore-glob="tests/models/" --ignore-glob='tests/special*' --ignore-glob="tests/experimental" --ignore-glob="tests/workers/reward_model" --ignore-glob="*test_shared_memory*" --ignore-glob="tests/workers/rollout/rollout_trtllm" --ignore-glob="*test_bucketed_weight_transfer*" tests/ + - name: Testing LinearCrossEntropyTP Correctness, Computation Time and Memory Consumption + run: | + LOW_MEMORY=True torchrun --standalone --nnodes=1 --nproc-per-node=8 tests/utils/test_special_linear_cross_entropy_tp.py + - name: Testing Megatron KL Loss TP Correctness + run: | + torchrun --standalone --nnodes=1 --nproc-per-node=2 tests/utils/test_special_megatron_kl_loss_tp.py + + cleanup: + runs-on: ubuntu-latest + needs: [setup, gpu_unit_tests] + if: always() + steps: + - id: destroy-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "destroy" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-task-id: "${{ needs.setup.outputs.mlp-task-id }}" diff --git a/verl/.github/workflows/model.yml b/verl/.github/workflows/model.yml new file mode 100644 index 0000000000000000000000000000000000000000..70c16fe34e43cf0f4581f664f755e1fb2e4c5e8a --- /dev/null +++ b/verl/.github/workflows/model.yml @@ -0,0 +1,183 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. +# name: Check PR Title + +name: model + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch + push: + branches: + - main + - v0.* + pull_request: + branches: + - main + - v0.* + paths: + - "verl/**/*.py" + # Entrypoints + - ".github/workflows/model.yml" + - "tests/special_distributed/test_fsdp_ckpt.py" + - "tests/special_distributed/test_tensor_dict.py" + - "tests/models/**" + - "tests/special_distributed/run_all.sh" + +# Declare permissions just read content. +permissions: + contents: read + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +env: + IMAGE: "verl-ci-cn-beijing.cr.volces.com/verlai/verl:vllm018.dev1" + DYNAMIC_RUNNER_ENDPOINT: "https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner" + +jobs: + setup: + if: github.repository_owner == 'verl-project' + runs-on: ubuntu-latest + outputs: + runner-label: ${{ steps.create-runner.outputs.runner-label }} + mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }} + steps: + - uses: actions/checkout@v4 + - id: create-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "create" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-image: "${{ env.IMAGE }}" + + model_rmpad: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 20 # Increase this timeout value as needed + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository and upgrade to latest transformers + run: | + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + - name: Running rmpad model tests on 8 L20 GPUs + flash_attn 2.5.8 + run: | + pytest -s tests/models/test_transformer.py + - name: Running rmpad model tests on 8 L20 GPUs + latest flash_attn + run: | + pytest -s tests/models/test_transformer.py + - name: Running FSDP rmpad model tests on 8 L20 GPUs + latest flash_attn + run: | + STRATEGY=fsdp torchrun --nproc_per_node=8 tests/special_distributed/test_fsdp_ckpt.py + - name: Running transformers ulysses tests on 8 L20 GPUs + latest transformers + run: | + torchrun --nproc_per_node=8 -m pytest tests/models/test_transformers_ulysses.py + - name: Running transformers ulysses tests on 8 L20 GPUs + transformers 4.54.1 + run: | + pip3 install transformers==4.54.1 + torchrun --nproc_per_node=8 -m pytest tests/models/test_transformers_ulysses.py + - name: Run distributed test + run: | + bash tests/special_distributed/run_all.sh + - name: Clean up and recover transformers + run: | + pip3 install --upgrade "transformers==5.3.0" + + # TODO: Move this back to model_rmpad once FSDP2 is stable. + # NOTE: List as an independent job to make rerun easier. + model_rmpad_fsdp2_unstable: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 20 # Increase this timeout value as needed + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository and upgrade to latest transformers/flash_attn + run: | + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + - name: Running FSDP2 rmpad model tests on 8 L20 GPUs + latest flash_attn + run: | + STRATEGY=fsdp2 torchrun --nproc_per_node=8 tests/special_distributed/test_fsdp_ckpt.py + + model_engine: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 20 # Increase this timeout value as needed + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + pip3 install --upgrade "accelerate>=1.13.0" + - name: Download model config files + run: | + hf download Qwen/Qwen2.5-0.5B-Instruct --local-dir $HOME/models/Qwen/Qwen2.5-0.5B-Instruct + - name: Running mcore engine tests on 8 L20 GPUs + run: | + ray stop --force + pytest -s -x tests/models/test_engine.py + + cleanup: + runs-on: ubuntu-latest + needs: [setup, model_rmpad, model_rmpad_fsdp2_unstable, model_engine] + if: always() + steps: + - id: destroy-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "destroy" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-task-id: "${{ needs.setup.outputs.mlp-task-id }}" diff --git a/verl/.github/workflows/model_ascend.yml b/verl/.github/workflows/model_ascend.yml new file mode 100644 index 0000000000000000000000000000000000000000..f797cca5c9422010b6454eed2c1e23e1d5efdf30 --- /dev/null +++ b/verl/.github/workflows/model_ascend.yml @@ -0,0 +1,135 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. +# name: Check PR Title + +name: model_ascend + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch + push: + branches: + - main + - v0.* + pull_request: + branches: + - main + - v0.* + paths: + - "verl/**/*.py" + # Entrypoints + - ".github/workflows/model_ascend.yml" + - "tests/special_distributed/test_fsdp_ckpt.py" + - "tests/special_distributed/test_tensor_dict.py" + - "tests/models/**" + - "tests/special_distributed/run_all.sh" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +permissions: + contents: read + +jobs: + model_rmpad_ascend: + if: github.repository_owner == 'verl-project' + runs-on: linux-aarch64-a2b3-8 + timeout-minutes: 60 # Increase this timeout value as needed + container: + image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest + options: >- + --shm-size 16g + env: + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - name: Check npu and CANN info + run: | + cat /usr/local/Ascend/ascend-toolkit/latest/"$(uname -i)"-linux/ascend_toolkit_install.info + npu-smi info + - name: Check initial pip list from image + run: | + pip list + - name: Checkout verl-project/verl repo + uses: actions/checkout@v4 + with: + fetch-depth: 0 + clean: true + - name: Install the current repository + run: | + pip install --no-deps -e .[test] + - name: Check final pip list + run: | + pip list + - name: Prepare weights + run: | + ln -s /root/.cache/models ~/models + - name: Running rmpad model tests on 8 NPUs + run: | + pytest -s tests/models/test_transformer.py + - name: Running FSDP rmpad model tests on 8 NPUs + run: | + STRATEGY=fsdp torchrun --nproc_per_node=8 tests/special_distributed/test_fsdp_ckpt.py + - name: Running transformers ulysses tests on 8 NPUs + run: | + torchrun --nproc_per_node=8 -m pytest tests/models/test_transformers_ulysses.py + - name: Run distributed test + run: | + bash tests/special_distributed/run_all.sh + + # TODO: Move this back to model_rmpad once FSDP2 is stable. + # NOTE: List as an independent job to make rerun easier. + model_rmpad_fsdp2_unstable_ascend: + if: github.repository_owner == 'verl-project' + runs-on: linux-aarch64-a2b3-8 + timeout-minutes: 60 + container: + image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest + options: >- + --shm-size 16g + env: + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip install --no-deps -e .[test] + - name: Prepare weights + run: | + ln -s /root/.cache/models ~/models + - name: Running FSDP2 rmpad model tests on 8 NPUs + run: | + STRATEGY=fsdp2 torchrun --nproc_per_node=8 tests/special_distributed/test_fsdp_ckpt.py diff --git a/verl/.github/workflows/nightly_ascend.yml b/verl/.github/workflows/nightly_ascend.yml new file mode 100644 index 0000000000000000000000000000000000000000..cc059e48d1cad84b15dad8b065437ddc5fc3a3d4 --- /dev/null +++ b/verl/.github/workflows/nightly_ascend.yml @@ -0,0 +1,260 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + +name: nightly_ci_ascend + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch + # For push, for now only anti-patterns are specified so it is more conservative + # and achieves higher coverage. + schedule: + - cron: "0 17 * * *" + +# Declare permissions just read content. +permissions: + contents: read + +jobs: + # Test ppo qwen3-8b fsdp+vllm + nightlyCI_ppo-qwen3-8b-fsdp-vllm_ascend: + if: github.repository_owner == 'verl-project' + runs-on: linux-aarch64-a2b3-8 + timeout-minutes: 180 # Increase this timeout value as needed + container: + image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest + options: >- + --shm-size 16g + env: + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - name: Check npu and CANN info + run: | + cat /usr/local/Ascend/ascend-toolkit/latest/"$(uname -i)"-linux/ascend_toolkit_install.info + npu-smi info + - name: Check initial pip list from image + run: | + pip list + - name: Checkout verl-project/verl repo + uses: actions/checkout@v4 + with: + fetch-depth: 0 + clean: true + - name: Install the current repository + run: | + pip install --no-deps -e . + - name: Check final pip list + run: | + pip list + - name: Prepare weights + run: | + ln -s /root/.cache/models ~/models + - name: Prepare GSM8K dataset + run: | + python examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k + - name: Running nightlyCI_ppo-qwen3-8b-fsdp-vllm_ascend + run: | + ray stop --force + bash tests/special_npu/nightly_ci_ascend/run_ppo_qwen3-8b_fsdp_npu.sh + + # Test grpo qwen25-7b-Instruct fsdp+vllm + nightlyCI_grpo-qwen25-7b-Instruct-fsdp-vllm_ascend: + if: github.repository_owner == 'verl-project' + runs-on: linux-aarch64-a2b3-8 + timeout-minutes: 180 # Increase this timeout value as needed + container: + image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest + options: >- + --shm-size 16g + env: + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - name: Check npu and CANN info + run: | + cat /usr/local/Ascend/ascend-toolkit/latest/"$(uname -i)"-linux/ascend_toolkit_install.info + npu-smi info + - name: Check initial pip list from image + run: | + pip list + - name: Checkout verl-project/verl repo + uses: actions/checkout@v4 + with: + fetch-depth: 0 + clean: true + - name: Install the current repository + run: | + pip install --no-deps -e . + - name: Check final pip list + run: | + pip list + - name: Prepare weights + run: | + ln -s /root/.cache/models ~/models + - name: Prepare GSM8K dataset + run: | + python examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k + - name: Running nightlyCI_grpo-qwen25-7b-Instruct-fsdp-vllm_ascend + run: | + ray stop --force + bash tests/special_npu/nightly_ci_ascend/run_grpo_qwen25-7b-instruct_fsdp_npu.sh + + # Test grpo qwen25-vl-3b-Instruct fsdp+vllm + nightlyCI_grpo-qwen25-vl-3b-Instruct-fsdp-vllm_ascend: + if: github.repository_owner == 'verl-project' + runs-on: linux-aarch64-a2b3-8 + timeout-minutes: 180 # Increase this timeout value as needed + container: + image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest + options: >- + --shm-size 16g + env: + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - name: Check npu and CANN info + run: | + cat /usr/local/Ascend/ascend-toolkit/latest/"$(uname -i)"-linux/ascend_toolkit_install.info + npu-smi info + - name: Check initial pip list from image + run: | + pip list + - name: Checkout verl-project/verl repo + uses: actions/checkout@v4 + with: + fetch-depth: 0 + clean: true + - name: Install the current repository + run: | + pip install --no-deps -e . + - name: Check final pip list + run: | + pip list + - name: Prepare weights + run: | + ln -s /root/.cache/models ~/models + - name: Preprocess geo3k dataset + run: | + python examples/data_preprocess/geo3k.py --local_dataset_path ${HOME}/.cache/datasets/hiyouga/geometry3k + - name: Running nightlyCI_grpo-qwen25-vl-3b-Instruct-fsdp-vllm_ascend + run: | + ray stop --force + bash tests/special_npu/nightly_ci_ascend/run_grpo_qwen25-vl-3b-instruct_fsdp_npu.sh + + # Test dapo moonlight-16b megatron vllm + nightlyCI_dapo-moonlight-16b-megatron-vllm_ascend: + if: github.repository_owner == 'verl-project' + runs-on: linux-aarch64-a2b3-8 + timeout-minutes: 180 # Increase this timeout value as needed + container: + image: swr.ap-southeast-1.myhuaweicloud.com/base_image/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest + options: >- + --shm-size 16g + env: + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - name: Check npu and CANN info + run: | + cat /usr/local/Ascend/ascend-toolkit/latest/"$(uname -i)"-linux/ascend_toolkit_install.info + npu-smi info + - name: Check initial pip list from image + run: | + pip list + - name: Checkout verl-project/verl repo + uses: actions/checkout@v4 + with: + fetch-depth: 0 + submodules: recursive + clean: true + - name: Install the current repository + run: | + pip install -r requirements-npu.txt + pip install --no-deps -e . + - name: Check final pip list + run: | + pip list + - name: Prepare weights + run: | + ln -s /root/.cache/models ~/models + - name: Preprocess geo3k dataset + run: | + python examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k + - name: update mbridge + run: | + # get mbridge path + MBRIDGE_PATH=$(pip show mbridge | grep Location | awk '{print $2}') + # cuda to npu + TARGET_FILE="${MBRIDGE_PATH}/mbridge/models/ext/deepseek_v3/dequant_fp8_safetensor_io.py" + sed -i '34s/cuda/npu/;51s/cuda/npu/' "$TARGET_FILE" + - name: Running nightlyCI_dapo-moonlight-16b-megatron-vllm_ascend + run: | + ray stop --force + export HCCL_OP_EXPANSION_MODE="AIV" + bash tests/special_npu/nightly_ci_ascend/run_dapo_moonlight-16b_megatron_npu.sh + + # Test gspo qwen3-30b megatron vllm + nightlyCI_gspo-qwen3-30b-megatron-vllm_ascend: + if: github.repository_owner == 'verl-project' + runs-on: linux-aarch64-a3-16 + timeout-minutes: 180 # Increase this timeout value as needed + container: + image: swr.ap-southeast-1.myhuaweicloud.com/base_image/ascend-ci/verl/verl:verl-8.5.0-a3-ubuntu22.04-py3.11-latest + options: >- + --shm-size 60g + env: + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - name: Check npu and CANN info + run: | + cat /usr/local/Ascend/ascend-toolkit/latest/"$(uname -i)"-linux/ascend_toolkit_install.info + npu-smi info + - name: Check initial pip list from image + run: | + pip list + - name: Checkout verl-project/verl repo + uses: actions/checkout@v4 + with: + fetch-depth: 0 + submodules: recursive + clean: true + - name: Prepare weights + run: | + ln -s /root/.cache/models ~/models + - name: Preprocess geo3k dataset + run: | + python examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k + - name: Running nightlyCI_gspo-qwen3-30b-megatron-vllm_ascend + run: | + ray stop --force + bash tests/special_npu/nightly_ci_ascend/run_gspo_qwen3_30b_megatron_npu.sh diff --git a/verl/.github/workflows/npu_unit_tests.yml b/verl/.github/workflows/npu_unit_tests.yml new file mode 100644 index 0000000000000000000000000000000000000000..8c9f13669f5c4ce2cb1e9e792bb22cc9f8c9fea6 --- /dev/null +++ b/verl/.github/workflows/npu_unit_tests.yml @@ -0,0 +1,120 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - `npu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix on ascend device. +# - Since cpu/gpu/npu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + +name: NPU unit tests + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch + push: + branches: + - main + - v0.* + paths: + - "**/*.py" + - .github/workflows/npu_unit_tests.yml + pull_request: + branches: + - main + paths: + # The order that you define paths patterns matters: + # A matching negative pattern (prefixed with !) after a positive match will exclude the path. + # A matching positive pattern after a negative match will include the path again. + - "**/*.py" + # Other entrypoints + - "!examples/**" + - "!tests/special_sanity/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" + - "!recipe/**" + # Entrypoints + - .github/workflows/npu_unit_tests.yml + - "tests/**test_*.py" + # Ignore CPU tests + - "!tests/*_on_cpu.py" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +jobs: + npu_unit_tests: + if: github.repository_owner == 'verl-project' + runs-on: linux-aarch64-a2b3-8 + timeout-minutes: 60 # Increase this timeout value as needed + container: + image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest + options: >- + --shm-size 16g + env: + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - name: Check npu and CANN info + run: | + cat /usr/local/Ascend/ascend-toolkit/latest/"$(uname -i)"-linux/ascend_toolkit_install.info + npu-smi info + - name: Check initial pip list from image + run: | + pip list + - name: Checkout verl-project/verl repo + uses: actions/checkout@v4 + with: + fetch-depth: 0 + clean: true + - name: Install the current repository + run: | + pip install --no-deps -e .[test] + pip install mlflow pytest-asyncio + - name: Check final pip list + run: | + pip list + - name: Prepare weights + run: | + ln -s /root/.cache/models ~/models + - name: Run all NPU unit tests + run: | + pytest -s -x --ignore-glob="*test_special_*.py" --ignore-glob="*on_cpu.py" --ignore-glob="*test_vllm*" --ignore-glob="*_sglang*" --ignore-glob="*_hf_rollout*" --ignore-glob="tests/models/" --ignore-glob="tests/special*" --ignore-glob="tests/experimental" --ignore-glob="tests/workers/reward_model" --ignore-glob="*test_rvdz*" --ignore-glob="*test_ray_collectives*" --ignore-glob="*test_nvtx_profile*" --ignore-glob="tests/checkpoint_engine" --ignore-glob="*test_shared_memory*" --ignore-glob="tests/workers/rollout/rollout_trtllm" --ignore-glob="*test_fsdp_lora_merge*" --ignore-glob="*test_activation_offload*" --ignore-glob="*test_normalize_peft_param_name.py*" tests/ + - name: Testing activation offload + run: | + pytest -s -x tests/utils/test_activation_offload.py + - name: Testing normalize peft param name + run: | + pytest -s -x tests/utils/test_normalize_peft_param_name.py + - name: Running NPU profiling unit tests + run: | + pytest -s -x tests/utils/test_special_mstx_profile.py diff --git a/verl/.github/workflows/pre-commit.yml b/verl/.github/workflows/pre-commit.yml new file mode 100644 index 0000000000000000000000000000000000000000..4f6aa4bdf0d10048057777269df9507188b3a264 --- /dev/null +++ b/verl/.github/workflows/pre-commit.yml @@ -0,0 +1,41 @@ +# c.f. https://github.com/pre-commit/action?tab=readme-ov-file#using-this-action +name: pre-commit + +# No need to avoid / cancel lightweight pre-commit jobs +on: + schedule: + - cron: "0 0 * * 0" + pull_request: + push: + branches: + - main + - v0.* + # Allow manual triggering + workflow_dispatch: + +# Declare permissions just read content. +permissions: + contents: read + +jobs: + pre-commit: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.12"] + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0 + with: + python-version: ${{ matrix.python-version }} + - name: Install the current repository + run: | + pip install pre-commit hydra-core + pip install --no-deps -e . + - name: Set ruff --output-format=github + run: | + sed -i 's/--output-format=full/--output-format=github/' .pre-commit-config.yaml + git add .pre-commit-config.yaml + # Check "--all-files" by default + - uses: pre-commit/action@v3.0.1 diff --git a/verl/.github/workflows/precommit-autofix.yml b/verl/.github/workflows/precommit-autofix.yml new file mode 100644 index 0000000000000000000000000000000000000000..d235da90cd21ce429559eccf962fdb4eee7a5252 --- /dev/null +++ b/verl/.github/workflows/precommit-autofix.yml @@ -0,0 +1,52 @@ +name: scheduled pre-commit autofix + +on: + schedule: + # Every hour + - cron: "0 * * * *" + workflow_dispatch: + +permissions: + contents: write + pull-requests: write + +jobs: + precommit: + if: github.repository_owner == 'verl-project' + runs-on: ubuntu-latest + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.10" + + - name: Install pre-commit + run: | + python -m pip install --upgrade pip + pip install pre-commit hydra-core + + - name: Run pre-commit + run: | + pre-commit run --all-files || true + + - name: Create or update PR + uses: peter-evans/create-pull-request@v6 + with: + branch: bot/precommit-autofix + delete-branch: true + title: "[ci] chore: scheduled pre-commit autofix" + commit-message: "chore: auto-fix pre-commit issues" + body: | + This PR was created automatically by a scheduled GitHub Action. + + - Runs `pre-commit run --all-files` + - Triggered hourly + labels: | + automated + pre-commit diff --git a/verl/.github/workflows/reward_model_sglang.yml b/verl/.github/workflows/reward_model_sglang.yml new file mode 100644 index 0000000000000000000000000000000000000000..965996e9c4340a3e98ab6ed92908fe4dac742db5 --- /dev/null +++ b/verl/.github/workflows/reward_model_sglang.yml @@ -0,0 +1,130 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. +# name: Check PR Title + +name: reward_model_sglang + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch + push: + branches: + - main + - v0.* + pull_request: + branches: + - main + - v0.* + paths: + - "verl/**/*.py" + # Entrypoints + - ".github/workflows/reward_model_sglang.yml" + - "tests/experimental/reward_loop/**" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +env: + IMAGE: "verl-ci-cn-beijing.cr.volces.com/verlai/verl:sgl059.dev2" + DYNAMIC_RUNNER_ENDPOINT: "https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner" + +jobs: + setup: + if: github.repository_owner == 'verl-project' + runs-on: ubuntu-latest + outputs: + runner-label: ${{ steps.create-runner.outputs.runner-label }} + mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }} + steps: + - uses: actions/checkout@v4 + - id: create-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "create" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-image: "${{ env.IMAGE }}" + + reward_model_sglang: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 30 # Increase this timeout value as needed + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK: "True" + NCCL_SHM_DISABLE: "1" + NCCL_P2P_DISABLE: "1" + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + pip3 install sglang-router==0.2.2 + - name: Prepare gsm8k dataset + run: | + ray stop --force + python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k --local_dir ${HOME}/data/gsm8k + - name: Running sglang generative reward model tests on 8 L20 GPUs + run: | + ROLLOUT_NAME=sglang pytest -s -x tests/experimental/reward_loop/test_reward_model_genrm.py + - name: Running sglang discriminative reward model tests on 8 L20 GPUs + run: | + ROLLOUT_NAME=sglang pytest -s -x tests/experimental/reward_loop/test_reward_model_disrm.py + # FIXME(@yyDing1): broken + # - name: Running sglang agent loop with reward manager tests on 8 L20 GPUs + # run: | + # ROLLOUT_NAME=sglang pytest -s -x tests/experimental/reward_loop/test_agent_reward_loop_standalone.py + # - name: Running sglang agent loop with reward model colocate tests on 8 L20 GPUs + # run: | + # ROLLOUT_NAME=sglang pytest -s -x tests/experimental/reward_loop/test_agent_reward_loop_colocate.py + + cleanup: + runs-on: ubuntu-latest + needs: [setup, reward_model_sglang] + if: always() + steps: + - id: destroy-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "destroy" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-task-id: "${{ needs.setup.outputs.mlp-task-id }}" diff --git a/verl/.github/workflows/reward_model_sglang_ascend.yml b/verl/.github/workflows/reward_model_sglang_ascend.yml new file mode 100644 index 0000000000000000000000000000000000000000..e143336d110c4c5ba9217f1b53e49c053b55bf4e --- /dev/null +++ b/verl/.github/workflows/reward_model_sglang_ascend.yml @@ -0,0 +1,117 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. +# name: Check PR Title + +name: reward_model_sglang_ascend + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch + push: + branches: + - main + - v0.* + pull_request: + branches: + - main + - v0.* + paths: + - "verl/**/*.py" + # Entrypoints + - ".github/workflows/reward_model_sglang_ascend.yml" + - "tests/experimental/reward_loop/**" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +jobs: + reward_model_sglang_ascend: + if: github.repository_owner == 'verl-project' + runs-on: linux-aarch64-a3-8 + timeout-minutes: 60 # Increase this timeout value as needed + container: + image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-sglang-8.5.0-a3-ubuntu22.04-py3.11-latest + options: >- + --shm-size 16g + env: + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES: "1" + ASCEND_RT_VISIBLE_DEVICES: "0,1,2,3,4,5,6,7" + HCCL_HOST_SOCKET_PORT_RANGE: "60000-60050" + HCCL_NPU_SOCKET_PORT_RANGE: "61000-61050" + steps: + - name: Check npu and CANN info + run: | + cat /usr/local/Ascend/ascend-toolkit/latest/"$(uname -i)"-linux/ascend_toolkit_install.info + npu-smi info + - name: Check initial pip list from image + run: | + pip list + - name: Checkout verl-project/verl repo + uses: actions/checkout@v4 + with: + fetch-depth: 0 + clean: true + - name: Install the current repository + run: | + pip3 install pytest pytest-asyncio + pip3 install --no-deps --no-build-isolation -e . + pip3 install sglang-router==0.2.2 + - name: Check final pip list + run: | + pip list + - name: Prepare weights + run: | + ln -s /root/.cache/models ~/models + - name: Prepare gsm8k dataset + run: | + ray stop --force + python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k --local_dir ${HOME}/data/gsm8k + - name: Running sglang generative reward model tests on 8 NPUs + run: | + ROLLOUT_NAME=sglang pytest -s -x tests/experimental/reward_loop/test_reward_model_genrm.py + - name: Running sglang discriminative reward model tests on 8 NPUs + run: | + ROLLOUT_NAME=sglang pytest -s -x tests/experimental/reward_loop/test_reward_model_disrm.py + # FIXME(@yyDing1): broken + # - name: Running sglang agent loop with reward manager tests on 8 NPUs + # run: | + # ROLLOUT_NAME=sglang pytest -s -x tests/experimental/reward_loop/test_agent_reward_loop_standalone.py + # - name: Running sglang agent loop with reward model colocate tests on 8 NPUs + # run: | + # ROLLOUT_NAME=sglang pytest -s -x tests/experimental/reward_loop/test_agent_reward_loop_colocate.py diff --git a/verl/.github/workflows/reward_model_vllm.yml b/verl/.github/workflows/reward_model_vllm.yml new file mode 100644 index 0000000000000000000000000000000000000000..278246a5e2d2ef2c9efb963a0c7213437162b87f --- /dev/null +++ b/verl/.github/workflows/reward_model_vllm.yml @@ -0,0 +1,129 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. +# name: Check PR Title + +name: reward_model_vllm + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch + push: + branches: + - main + - v0.* + pull_request: + branches: + - main + - v0.* + paths: + - "verl/**/*.py" + # Entrypoints + - ".github/workflows/reward_model_vllm.yml" + - "tests/experimental/reward_loop/**" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +env: + IMAGE: "verl-ci-cn-beijing.cr.volces.com/verlai/verl:vllm018.dev1" + DYNAMIC_RUNNER_ENDPOINT: "https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner" + +jobs: + setup: + if: github.repository_owner == 'verl-project' + runs-on: ubuntu-latest + outputs: + runner-label: ${{ steps.create-runner.outputs.runner-label }} + mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }} + steps: + - uses: actions/checkout@v4 + - id: create-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "create" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-image: "${{ env.IMAGE }}" + + reward_model_vllm: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 30 # Increase this timeout value as needed + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK: "True" + NCCL_SHM_DISABLE: "1" + NCCL_P2P_DISABLE: "1" + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + - name: Prepare gsm8k dataset + run: | + ray stop --force + python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k --local_dir ${HOME}/data/gsm8k + - name: Running vllm generative reward model tests on 8 L20 GPUs + run: | + ROLLOUT_NAME=vllm pytest -s -x tests/experimental/reward_loop/test_reward_model_genrm.py + - name: Running vllm discriminative reward model tests on 8 L20 GPUs + run: | + ROLLOUT_NAME=vllm pytest -s -x tests/experimental/reward_loop/test_reward_model_disrm.py + + - name: Running vllm agent loop with reward manager tests on 8 L20 GPUs + run: | + ROLLOUT_NAME=vllm pytest -s -x tests/experimental/reward_loop/test_agent_reward_loop_standalone.py + - name: Running vllm agent loop with reward model colocate tests on 8 L20 GPUs + run: | + ROLLOUT_NAME=vllm pytest -s -x tests/experimental/reward_loop/test_agent_reward_loop_colocate.py + + cleanup: + runs-on: ubuntu-latest + needs: [setup, reward_model_vllm] + if: always() + steps: + - id: destroy-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "destroy" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-task-id: "${{ needs.setup.outputs.mlp-task-id }}" diff --git a/verl/.github/workflows/reward_model_vllm_ascend.yml b/verl/.github/workflows/reward_model_vllm_ascend.yml new file mode 100644 index 0000000000000000000000000000000000000000..60507dddf3aeb8c90a70c9434d6ee7b007fa225f --- /dev/null +++ b/verl/.github/workflows/reward_model_vllm_ascend.yml @@ -0,0 +1,112 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. +# name: Check PR Title + +name: reward_model_vllm_ascend + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch + push: + branches: + - main + - v0.* + pull_request: + branches: + - main + - v0.* + paths: + - "verl/**/*.py" + # Entrypoints + - ".github/workflows/reward_model_vllm_ascend.yml" + - "tests/experimental/reward_loop/**" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +jobs: + reward_model_vllm_ascend: + if: github.repository_owner == 'verl-project' + runs-on: linux-aarch64-a2b3-8 + timeout-minutes: 60 # Increase this timeout value as needed + container: + image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest + options: >- + --shm-size 16g + env: + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - name: Check npu and CANN info + run: | + cat /usr/local/Ascend/ascend-toolkit/latest/"$(uname -i)"-linux/ascend_toolkit_install.info + npu-smi info + - name: Check initial pip list from image + run: | + pip list + - name: Checkout verl-project/verl repo + uses: actions/checkout@v4 + with: + fetch-depth: 0 + clean: true + - name: Install the current repository + run: | + pip install --no-deps -e .[test] + - name: Check final pip list + run: | + pip list + - name: Prepare weights + run: | + ln -s /root/.cache/models ~/models + - name: Prepare gsm8k dataset + run: | + ray stop --force + python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k --local_dir ${HOME}/data/gsm8k + - name: Running vllm generative reward model tests on 8 NPUs + run: | + ROLLOUT_NAME=vllm pytest -s -x tests/experimental/reward_loop/test_reward_model_genrm.py + - name: Running vllm discriminative reward model tests on 8 NPUs + run: | + ROLLOUT_NAME=vllm pytest -s -x tests/experimental/reward_loop/test_reward_model_disrm.py + - name: Running vllm agent loop with reward manager tests on 8 NPUs + run: | + ROLLOUT_NAME=vllm pytest -s -x tests/experimental/reward_loop/test_agent_reward_loop_standalone.py + - name: Running vllm agent loop with reward model colocate tests on 8 NPUs + run: | + export HCCL_HOST_SOCKET_PORT_RANGE=auto + export HCCL_NPU_SOCKET_PORT_RANGE=auto + ROLLOUT_NAME=vllm pytest -s -x tests/experimental/reward_loop/test_agent_reward_loop_colocate.py diff --git a/verl/.github/workflows/sanity.yml b/verl/.github/workflows/sanity.yml new file mode 100644 index 0000000000000000000000000000000000000000..b9ec384f400318c69961a100241264b36ea41014 --- /dev/null +++ b/verl/.github/workflows/sanity.yml @@ -0,0 +1,86 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. +# name: Check PR Title + +name: sanity + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch + push: + branches: + - main + - v0.* + pull_request: + branches: + - main + - v0.* + paths: + - "**/*.py" + - .github/workflows/sanity.yml + - "tests/special_sanity/**" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +jobs: + sanity: + runs-on: ubuntu-latest + timeout-minutes: 5 # Increase this timeout value as needed + strategy: + matrix: + python-version: ["3.10"] + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0 + with: + python-version: ${{ matrix.python-version }} + - name: Install the current repository + run: | + pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cpu + pip3 install -r requirements.txt + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + # Most sanity checks now run via pre-commit (see .pre-commit-config.yaml + # and .github/workflows/pre-commit.yml). Only checks that need the full + # verl installation or CI-only context remain here. + - name: Run sanity test + run: | + pytest -s -x tests/special_sanity + - name: Assert documentation requirement for functions + run: python3 tests/special_sanity/validate_imported_docs.py diff --git a/verl/.github/workflows/scorecard.yml b/verl/.github/workflows/scorecard.yml new file mode 100644 index 0000000000000000000000000000000000000000..176d15ae2bd470752daaf138fa5aaa90641738e9 --- /dev/null +++ b/verl/.github/workflows/scorecard.yml @@ -0,0 +1,66 @@ +# This workflow uses actions that are not certified by GitHub. They are provided +# by a third-party and are governed by separate terms of service, privacy +# policy, and support documentation. + +name: Scorecard supply-chain security +on: + # For Branch-Protection check. Only the default branch is supported. See + # https://github.com/ossf/scorecard/blob/main/docs/checks.md#branch-protection + branch_protection_rule: + # To guarantee Maintained check is occasionally updated. See + # https://github.com/ossf/scorecard/blob/main/docs/checks.md#maintained + schedule: + - cron: "27 7 * * 1" + push: + branches: + - main + - v0.* + +# Declare default permissions as read only. +permissions: read-all + +jobs: + analysis: + name: Scorecard analysis + runs-on: ubuntu-latest + permissions: + # Needed to upload the results to code-scanning dashboard. + security-events: write + # Needed to publish results and get a badge (see publish_results below). + id-token: write + # Uncomment the permissions below if installing in a private repository. + # contents: read + # actions: read + + steps: + - name: "Checkout code" + uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # v4.1.1 + with: + persist-credentials: false + + - name: "Run analysis" + uses: ossf/scorecard-action@0864cf19026789058feabb7e87baa5f140aac736 # v2.3.1 + with: + results_file: results.sarif + results_format: sarif + # (Optional) "write" PAT token. Uncomment the `repo_token` line below if: + # - you want to enable the Branch-Protection check on a *public* repository, or + # - you are installing Scorecard on a *private* repository + # To create the PAT, follow the steps in https://github.com/ossf/scorecard-action?tab=readme-ov-file#authentication-with-fine-grained-pat-optional. + # repo_token: ${{ secrets.SCORECARD_TOKEN }} + + # Public repositories: + # - Publish results to OpenSSF REST API for easy access by consumers + # - Allows the repository to include the Scorecard badge. + # - See https://github.com/ossf/scorecard-action#publishing-results. + # For private repositories: + # - `publish_results` will always be set to `false`, regardless + # of the value entered here. + publish_results: true + + # Upload the results to GitHub's code scanning dashboard (optional). + # Commenting out will disable upload of results to your repo's Code Scanning dashboard + - name: "Upload to code-scanning" + uses: github/codeql-action/upload-sarif@9e8d0789d4a0fa9ceb6b1738f7e269594bdd67f0 #v3.28.9 + with: + sarif_file: results.sarif diff --git a/verl/.github/workflows/secrets_scan.yml b/verl/.github/workflows/secrets_scan.yml new file mode 100644 index 0000000000000000000000000000000000000000..298ed16c668c67facdb6af2119878da576f5bdf5 --- /dev/null +++ b/verl/.github/workflows/secrets_scan.yml @@ -0,0 +1,22 @@ +on: + push: + branches: + - main + - v0.* + pull_request: + +permissions: + contents: read + +jobs: + test: + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # v4.1.1 + with: + fetch-depth: 0 + - name: Secret Scanning + uses: trufflesecurity/trufflehog@7dc056a193116ba8d82154bf0549381c8fb8545c # v3.88.14 + with: + extra_args: --results=verified,unknown diff --git a/verl/.github/workflows/sgl.yml b/verl/.github/workflows/sgl.yml new file mode 100644 index 0000000000000000000000000000000000000000..a0a653b86859b094bc0c61e66617e3483cd0db1f --- /dev/null +++ b/verl/.github/workflows/sgl.yml @@ -0,0 +1,162 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + +name: sgl + +on: + # workflow_dispatch: # Manual + # Trigger the workflow on push or pull request, + # but only for the main branch + push: + branches: + - main + - v0.* + paths: + - "**/*.py" + - .github/workflows/sgl.yml + pull_request: + branches: + - main + - v0.* + paths: + - "**/*.py" + # Other entrypoints + - "!examples/**" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" # FSDP + - "!verl/workers/**/*dp_*.py" + # Megatron + - "!verl/workers/**/megatron_*.py" + # vLLM + - "!**/*vllm*" + + # Entrypoints + - ".github/workflows/sgl.yml" + - "tests/rollout/*sglang*" + - "tests/rollout/async_rollout_utils.py" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +env: + IMAGE: "verl-ci-cn-beijing.cr.volces.com/verlai/verl:sgl059.dev2" + DYNAMIC_RUNNER_ENDPOINT: "https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner" + +jobs: + setup: + if: github.repository_owner == 'verl-project' + runs-on: ubuntu-latest + outputs: + runner-label: ${{ steps.create-runner.outputs.runner-label }} + mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }} + steps: + - uses: actions/checkout@v4 + - id: create-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "create" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-image: "${{ env.IMAGE }}" + + sgl: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 35 # Increase this timeout value as needed + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: 1 + SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK: "True" + NCCL_SHM_DISABLE: "1" + NCCL_P2P_DISABLE: "1" + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip3 install cupy-cuda12x==13.6.0 pytest-asyncio + pip3 install hf_transfer pytest-asyncio + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + - name: Prepare gsm8k dataset + run: | + ray stop --force + python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k + - name: Test the latest SGLang Rollout async with agent loop + run: | + ROLLOUT_NAME=sglang pytest -svvv tests/experimental/agent_loop + + sgl_checkpoint_engine: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 35 # Increase this timeout value as needed + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: 1 + SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK: "True" + NCCL_SHM_DISABLE: "1" + NCCL_P2P_DISABLE: "1" + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip3 install cupy-cuda12x==13.6.0 pytest-asyncio + pip3 install hf_transfer pytest-asyncio + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + - name: Test SGLang ServerAdapter with Checkpoint Engine (NCCL) + run: | + ROLLOUT_NAME=sglang pytest -svvv tests/checkpoint_engine/test_special_server_adapter.py + + cleanup: + runs-on: ubuntu-latest + needs: [setup, sgl, sgl_checkpoint_engine] + if: always() + steps: + - id: destroy-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "destroy" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-task-id: "${{ needs.setup.outputs.mlp-task-id }}" diff --git a/verl/.github/workflows/sgl_ascend.yml b/verl/.github/workflows/sgl_ascend.yml new file mode 100644 index 0000000000000000000000000000000000000000..acfdc5a8e1ee7812a244b390a4964850f714fd40 --- /dev/null +++ b/verl/.github/workflows/sgl_ascend.yml @@ -0,0 +1,126 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + +name: sglang_ascend + +on: + # workflow_dispatch: # Manual + # Trigger the workflow on push or pull request, + # but only for the main branch + push: + branches: + - main + - v0.* + paths: + - "**/*.py" + - .github/workflows/sgl_ascend.yml + pull_request: + branches: + - main + - v0.* + paths: + - "**/*.py" + # Other entrypoints + - "!examples/**" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" # FSDP + - "!verl/workers/**/*dp_*.py" + # Megatron + - "!verl/workers/**/megatron_*.py" + # vLLM + - "!**/*vllm*" + + # Entrypoints + - ".github/workflows/sgl_ascend.yml" + - "tests/rollout/*sglang*" + - "tests/rollout/async_rollout_utils.py" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +jobs: + sglang_ascend: + if: github.repository_owner == 'verl-project' + runs-on: linux-aarch64-a3-8 + timeout-minutes: 90 # Increase this timeout value as needed + container: + image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-sglang-8.5.0-a3-ubuntu22.04-py3.11-latest + options: >- + --shm-size 16g + env: + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES: "1" + ASCEND_RT_VISIBLE_DEVICES: "0,1,2,3,4,5,6,7" + HCCL_HOST_SOCKET_PORT_RANGE: "60000-60050" + HCCL_NPU_SOCKET_PORT_RANGE: "61000-61050" + ENGINE: sglang + steps: + - name: Check npu and CANN info + run: | + cat /usr/local/Ascend/ascend-toolkit/latest/"$(uname -i)"-linux/ascend_toolkit_install.info + npu-smi info + - name: Check initial pip list from image + run: | + pip list + - name: Checkout verl-project/verl repo + uses: actions/checkout@v4 + with: + fetch-depth: 0 + clean: true + - name: Install the current repository + run: | + pip3 install pytest pytest-asyncio + pip3 install megatron-bridge --no-deps + pip3 install git+https://github.com/ISEEKYAN/mbridge.git@main --no-deps --no-build-isolation + pip3 install --no-deps --no-build-isolation -e . + - name: Check final pip list + run: | + pip list + - name: Prepare weights + run: | + ln -s /root/.cache/models ~/models + - name: Prepare gsm8k dataset + run: | + python examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k + - name: Test the latest vLLM Rollout async with agent loop + run: | + export HCCL_HOST_SOCKET_PORT_RANGE=auto + export HCCL_NPU_SOCKET_PORT_RANGE=auto + export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 + ROLLOUT_NAME=sglang pytest -svvv --ignore=tests/experimental/agent_loop/test_multi_modal.py -k "not test_standalone_rollout[4]" tests/experimental/agent_loop diff --git a/verl/.github/workflows/type-coverage-check.yml b/verl/.github/workflows/type-coverage-check.yml new file mode 100644 index 0000000000000000000000000000000000000000..268f0c672f0f87e8437d7dc964ee464922ec5d4e --- /dev/null +++ b/verl/.github/workflows/type-coverage-check.yml @@ -0,0 +1,31 @@ +name: Type Annotation and Docstring Coverage + +on: + pull_request: + paths: + - '**/*.py' + - '.github/workflows/type-coverage-check.yml' + +jobs: + type-coverage-check: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 # 🚨 Important: fetch full history so `origin/main` is available + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.10' + + - name: Install dependencies + run: | + pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cpu + pip3 install -r requirements.txt + pip3 install --no-deps -e . + - name: Run type annotation coverage check + run: | + python3 tests/special_sanity/type_coverage_check.py + - name: Run docstring coverage check + run: | + python3 tests/special_sanity/check_api_docs.py verl diff --git a/verl/.github/workflows/vllm.yml b/verl/.github/workflows/vllm.yml new file mode 100644 index 0000000000000000000000000000000000000000..b68fca01147a58e20307291a02f585322fcefc21 --- /dev/null +++ b/verl/.github/workflows/vllm.yml @@ -0,0 +1,165 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + +name: vllm + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch + push: + branches: + - main + - v0.* + pull_request: + branches: + - main + - v0.* + paths: + - "**/*.py" + # Other entrypoints + - "!examples/**" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" + # FSDP + - "!verl/workers/**/*dp_*.py" + # Megatron + - "!verl/workers/**/megatron_*.py" + # SGLang + - "!**/*sglang*" + # Entrypoints + - ".github/workflows/vllm.yml" + - "tests/special_e2e/generation" + - "tests/workers/rollout" + - "verl/trainer/main_generation.py" + - "verl/trainer/config/generation.yaml" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +env: + IMAGE: "verl-ci-cn-beijing.cr.volces.com/verlai/verl:vllm018.dev1" + DYNAMIC_RUNNER_ENDPOINT: "https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner" + +jobs: + setup: + if: github.repository_owner == 'verl-project' + runs-on: ubuntu-latest + outputs: + runner-label: ${{ steps.create-runner.outputs.runner-label }} + mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }} + steps: + - uses: actions/checkout@v4 + - id: create-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "create" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-image: "${{ env.IMAGE }}" + + vllm: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 35 # Increase this timeout value as needed + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + # - name: Download Model to Use + # run: | + # hf download Qwen/Qwen2.5-0.5B-Instruct --local-dir ${HOME}/models/Qwen/Qwen2.5-0.5B-Instruct + # hf download Qwen/Qwen2.5-1.5B-Instruct --local-dir ${HOME}/models/Qwen/Qwen2.5-1.5B-Instruct + # hf download Qwen/Qwen2.5-VL-3B-Instruct --local-dir ${HOME}/models/Qwen/Qwen2.5-VL-3B-Instruct + # hf download OldKingMeister/Qwen2.5-1.5B-Instruct-YaRN --local-dir ${HOME}/models/OldKingMeister/Qwen2.5-1.5B-Instruct-YaRN + # export HF_HUB_OFFLINE=1 + - name: Prepare gsm8k dataset + run: | + ray stop --force + python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k + - name: Test the latest vLLM Rollout async with agent loop + run: | + ROLLOUT_NAME=vllm pytest -svvv tests/experimental/agent_loop + - name: Test vllm server abort functionality + run: | + pytest tests/workers/rollout/rollout_vllm/test_vllm_abort.py -v -s + + vllm_checkpoint_engine: + needs: setup + runs-on: ["${{ needs.setup.outputs.runner-label || 'L20x8' }}"] + timeout-minutes: 35 # Increase this timeout value as needed + env: + HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }} + NO_PROXY: "localhost,127.0.0.1,hf-mirror.com" + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + fetch-depth: 0 + - name: Install the current repository + run: | + pip3 install pytest-asyncio + pip3 install -r requirements-test.txt + pip3 install --no-deps -e . + pip3 install cupy-cuda12x==13.6.0 + - name: Test vLLM ServerAdapter with Checkpoint Engine (NCCL) + run: | + ROLLOUT_NAME=vllm pytest -svvv tests/checkpoint_engine/test_special_server_adapter.py + - name: Test bucketed weight transfer + run: | + pytest -svvv tests/utils/test_bucketed_weight_transfer.py + + cleanup: + runs-on: ubuntu-latest + needs: [setup, vllm, vllm_checkpoint_engine] + if: always() + steps: + - id: destroy-runner + uses: volcengine/vemlp-github-runner@v1 + with: + mode: "destroy" + faas-url: "${{ env.DYNAMIC_RUNNER_ENDPOINT }}" + mlp-task-id: "${{ needs.setup.outputs.mlp-task-id }}" diff --git a/verl/.github/workflows/vllm_ascend.yml b/verl/.github/workflows/vllm_ascend.yml new file mode 100644 index 0000000000000000000000000000000000000000..50ad7745d877fb2891c6275d1574c562856f67aa --- /dev/null +++ b/verl/.github/workflows/vllm_ascend.yml @@ -0,0 +1,121 @@ +# # Tests layout + +# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +# - `tests/trainer` for testing functionality related to `verl/trainer` +# - `tests/models` for testing functionality related to `verl/models` +# - ... + +# There are a few folders with `special_` prefix, created for special purposes: +# - `special_distributed`: unit tests that must run with multiple GPUs +# - `special_e2e`: end-to-end tests with training/generation scripts +# - `special_npu`: tests for NPUs +# - `special_sanity`: a suite of quick sanity tests +# - `special_standalone`: a set of test that are designed to run in dedicated environments + +# Accelerators for tests +# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# # Workflow layout + +# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +# 3. End-to-end tests: `e2e_*.yml` +# 4. Unit tests +# - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` +# - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. +# - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when +# - new workflow yaml is added to `.github/workflows` +# - new tests are added to workflow mentioned in 2. + +name: vllm_ascend + +on: + # Trigger the workflow on push or pull request, + # but only for the main branch + push: + branches: + - main + - v0.* + pull_request: + branches: + - main + - v0.* + paths: + - "**/*.py" + # Other entrypoints + - "!examples/**" + - "!tests/**" + - "!verl/trainer/main_*.py" + - "!verl/trainer/fsdp_sft_trainer.py" + # FSDP + - "!verl/workers/**/*dp_*.py" + # Megatron + - "!verl/workers/**/megatron_*.py" + # SGLang + - "!**/*sglang*" + # Entrypoints + - ".github/workflows/vllm_ascend.yml" + - "tests/special_e2e/generation" + - "tests/workers/rollout" + - "verl/trainer/main_generation.py" + - "verl/trainer/config/generation.yaml" + +# Cancel jobs on the same ref if a new one is triggered +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +# Declare permissions just read content. +permissions: + contents: read + +jobs: + vllm_ascend: + if: github.repository_owner == 'verl-project' + runs-on: linux-aarch64-a2b3-8 + timeout-minutes: 60 # Increase this timeout value as needed + container: + image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest + options: >- + --shm-size 16g + env: + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable + steps: + - name: Check npu and CANN info + run: | + cat /usr/local/Ascend/ascend-toolkit/latest/"$(uname -i)"-linux/ascend_toolkit_install.info + npu-smi info + - name: Check initial pip list from image + run: | + pip list + - name: Checkout verl-project/verl repo + uses: actions/checkout@v4 + with: + fetch-depth: 0 + clean: true + - name: Install the current repository + run: | + pip3 install --no-deps -e .[test] + pip3 install pytest-asyncio + - name: Check final pip list + run: | + pip list + - name: Prepare weights + run: | + ln -s /root/.cache/models ~/models + - name: Prepare gsm8k dataset + run: | + ray stop --force + python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k + - name: Test the latest vLLM Rollout async with agent loop + run: | + export HCCL_HOST_SOCKET_PORT_RANGE=auto + export HCCL_NPU_SOCKET_PORT_RANGE=auto + export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 + ROLLOUT_NAME=vllm pytest -svvv tests/experimental/agent_loop + - name: Test vllm server abort functionality + run: | + pytest tests/workers/rollout/rollout_vllm/test_vllm_abort.py -v -s diff --git a/verl/.gitignore b/verl/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..b4d8d2f9c2b57c71277ae10ec11777312a42e452 --- /dev/null +++ b/verl/.gitignore @@ -0,0 +1,133 @@ +**/*.pt +**/checkpoints +**/wget-log +**/_build/ +**/*.ckpt +**/outputs +**/*.tar.gz +**/playground +**/wandb + +/pyrightconfig.json + +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class +dataset/* +tensorflow/my_graph/* +.idea/ +# C extensions +*.so + +# Distribution / packaging +.Python +# env/ +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +tmp/ +*.egg-info/ +.installed.cfg +*.egg + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*,cover +.hypothesis/ +pytest.ini +output.txt + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +# IPython Notebook +.ipynb_checkpoints + +# pyenv +.python-version + +# celery beat schedule file +celerybeat-schedule + +# dotenv +.env + +# virtualenv +venv/ +.venv/ +ENV/ +!**/workers/env/ + +# Spyder project settings +.spyderproject + +# Rope project settings +.ropeproject + +# vscode +.vscode + +# Mac +.DS_Store + +# vim +*.swp + +# emacs +*~ + +# ckpt +*.lock + +# data +*.parquet + + +# local logs +logs +log +outputs +.history \ No newline at end of file diff --git a/verl/.gitmodules b/verl/.gitmodules new file mode 100644 index 0000000000000000000000000000000000000000..d5dd7a6aa577ccb64650ca389b699e04fd7af259 --- /dev/null +++ b/verl/.gitmodules @@ -0,0 +1,3 @@ +[submodule "recipe"] + path = recipe + url = https://github.com/verl-project/verl-recipe.git diff --git a/verl/.pre-commit-config.yaml b/verl/.pre-commit-config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..39c0910c348ee43e903055044c9a6edd7d6e7556 --- /dev/null +++ b/verl/.pre-commit-config.yaml @@ -0,0 +1,75 @@ +repos: + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: "v0.12.2" + hooks: + - id: ruff + args: ["--fix", "--show-fixes", "--output-format=full"] + exclude: ^.*\.(ipynb)$ + - id: ruff-format + + - repo: https://github.com/pre-commit/mirrors-mypy + rev: "v1.17.0" + hooks: + - id: mypy + + - repo: local + hooks: + - id: autogen-trainer-cfg + name: Generate and verify verl/trainer/config/_generated_*.yaml + entry: scripts/generate_trainer_config.sh + language: script + pass_filenames: false + + - id: check-docs-time-info + name: Check docs Last updated info + entry: python3 tests/special_sanity/check_docs_time_info.py + language: python + pass_filenames: false + + - id: check-docstrings + name: Check doc string coverage + entry: python3 tests/special_sanity/check_docstrings.py + language: python + pass_filenames: false + + - id: check-license + name: Check license + entry: python3 tests/special_sanity/check_license.py --directories . + language: python + pass_filenames: false + + - id: check-device-api-usage + name: Check device API usage + entry: python3 tests/special_sanity/check_device_api_usage.py --directory ./verl + language: python + pass_filenames: false + + - id: check-dataproto-usage + name: Check DataProto usage + entry: python3 tests/special_sanity/check_dataproto_usage.py -d ./verl/workers/engine + language: python + pass_filenames: false + + - id: validate-structure + name: Validate test structure + entry: python3 tests/special_sanity/validate_structure.py + language: python + pass_filenames: false + + - id: check-naming-conventions + name: Check naming conventions + entry: sh -c 'fail=0; if grep -rIn --exclude-dir=.git --exclude-dir=.github --exclude-dir=venv --exclude-dir=__pycache__ --exclude=.pre-commit-config.yaml "veRL" .; then echo "Please use verl instead of veRL"; fail=1; fi; if grep -rIn --exclude-dir=.git --exclude-dir=.github --exclude-dir=venv --exclude-dir=__pycache__ --exclude=ascend_sglang_best_practices.rst --exclude=.pre-commit-config.yaml -E "Sglang|sgLang|sglAng|sglaNg|sglanG" .; then echo "Please use SGLang or sglang"; fail=1; fi; exit $fail' + language: system + pass_filenames: false + + - id: check-example-naming + name: Check example script naming convention + entry: python3 tests/special_sanity/check_example_naming.py --root examples + language: python + pass_filenames: false + + - id: compileall + name: Compile all python files + entry: sh -c 'PYTHONWARNINGS=error python3 -m compileall -q . -x "(^|[\\/])(\.venv|venv|\.git|verl[\\/]experimental[\\/]vla)([\\/]|$)"' + language: python + pass_filenames: false diff --git a/verl/.readthedocs.yaml b/verl/.readthedocs.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0016868541a2a0667ef40ae6a9d861bcd26b9316 --- /dev/null +++ b/verl/.readthedocs.yaml @@ -0,0 +1,19 @@ +# Read the Docs configuration file +# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details + +version: 2 + +build: + os: ubuntu-22.04 + tools: + python: "3.11" + rust: "1.70" + +sphinx: + configuration: docs/conf.py + +python: + install: + - requirements: docs/requirements-docs.txt + - method: pip + path: . diff --git a/verl/.vscode/settings.json b/verl/.vscode/settings.json new file mode 100644 index 0000000000000000000000000000000000000000..705533538d75025d58f108b0d5ae1ec7b5a470b5 --- /dev/null +++ b/verl/.vscode/settings.json @@ -0,0 +1,15 @@ +{ + "[python]": { + "editor.defaultFormatter": "charliermarsh.ruff", + "editor.codeActionsOnSave": { + "source.organizeImports": "always", + } + }, + "files.associations": { + "array": "cpp", + "string_view": "cpp", + "initializer_list": "cpp", + "utility": "cpp" + }, + "iis.configDir": "" +} \ No newline at end of file diff --git a/verl/AGENTS.md b/verl/AGENTS.md new file mode 100644 index 0000000000000000000000000000000000000000..f44706f26edd31c83f1102dc829e8faeac90ea87 --- /dev/null +++ b/verl/AGENTS.md @@ -0,0 +1,87 @@ +# Agent Instructions for verl + +> These instructions apply to **all** AI-assisted contributions to `verl-project/verl`. +> Breaching these guidelines can result in automatic banning. + +## 1. Contribution Policy (Mandatory) + +### Duplicate-work checks + +Before proposing a PR, run these checks: + +```bash +gh issue view --repo verl-project/verl --comments +gh pr list --repo verl-project/verl --state open --search " in:body" +gh pr list --repo verl-project/verl --state open --search "" +``` + +- If an open PR already addresses the same fix, do not open another. +- If your approach is materially different, explain the difference in the issue. + +### No low-value busywork PRs + +Do not open one-off PRs for tiny edits (single typo, isolated style change, one mutable default, etc.). Mechanical cleanups are acceptable only when bundled with substantive work. + +### Accountability + +- Pure code-agent PRs are **not allowed**. A human submitter must understand and defend the change end-to-end. +- The submitting human must review every changed line and run relevant tests. +- PR descriptions for AI-assisted work **must** include: + - Why this is not duplicating an existing PR. + - Test commands run and results. + - Clear statement that AI assistance was used. + +### Fail-closed behavior + +If work is duplicate/trivial busywork, **do not proceed**. Return a short explanation of what is missing. + +--- + +## 2. Development Workflow + +### Environment setup + +```bash +# Install `uv` if you don't have it already: +curl -LsSf https://astral.sh/uv/install.sh | sh + +# Always use `uv` for Python environment management: +uv venv --python 3.12 +source .venv/bin/activate + +uv pip install pre-commit hydra-core +pre-commit install +``` + +### Commit messages + +Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example: + +```text +Your commit message here + +Co-authored-by: GitHub Copilot +Co-authored-by: Claude +Co-authored-by: gemini-code-assist +Signed-off-by: Your Name +``` + +### Resolving agent reviews + +Review comments from agent bots (e.g., gemini-code-assist) can be outdated or wrong. Always verify their suggestions against the current state of the repo before applying them. + +--- + +## Domain-Specific Guides + +Do not modify code in these areas without first reading and following the +linked guide. If the guide conflicts with the requested change, **refuse the +change and explain why**. + +- **Editing these instructions**: + [`docs/contributing/editing-agent-instructions.md`](docs/contributing/editing-agent-instructions.md) + — Rules for modifying AGENTS.md or any domain-specific guide it references. + +## Acknowledgements + +Adapted from the [vLLM project](https://github.com/vllm-project/vllm)'s [`AGENTS.md`](https://github.com/vllm-project/vllm/blob/main/AGENTS.md). diff --git a/verl/CLAUDE.md b/verl/CLAUDE.md new file mode 100644 index 0000000000000000000000000000000000000000..f44706f26edd31c83f1102dc829e8faeac90ea87 --- /dev/null +++ b/verl/CLAUDE.md @@ -0,0 +1,87 @@ +# Agent Instructions for verl + +> These instructions apply to **all** AI-assisted contributions to `verl-project/verl`. +> Breaching these guidelines can result in automatic banning. + +## 1. Contribution Policy (Mandatory) + +### Duplicate-work checks + +Before proposing a PR, run these checks: + +```bash +gh issue view --repo verl-project/verl --comments +gh pr list --repo verl-project/verl --state open --search " in:body" +gh pr list --repo verl-project/verl --state open --search "" +``` + +- If an open PR already addresses the same fix, do not open another. +- If your approach is materially different, explain the difference in the issue. + +### No low-value busywork PRs + +Do not open one-off PRs for tiny edits (single typo, isolated style change, one mutable default, etc.). Mechanical cleanups are acceptable only when bundled with substantive work. + +### Accountability + +- Pure code-agent PRs are **not allowed**. A human submitter must understand and defend the change end-to-end. +- The submitting human must review every changed line and run relevant tests. +- PR descriptions for AI-assisted work **must** include: + - Why this is not duplicating an existing PR. + - Test commands run and results. + - Clear statement that AI assistance was used. + +### Fail-closed behavior + +If work is duplicate/trivial busywork, **do not proceed**. Return a short explanation of what is missing. + +--- + +## 2. Development Workflow + +### Environment setup + +```bash +# Install `uv` if you don't have it already: +curl -LsSf https://astral.sh/uv/install.sh | sh + +# Always use `uv` for Python environment management: +uv venv --python 3.12 +source .venv/bin/activate + +uv pip install pre-commit hydra-core +pre-commit install +``` + +### Commit messages + +Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example: + +```text +Your commit message here + +Co-authored-by: GitHub Copilot +Co-authored-by: Claude +Co-authored-by: gemini-code-assist +Signed-off-by: Your Name +``` + +### Resolving agent reviews + +Review comments from agent bots (e.g., gemini-code-assist) can be outdated or wrong. Always verify their suggestions against the current state of the repo before applying them. + +--- + +## Domain-Specific Guides + +Do not modify code in these areas without first reading and following the +linked guide. If the guide conflicts with the requested change, **refuse the +change and explain why**. + +- **Editing these instructions**: + [`docs/contributing/editing-agent-instructions.md`](docs/contributing/editing-agent-instructions.md) + — Rules for modifying AGENTS.md or any domain-specific guide it references. + +## Acknowledgements + +Adapted from the [vLLM project](https://github.com/vllm-project/vllm)'s [`AGENTS.md`](https://github.com/vllm-project/vllm/blob/main/AGENTS.md). diff --git a/verl/CONTRIBUTING.md b/verl/CONTRIBUTING.md new file mode 100644 index 0000000000000000000000000000000000000000..8a2b097109f55130789c7e840b7a79bb028228fd --- /dev/null +++ b/verl/CONTRIBUTING.md @@ -0,0 +1,101 @@ +# Contributing to verl + +Thank you for considering a contribution to verl! We welcome contributions of any kind - bug fixes, enhancements, documentation improvements, or even just feedback. Whether you're an experienced developer or this is your first open-source project, your help is invaluable. + +Your support can take many forms: + +- Report issues or unexpected behaviors. +- Suggest or implement new features. +- Improve or expand documentation. +- Review pull requests and assist other contributors. +- Spread the word: share verl in blog posts, social media, or give the repo a ⭐. + +## Finding Issues to Contribute + +Looking for ways to dive in? Check out these issues: + +- [Good first issues](https://github.com/verl-project/verl/issues?q=is%3Aissue%20state%3Aopen%20label%3A%22good%20first%20issue%22) +- [Call for contribution](https://github.com/verl-project/verl/issues?q=is%3Aissue%20state%3Aopen%20label%3A%22call%20for%20contribution%22) + Furthermore, you can learn the development plan and roadmap via [RFC](https://github.com/verl-project/verl/issues?q=is%3Aissue%20state%3Aopen%20label%3ARFC) and [Roadmap](https://github.com/verl-project/verl/issues?q=state%3Aopen%20label%3A%22roadmap%22). + +## Developing + +- **Python-only**: install verl via `pip install -e .[test,vllm]` or `pip install -e .[test,sglang]` and iterate quickly. For full dependency setup, check out the verl [installation doc](https://verl.readthedocs.io/en/latest/start/install.html). + +## Code Linting and Formatting + +We rely on pre-commit to keep our code consistent. To set it up: + +```bash +pip install pre-commit hydra-core +pre-commit install +# for staged changes +pre-commit run +# for all files in the repo +pre-commit run --all-files +# run a specific hook with pre-commit +# pre-commit run --all-files --show-diff-on-failure --color=always +pre-commit run --all-files --show-diff-on-failure --color=always ruff +pre-commit run --all-files --show-diff-on-failure --color=always autogen-trainer-cfg +``` + +## Testing + +Our test suites run on GitHub Actions. Check these workflows for details: + +- [GPU unit tests](https://github.com/verl-project/verl/blob/main/.github/workflows/gpu_unit_tests.yml) +- [CPU unit tests](https://github.com/verl-project/verl/blob/main/.github/workflows/cpu_unit_tests.yml) +- [vLLM tests](https://github.com/verl-project/verl/blob/main/.github/workflows/vllm.yml) +- [SGLang tests](https://github.com/verl-project/verl/blob/main/.github/workflows/sgl.yml) + +### Adding CI tests + +If possible, please add CI test(s) for your new feature: + +1. Find the most relevant workflow yml file, which usually corresponds to a `hydra` default config (e.g. `ppo_trainer`, `ppo_megatron_trainer`, `sft_trainer`, etc). +2. Add related path patterns to the `paths` section if not already included. +3. Minimize the workload of the test script(s) (see existing scripts for examples). + +## Building the Docs + +``` +# Ensure verl is on your PYTHONPATH, e.g.: +pip install -e .[test] + +# Install documentation dependencies +cd docs +pip install -r requirements-docs.txt + +# Generate HTML docs +make clean +make html + +# Preview locally +python -m http.server -d _build/html/ +``` + +Open your browser at http://localhost:8000 to explore the docs. + +## Pull Requests & Code Reviews + +Thanks for submitting a PR! To streamline reviews: + +- Follow our Pull Request Template for title format and checklist. +- Adhere to our pre-commit lint rules and ensure all checks pass. +- Update docs for any user-facing changes. +- Add or update tests in the CI workflows, or explain why tests aren't applicable. + +## AI-Assisted Contributions + +See + +- [`AGENTS.md`](AGENTS.md) for rules that all AI coding agents must follow +- [`editing-agent-instructions.md`](docs/contributing/editing-agent-instructions.md) for guidelines on editing agent instructions. + +## License + +See the [LICENSE](https://github.com/verl-project/verl/blob/main/LICENSE) file for full details. + +## Thank You + +We appreciate your contributions to verl. Your efforts help make the project stronger and more user-friendly. Happy coding! diff --git a/verl/LICENSE b/verl/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..d645695673349e3947e8e5ae42332d0ac3164cd7 --- /dev/null +++ b/verl/LICENSE @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/verl/Notice.txt b/verl/Notice.txt new file mode 100644 index 0000000000000000000000000000000000000000..ade439da525ac3f82936e131a1ae386f43207fd8 --- /dev/null +++ b/verl/Notice.txt @@ -0,0 +1 @@ +Copyright 2023-2024 Bytedance Ltd. and/or its affiliates \ No newline at end of file diff --git a/verl/README.md b/verl/README.md new file mode 100644 index 0000000000000000000000000000000000000000..0f841ea6b5402fa877e1307046dcb66e279fa62f --- /dev/null +++ b/verl/README.md @@ -0,0 +1,307 @@ +
+ 👋 Hi, everyone! + verl is a RL training library initiated by ByteDance Seed team and maintained by the verl community. +
+
+
+ +
+ +Ask DeepWiki.com +[![GitHub Repo stars](https://img.shields.io/github/stars/verl-project/verl)](https://github.com/verl-project/verl/stargazers) +[![Twitter](https://img.shields.io/twitter/follow/verl_project)](https://twitter.com/verl_project) + + +[![Documentation](https://img.shields.io/badge/documentation-blue)](https://verl.readthedocs.io/en/latest/) + + +
+ +![seed logo](https://github.com/user-attachments/assets/c42e675e-497c-4508-8bb9-093ad4d1f216) + +

verl: Volcano Engine Reinforcement Learning for LLMs

+ +verl is a flexible, efficient and production-ready RL training library for large language models (LLMs). + +verl is the open-source version of **[HybridFlow: A Flexible and Efficient RLHF Framework](https://arxiv.org/abs/2409.19256v2)** paper. + +verl is flexible and easy to use with: + +- **Easy extension of diverse RL algorithms**: The hybrid-controller programming model enables flexible representation and efficient execution of complex post-training dataflows. Build RL dataflows such as GRPO, PPO in a few lines of code. + +- **Seamless integration of existing LLM infra with modular APIs**: Decouples computation and data dependencies, enabling seamless integration with existing LLM frameworks, such as FSDP, Megatron-LM, vLLM, SGLang, etc + +- **Flexible device mapping**: Supports various placement of models onto different sets of GPUs for efficient resource utilization and scalability across different cluster sizes. + +- Ready integration with popular HuggingFace models + +verl is fast with: + +- **State-of-the-art throughput**: SOTA LLM training and inference engine integrations and SOTA RL throughput. + +- **Efficient actor model resharding with 3D-HybridEngine**: Eliminates memory redundancy and significantly reduces communication overhead during transitions between training and generation phases. + +
+ verl-arch.png +
+ +

+ +## News +- [2026/05] [VeRL-Omni](https://github.com/verl-project/verl-omni) is pre-released: a unified RL stack for diffusion and omni-modal model post-training built on top of verl. Read the [blog post](https://vllm.ai/blog/2026-05-14-verl-omni) for details. +- [2026/05] verl's zero-mismatch HuggingFace rollout [vexact](https://github.com/verl-project/vexact) is released: with batch-invariant kernels, shared model definition with FSDP, and out-of-box examples compatible with VeOmni. +- [2026/04] verl's Megatron backend LoRA and router replay support is showcased at [PyTorch Conference Europe 2026](https://pytorchconferenceeu2026.sched.com/event/2Juce/optimizing-reinforcement-learning-at-trillion-parameter-scale-songlin-jiang-aalto-university-mind-lab). +- [2026/03] verl is presented at NVIDIA GTC26: [session#1](https://www.nvidia.com/en-us/on-demand/session/gtc26-S81829/), [session#2](https://www.nvidia.com/en-us/on-demand/session/gtc26-S81620/) +- [2026/01] verl has been migrated to the [verl-project](https://github.com/verl-project) +- [2026/01] verl first meetup was successfully held in Shanghai on 01/10, hosted by Volcengine and NVIDIA, the slides has been uploaded to [verl-data](https://github.com/verl-project/verl-data). +- [2026/01] The `recipe` directory has been migrated to a dedicated repository: [verl-recipe](https://github.com/verl-project/verl-recipe) and added as a submodule. See https://github.com/verl-project/verl/pull/4795. It can be used as it was after `git submodule update --init --recursive recipe`. Note that [`transfer_queue`](verl/experimental/transfer_queue), [`fully_async_policy`](verl/experimental/fully_async_policy), [`one_step_off_policy`](verl/experimental/one_step_off_policy) and [`vla`](verl/experimental/vla) are kept under [`verl/experimental`](verl/experimental) since they are planned to be merged into the main library. Use them through `verl.experimental.{module}`. +- [2025/12] [Mind Lab](https://macaron.im/mindlab) successfully used [verl](https://github.com/verl-project/verl) and [Megatron-bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge) to train GRPO Lora for Trillion-parameter model on 64 H800 - See their [techblog](https://macaron.im/mindlab/research/building-trillion-parameter-reasoning-rl-with-10-gpus). +- [2025/10] verl is presented in the [PyTorch Conference 2025](https://pytorch.org/event/pytorch-conference-2025/). +- [2025/08] verl is presented in the [PyTorch Expert Exchange Webinar](https://www.youtube.com/watch?v=Vd79NmmqY3Q&t=2s). [Slides](https://github.com/eric-haibin-lin/verl-community/blob/main/slides/verl_talk_pytorch_2025_08.pdf) available. +- [2025/07] The [ReTool](https://arxiv.org/pdf/2504.11536) recipe is fully open sourced. [Blog](https://www.notion.so/verl-reTool-recipe-Using-multi-round-conversations-and-code-sandboxing-to-improve-the-math-of-large-23a8b5b7feba80b386b2e5b5e3c1cde0) +- [2025/07] The first verl meetup will be held at ICML Vancouver on July 16th! Please [join us](https://lu.ma/0ek2nyao) if you are at ICML! (onsite only) +- [2025/06] verl with Megatron backend enables large MoE models such as [DeepSeek-671B and Qwen3-235B](https://verl.readthedocs.io/en/latest/perf/dpsk.html). +- [2025/03] [DAPO](https://dapo-sia.github.io/) is the open-sourced SOTA RL algorithm that achieves 50 points on AIME 2024 based on the Qwen2.5-32B pre-trained model, surpassing the previous SOTA achieved by DeepSeek's GRPO (DeepSeek-R1-Zero-Qwen-32B). DAPO's training is fully powered by verl and the reproduction code is available in `recipe/dapo` now. +
more... +
    +
  • [2025/04] [Seed-Thinking-v1.5](https://github.com/ByteDance-Seed/Seed-Thinking-v1.5/blob/main/seed-thinking-v1.5.pdf) tech report is released! Trained with verl, Seed-Thinking-v1.5 achieves 86.7 on AIME 2024, 55.0 on Codeforces and 77.3 on GPQA, demonstrating excellent reasoning abilities in STEM and coding. Beyond reasoning tasks, the method demonstrates notable generalization across diverse domains.
    • +
  • [2025/07] verl keynote at [AWS AI Hours Singapore](https://pages.awscloud.com/aws-ai-hours-sg.html#agenda) on 7/8, verl & verl-agent project updates at [Agent for SWE meetup](https://lu.ma/e498qhsi) by LF AI & Data Singapore on 7/11.
    • +
  • [2025/06] verl team will provide latest project updates at [PyTorch Day China](https://www.lfasiallc.com/pytorch-day-china/) on June 7th. Meet our dev team in Beijing!
    • +
  • [2025/04] [VAPO](https://arxiv.org/pdf/2504.05118) (value-based augmented PPO) paper covers our latest RL method for reasoning models. Trained from Qwen-32B-base model, VAPO achieves 60.4 on AIME 2024, outperforming DAPO-32B.
    • +
  • [2025/05] [PF-PPO](https://arxiv.org/abs/2409.06957), accepted to ICML 2025, is now supported in verl! PF-PPO enhances policy learning efficiency and robustness by filtering potentially noisy reward signals and reusing high-quality experiences via a replay buffer.
    • +
  • [2025/04] We will give a tutorial about latest post-training techniques and programming guide for verl at [ICLR 2025 Expo](https://iclr.cc/virtual/2025/calendar?filter_events=Expo+Talk+Panel&filter_rooms=), [SCI-FM workshop](https://open-foundation-model.github.io/) and [LMSys afterparty](https://lu.ma/d23nyynm). Talk materials available [here](https://github.com/eric-haibin-lin/verl-community/tree/main/iclr25).
    • +
  • [2025/03] verl v0.3.0.post1 is released! See [release note](https://github.com/verl-project/verl/releases/) for details. It achieves [~1.4x speedup](https://tongyx361.github.io/blogs/posts/verl-intro/#/verl-flexible-and-efficient-rl-for-llms) compared to prev versions.
    • +
  • [2025/05] verl will be presented at [A2M Shanghai](https://a2m.msup.com.cn/home/?aid=4488&city=shanghai) on 5/16 - 5/17.
    • +
  • [2025/05] verl will be presented at [GOSIM x PyTorch Day 2025](https://paris2025.gosim.org/). See you in Paris!
    • +
  • [2025/03] We introduced the programming model of verl at the [vLLM Beijing Meetup](https://mp.weixin.qq.com/s/n77GibL2corAtQHtVEAzfg) and [verl intro and updates](https://github.com/eric-haibin-lin/verl-community/blob/main/slides/verl-lmsys-meetup.pdf) at the [SGLang-LMSYS Org Meetup](https://lu.ma/ntjrr7ig) in Sunnyvale mid-March.
    • +
  • [2025/03] We will present verl(HybridFlow) at EuroSys 2025. See you in Rotterdam!
    • +
  • [2025/02] verl v0.2.0.post2 is released!
    • +
  • [2025/02] We presented verl in the Bytedance/NVIDIA/Anyscale Ray Meetup. See you in San Jose!
    • +
  • [2025/01] [Doubao-1.5-pro](https://team.doubao.com/zh/special/doubao_1_5_pro) is released with SOTA-level performance on LLM & VLM. The RL scaling preview model is trained using verl, reaching OpenAI O1-level performance on math benchmarks (70.0 pass@1 on AIME).
    • +
  • [2024/12] verl is presented at Ray Forward 2024. Slides available here
    • +
  • [2024/12] The team presented Post-training LLMs: From Algorithms to Infrastructure at NeurIPS 2024. Slides and video available.
    • +
  • [2024/10] verl is presented at Ray Summit. Youtube video available.
    • +
  • [2024/08] HybridFlow (verl) is accepted to EuroSys 2025.
    • +
+
+ +## Key Features + +- **FSDP**, **FSDP2** and **Megatron-LM** for training. +- **vLLM**, **SGLang** and **HF Transformers** for rollout generation. +- Compatible with Hugging Face Transformers and Modelscope Hub: Qwen3.5, Qwen3, Qwen-2.5, Llama3.1, Gemma2, DeepSeek-LLM, etc +- Supervised fine-tuning. +- Reinforcement learning with [PPO](examples/ppo_trainer/), [GRPO](examples/grpo_trainer/), [GSPO](https://github.com/verl-project/verl-recipe/tree/main/gspo/), [ReMax](examples/remax_trainer/), [REINFORCE++](https://verl.readthedocs.io/en/latest/examples/config.html#algorithm), [RLOO](examples/rloo_trainer/), [PRIME](https://github.com/verl-project/verl-recipe/tree/main/prime/), [DAPO](https://github.com/verl-project/verl-recipe/tree/main/dapo/), [DrGRPO](https://github.com/verl-project/verl-recipe/tree/main/drgrpo), [KL_Cov & Clip_Cov](https://github.com/verl-project/verl-recipe/tree/main/entropy) etc. + - Support model-based reward and function-based reward (verifiable reward) for math, [coding](https://github.com/verl-project/verl-recipe/tree/main/dapo), etc + - Support vision-language models (VLMs) and [multi-modal RL](examples/grpo_trainer/run_qwen2_5_vl_7b_fsdp.sh) with Qwen2.5-vl, Kimi-VL + - [Multi-turn with tool calling](examples/tutorial/agent_loop_get_started/) +- LLM alignment recipes such as [Self-play preference optimization (SPPO)](https://github.com/verl-project/verl-recipe/tree/main/sppo) +- Flash attention 2, sequence packing, sequence parallelism via DeepSpeed Ulysses, [LoRA](examples/tuning/lora/run_qwen3_8b_fsdp.sh), [Liger-kernel](examples/sft/gsm8k/run_qwen2_5_0_5b_fsdp.sh) (`USE_LIGER=1`). +- Scales up to 671B models and hundreds of GPUs with [expert parallelism](https://github.com/verl-project/verl/pull/1467) +- Multi-gpu [LoRA RL](https://verl.readthedocs.io/en/latest/advance/ppo_lora.html) support to save memory. +- Experiment tracking with wandb, swanlab, mlflow and tensorboard. +- Hardware Support: Supports NVIDIA, AMD, [Ascend](https://github.com/verl-project/verl/blob/main/docs/ascend_tutorial/quick_start/ascend_quick_start.rst) + +## Getting Started + +Documentation + +**Quickstart:** + +- [Installation](https://verl.readthedocs.io/en/latest/start/install.html) +- [Quickstart](https://verl.readthedocs.io/en/latest/start/quickstart.html) +- [Programming Guide](https://verl.readthedocs.io/en/latest/hybrid_flow.html) & [Tech Talk](https://hcqnc.xetlk.com/sl/3vACOK) (in Chinese) +- [PPO in verl](https://verl.readthedocs.io/en/latest/algo/ppo.html) +- [GRPO in verl](https://verl.readthedocs.io/en/latest/algo/grpo.html) + +**Running a PPO example step-by-step:** + +- [Prepare Data for Post-Training](https://verl.readthedocs.io/en/latest/preparation/prepare_data.html) +- [Implement Reward Function for Dataset](https://verl.readthedocs.io/en/latest/preparation/reward_function.html) +- [PPO Example Architecture](https://verl.readthedocs.io/en/latest/examples/ppo_code_architecture.html) +- [Config Explanation](https://verl.readthedocs.io/en/latest/examples/config.html) + +**Reproducible algorithm baselines:** + +- [RL performance on coding, math](https://verl.readthedocs.io/en/latest/algo/baseline.html) + +**Algorithm recipes (`recipe/`):** + +- Optional workflows and baselines live under [`recipe/`](recipe/). Each recipe subdirectory includes a small **`REQUIRED_VERL.txt`** file describing the intended `verl` install: pinned recipes use a **tag or fixed git SHA**; rolling recipes record an explicit **`VERL_COMMIT`** (and related submodule / recipe-folder SHAs) so you can `pip install verl@git+…@` without guessing. See [`recipe/README.md`](recipe/README.md) for the full index and links. + +**For code explanation and advance usage (extension):** + +- PPO Trainer and Workers + + - [PPO Ray Trainer](https://verl.readthedocs.io/en/latest/workers/ray_trainer.html) + - [Model Engine](https://verl.readthedocs.io/en/latest/workers/model_engine.html) + - [Engine Workers (FSDP / Megatron-LM / Automodel / VeOmni / TorchTitan)](https://verl.readthedocs.io/en/latest/workers/engine_workers.html) + +- Advanced Usage and Extension + - [Add Models with the FSDP Backend](https://verl.readthedocs.io/en/latest/advance/fsdp_extension.html) + - [Add Models with the Megatron-LM Backend](https://verl.readthedocs.io/en/latest/advance/megatron_extension.html) + - [Multi-turn Rollout Support](https://verl.readthedocs.io/en/latest/sglang_multiturn/multiturn.html) + - [Search Tool Integration](https://verl.readthedocs.io/en/latest/sglang_multiturn/search_tool_example.html) + - [Sandbox Fusion Integration](https://verl.readthedocs.io/en/latest/examples/sandbox_fusion_example.html) + - [Extend to Other RL(HF) algorithms](https://verl.readthedocs.io/en/latest/advance/dpo_extension.html) + - [Ray API design tutorial](https://verl.readthedocs.io/en/latest/advance/placement.html) + +**Blogs from the community** + +- [When Reasoning Models Break Tokenization: The Hidden Complexity of Multiturn Training](https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/rlhf/verl/multi-turn/fast_tokenization/multiturn_tokenization_and_masking.md) +- [verl deployment on AWS SageMaker](https://medium.com/@kaige.yang0110/run-verl-on-sagemaker-using-4x8-l40s-gpus-8e6d5c3c61d3) +- [verl x SGLang Multi-turn Code Walkthrough](https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/rlhf/verl/multi-turn/code-walk-through/readme_EN.md) +- [Optimizing SGLang Memory Usage in verl](https://hebiao064.github.io/rl-memory-management) +- [SGLang, verl, OpenBMB and Tsinghua University: Pioneering End-to-End Multi-Turn RLHF](https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/rlhf/verl/multi-turn/verl-multiturn-rollout-Release.md) +- [Reinforcement Learning from Human Feedback on AMD GPUs with verl and ROCm Integration](https://rocm.blogs.amd.com/artificial-intelligence/verl-large-scale/README.html) +- [veMLP x verl ：玩转强化学习训练](https://mp.weixin.qq.com/s/7nbqxk4knMGd-hQE9ls2tA) +- [使用 verl 进行 GRPO 分布式强化学习训练最佳实践](https://www.volcengine.com/docs/6459/1463942) +- [HybridFlow verl 原文浅析](https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/rlhf/verl/readme.md) +- [最高提升 20 倍吞吐量！豆包大模型团队发布全新 RLHF 框架，现已开源！](https://team.doubao.com/en/blog/%E6%9C%80%E9%AB%98%E6%8F%90%E5%8D%8720%E5%80%8D%E5%90%9E%E5%90%90%E9%87%8F-%E8%B1%86%E5%8C%85%E5%A4%A7%E6%A8%A1%E5%9E%8B%E5%9B%A2%E9%98%9F%E5%8F%91%E5%B8%83%E5%85%A8%E6%96%B0-rlhf-%E6%A1%86%E6%9E%B6-%E7%8E%B0%E5%B7%B2%E5%BC%80%E6%BA%90) + +## Performance Tuning Guide + +The performance is essential for on-policy RL algorithm. We have written a detailed [performance tuning guide](https://verl.readthedocs.io/en/latest/perf/perf_tuning.html) to help you optimize performance. + +## Upgrade to vLLM >= v0.8.2 + +verl now supports vLLM>=0.8.2 when using FSDP as the training backend. Please refer to [this document](https://github.com/verl-project/verl/blob/main/docs/README_vllm0.8.md) for the installation guide and more information. Please avoid vllm 0.7.x, which contains bugs that may lead to OOMs and unexpected errors. + +## Use Latest SGLang + +SGLang is fully supported with verl, and SGLang RL Group is working extensively on building unique features, including multi-turn agentic RL, VLM RLHF, server-based RL, and partial rollout. Please refer to [this document](https://verl.readthedocs.io/en/latest/workers/sglang_worker.html) for the installation guide and more information. + +## Upgrade to FSDP2 + +verl is fully embracing FSDP2! FSDP2 is recommended by torch distributed team, providing better throughput and memory usage, and is composible with other features (e.g. torch.compile). To enable FSDP2, simply use verl main and set the following options: + +``` +actor_rollout_ref.ref.strategy=fsdp2 +actor_rollout_ref.actor.strategy=fsdp2 +critic.strategy=fsdp2 +``` + +Furthermore, FSDP2 cpu offloading is compatible with gradient accumulation. You can turn it on to save memory with `actor_rollout_ref.actor.fsdp_config.offload_policy=True`. For more details, see https://github.com/verl-project/verl/pull/1026 + +## AMD Support (ROCm Kernel) + +verl now supports FSDP as the training engine (Megatron support coming soon) and both integrates with vLLM and SGLang as inference engines. Please refer to [this document](https://github.com/verl-project/verl/blob/main/docs/amd_tutorial/amd_build_dockerfile_page.rst) for the installation guide and more information, and [this document](https://github.com/verl-project/verl/blob/main/docs/amd_tutorial/amd_vllm_page.rst) for the vLLM performance tuning for ROCm. + +## Citation and acknowledgement + +If you find the project helpful, please cite: + +- [HybridFlow: A Flexible and Efficient RLHF Framework](https://arxiv.org/abs/2409.19256v2) +- [A Framework for Training Large Language Models for Code Generation via Proximal Policy Optimization](https://i.cs.hku.hk/~cwu/papers/gmsheng-NL2Code24.pdf) + +```bibtex +@article{sheng2024hybridflow, + title = {HybridFlow: A Flexible and Efficient RLHF Framework}, + author = {Guangming Sheng and Chi Zhang and Zilingfeng Ye and Xibin Wu and Wang Zhang and Ru Zhang and Yanghua Peng and Haibin Lin and Chuan Wu}, + year = {2024}, + journal = {arXiv preprint arXiv: 2409.19256} +} +``` + +verl is inspired by the design of Nemo-Aligner, Deepspeed-chat and OpenRLHF. The project is adopted and contributed by Bytedance, Anyscale, LMSys.org, [Alibaba Qwen team](https://github.com/QwenLM/), Shanghai AI Lab, Tsinghua University, UC Berkeley, UCLA, UIUC, University of Hong Kong, ke.com, [All Hands AI](https://www.all-hands.dev/), [ModelBest](http://modelbest.cn/), JD AI Lab, Microsoft Research, [StepFun](https://www.stepfun.com/), Amazon, LinkedIn, Meituan, [Camel-AI](https://www.camel-ai.org/), [OpenManus](https://github.com/OpenManus), Xiaomi, NVIDIA research, [Baichuan](https://www.baichuan-ai.com/home), [RedNote](https://www.xiaohongshu.com/), [SwissAI](https://www.swiss-ai.org/), [Moonshot AI (Kimi)](https://www.moonshot-ai.com/), Baidu, Snowflake, Skywork.ai, JetBrains, [IceSword Lab](https://www.iceswordlab.com), and many more. + +## Awesome Projects Built with `verl` + +Welcome to register your awesome project build with `verl` for other developers' reference! + +- [TinyZero](https://github.com/Jiayi-Pan/TinyZero): a reproduction of **DeepSeek R1 Zero** recipe for reasoning tasks ![GitHub Repo stars](https://img.shields.io/github/stars/Jiayi-Pan/TinyZero) +- [SkyThought](https://github.com/NovaSky-AI/SkyThought): RL training for Sky-T1-7B by NovaSky AI team. ![GitHub Repo stars](https://img.shields.io/github/stars/NovaSky-AI/SkyThought) +- [simpleRL-reason](https://github.com/hkust-nlp/simpleRL-reason): SimpleRL-Zoo: Investigating and Taming Zero Reinforcement Learning for Open Base Models in the Wild ![GitHub Repo stars](https://img.shields.io/github/stars/hkust-nlp/simpleRL-reason) +- [Easy-R1](https://github.com/hiyouga/EasyR1): **Multi-modal** RL training framework ![GitHub Repo stars](https://img.shields.io/github/stars/hiyouga/EasyR1) +- [RandOpt](https://github.com/sunrainyg/RandOpt): Neural Thickets: Diverse Task Experts Are Dense Around Pretrained Weights ![GitHub Repo stars](https://img.shields.io/github/stars/sunrainyg/RandOpt) +- [OpenManus-RL](https://github.com/OpenManus/OpenManus-RL): LLM Agents RL tuning framework for multiple agent environments. ![GitHub Repo stars](https://img.shields.io/github/stars/OpenManus/OpenManus-RL) +- [rllm](https://github.com/agentica-project/rllm): async RL training with [verl-pipeline](https://github.com/agentica-project/verl-pipeline) ![GitHub Repo stars](https://img.shields.io/github/stars/agentica-project/rllm) +- [RAGEN](https://github.com/ZihanWang314/ragen): a general-purpose reasoning **agent** training framework ![GitHub Repo stars](https://img.shields.io/github/stars/ZihanWang314/ragen) +- [Search-R1](https://github.com/PeterGriffinJin/Search-R1): RL with reasoning and **searching (tool-call)** interleaved LLMs ![GitHub Repo stars](https://img.shields.io/github/stars/PeterGriffinJin/Search-R1) +- [ReSearch](https://github.com/Agent-RL/ReSearch): Learning to **Re**ason with **Search** for LLMs via Reinforcement Learning ![GitHub Repo stars](https://img.shields.io/github/stars/Agent-RL/ReSearch) +- [Skywork-OR1](https://github.com/SkyworkAI/Skywork-OR1): Skywork open reaonser series ![GitHub Repo stars](https://img.shields.io/github/stars/SkyworkAI/Skywork-OR1) +- [ToRL](https://github.com/GAIR-NLP/ToRL): Scaling tool-integrated RL ![GitHub Repo stars](https://img.shields.io/github/stars/GAIR-NLP/ToRL) +- [Absolute Zero Reasoner](https://github.com/LeapLabTHU/Absolute-Zero-Reasoner): [A no human curated data self-play framework for reasoning](https://arxiv.org/abs/2505.03335) ![GitHub Repo stars](https://img.shields.io/github/stars/LeapLabTHU/Absolute-Zero-Reasoner) +- [verl-agent](https://github.com/langfengQ/verl-agent): A scalable training framework for **long-horizon LLM/VLM agents**, along with a new algorithm **GiGPO** ![GitHub Repo stars](https://img.shields.io/github/stars/langfengQ/verl-agent) +- [RL-Factory](https://github.com/Simple-Efficient/RL-Factory): An easy and efficient RL post-training framework for Agentic Learning ![GitHub Repo stars](https://img.shields.io/github/stars/Simple-Efficient/RL-Factory) +- [ReTool](https://retool-rl.github.io/): ReTool: reinforcement learning for strategic tool use in LLMs. Code release is in progress... +- [verl-tool](https://github.com/TIGER-AI-Lab/verl-tool): An unified and easy-to-extend tool-agent training framework based on verl![GitHub Repo stars](https://img.shields.io/github/stars/TIGER-AI-Lab/verl-tool) +- [PRIME](https://github.com/PRIME-RL/PRIME): Process reinforcement through implicit rewards ![GitHub Repo stars](https://img.shields.io/github/stars/PRIME-RL/PRIME) +- [MemAgent](https://github.com/BytedTsinghua-SIA/MemAgent): MemAgent: Reshaping Long-Context LLM with Multi-Conv RL based Memory Agent ![GitHub Repo stars](https://img.shields.io/github/stars/BytedTsinghua-SIA/MemAgent) +- [POLARIS](https://github.com/ChenxinAn-fdu/POLARIS): A Post-training recipe for scaling RL on Advanced Reasoning models ![GitHub Repo stars](https://img.shields.io/github/stars/ChenxinAn-fdu/POLARIS) +- [GUI-R1](https://github.com/ritzz-ai/GUI-R1): **GUI-R1**: A Generalist R1-style Vision-Language Action Model For **GUI Agents** ![GitHub Repo stars](https://img.shields.io/github/stars/ritzz-ai/GUI-R1) +- [DeepRetrieval](https://github.com/pat-jj/DeepRetrieval): RL Training of **Search Agent** with **Search/Retrieval Outcome** ![GitHub Repo stars](https://img.shields.io/github/stars/pat-jj/DeepRetrieval) +- [Code-R1](https://github.com/ganler/code-r1): Reproducing R1 for **Code** with Reliable Rewards ![GitHub Repo stars](https://img.shields.io/github/stars/ganler/code-r1) +- [DeepResearcher](https://github.com/GAIR-NLP/DeepResearcher): Scaling deep research via reinforcement learning in real-world environments ![GitHub Repo stars](https://img.shields.io/github/stars/GAIR-NLP/DeepResearcher) +- [VAGEN](https://github.com/RAGEN-AI/VAGEN): Training VLM agents with multi-turn reinforcement learning ![GitHub Repo stars](https://img.shields.io/github/stars/RAGEN-AI/VAGEN) +- [RM-R1](https://arxiv.org/abs/2505.02387): RL training of reasoning reward models ![GitHub Repo stars](https://img.shields.io/github/stars/RM-R1-UIUC/RM-R1) +- [Dr. MAS](https://arxiv.org/pdf/2602.08847): Stable **end-to-end RL** post-training for **multi-agent LLM systems** ![GitHub Repo stars](https://img.shields.io/github/stars/langfengQ/DrMAS) +- [LUFFY](https://arxiv.org/pdf/2504.14945): Learning to Reason under Off-Policy Guidance![GitHub Repo stars](https://img.shields.io/github/stars/ElliottYan/LUFFY) +- [DeepMath](https://github.com/zwhe99/DeepMath): DeepMath-103K data and series models for math reasoning![GitHub Repo stars](https://img.shields.io/github/stars/zwhe99/DeepMath) +- [PACS](https://github.com/ritzz-ai/PACS): Implicit Actor Critic Coupling via a Supervised Learning Framework for RLVR ![GitHub Repo stars](https://img.shields.io/github/stars/ritzz-ai/PACS) +- [Entropy Mechanism of RL](https://github.com/PRIME-RL/Entropy-Mechanism-of-RL): The Entropy Mechanism of Reinforcement Learning for Large Language Model Reasoning![GitHub Repo stars](https://img.shields.io/github/stars/PRIME-RL/Entropy-Mechanism-of-RL) +- [LLaSA-TTS-GRPO](https://github.com/channel-io/ch-tts-llasa-rl-grpo): TTS fine-tuning with GRPO optimization based on LLASA models ![GitHub Repo stars](https://img.shields.io/github/stars/channel-io/ch-tts-llasa-rl-grpo) +- [PF-PPO](https://arxiv.org/abs/2409.06957): Policy Filtration for PPO based on the reliability of reward signals for more efficient and robust RLHF. +- [RACRO](https://github.com/gyhdog99/RACRO2): Build multi-modal reasoning models via decoupling it into query-conditioned captioning and text-only reasoning ![GitHub Repo stars](https://img.shields.io/github/stars/gyhdog99/RACRO2) +- [Agent Lightning](https://github.com/microsoft/agent-lightning): A flexible and extensible framework that enables seamless agent optimization for any existing agent framework. ![GitHub Repo stars](https://img.shields.io/github/stars/microsoft/agent-lightning) +- [VTool-R1](https://github.com/VTOOL-R1/vtool-r1): VLMs Learn to Think with Images via Reinforcement Learning on Multimodal Tool Use. ![GitHub Repo stars](https://img.shields.io/github/stars/VTOOL-R1/vtool-r1) +- [Kimina-Prover-RL](https://github.com/project-numina/kimina-prover-rl/tree/main/recipe/kimina_prover_rl): Training pipeline for formal theorem proving, based on a paradigm inspired by DeepSeek-R1. +- [RL-PLUS](https://github.com/YihongDong/RL-PLUS): Countering Capability Boundary Collapse of LLMs in Reinforcement Learning with Hybrid-policy Optimization. +- [rStar2-Agent](https://github.com/microsoft/rStar): Using reinforcement learning with multi-step tool-calling for math tasks, rStar2-Agent-14B reaches frontier-level math reasoning in just 510 RL training steps ![GitHub Repo stars](https://img.shields.io/github/stars/microsoft/rStar) +- [Vision-SR1](https://github.com/zli12321/Vision-SR1): Self-Rewarding Vision-Language Model via Reasoning Decomposition ![GitHub Repo stars](https://img.shields.io/github/stars/zli12321/Vision-SR1) +- [SimpleVLA-RL](https://github.com/PRIME-RL/SimpleVLA-RL): SimpleVLA-RL: A Simple yet Effective Vision-Language Action Model for Reinforcement Learning ![GitHub Repo stars](https://img.shields.io/github/stars/PRIME-RL/SimpleVLA-RL) +- [Table-R1](https://github.com/Table-R1/Table-R1): Table-R1: Inference-Time Scaling for Table Reasoning ![GitHub Repo stars](https://img.shields.io/github/stars/Table-R1/Table-R1) +- [Revisual-R1](https://github.com/CSfufu/Revisual-R1): Revisual-R1: Advancing Multimodal Reasoning From Optimized Cold Start to Staged Reinforcement Learning ![GitHub Repo stars](https://img.shields.io/github/stars/CSfufu/Revisual-R1) +- [ARES](https://github.com/shawn0728/ARES): ARES: Multimodal Adaptive Reasoning via Difficulty-Aware Token-Level Entropy Shaping ![GitHub Repo stars](https://img.shields.io/github/stars/shawn0728/ARES) +- [Meta-Bandit-LLM](https://github.com/sanxing-chen/meta-bandit-llm): Meta-Bandit-LLM: Long-horizon multiturn interactive training for meta-bandit agents ![GitHub Repo stars](https://img.shields.io/github/stars/sanxing-chen/meta-bandit-llm) +- [PokeeResearch](https://github.com/Pokee-AI/PokeeResearchOSS): PokeeResearch: State-of-the-art 7B DeepResearch Agent that leverages web search and content reading capabilities to answer complex questions using the most up-to-date information available online. ![Github Repo Stars](https://img.shields.io/github/stars/Pokee-AI/PokeeResearchOSS) +- [Search Self-play](https://github.com/Alibaba-Quark/SSP): Pushing the Frontier of Agent Capability without Supervision ![GitHub Repo stars](https://img.shields.io/github/stars/Alibaba-Quark/SSP) +- [OneThinker](https://github.com/tulerfeng/OneThinker): All-in-one Reasoning Model for Image and Video ![GitHub Repo stars](https://img.shields.io/github/stars/tulerfeng/OneThinker) +- [OpenTinker](https://github.com/open-tinker/OpenTinker): Democratizing Agentic Reinforcement Learning as a Service ![GitHub Repo stars](https://img.shields.io/github/stars/open-tinker/OpenTinker) +- [FlowRL](https://github.com/Xuekai-Zhu/FlowRL): Matching reward distributions via **flow balance** for diverse exploration and generalizable reasoning ![GitHub Repo stars](https://img.shields.io/github/stars/Xuekai-Zhu/FlowRL) +- [Logic-RL](https://github.com/Unakar/Logic-RL): a reproduction of DeepSeek R1 Zero on 2K Tiny Logic Puzzle Dataset. ![GitHub Repo stars](https://img.shields.io/github/stars/Unakar/Logic-RL) +- [Seed-Coder](https://github.com/ByteDance-Seed/Seed-Coder): RL training of Seed-Coder boosts performance on competitive programming ![GitHub Repo stars](https://img.shields.io/github/stars/ByteDance-Seed/Seed-Coder) +- [all-hands/openhands-lm-32b-v0.1](https://www.all-hands.dev/blog/introducing-openhands-lm-32b----a-strong-open-coding-agent-model): A strong, open coding agent model, trained with [multi-turn fine-tuning](https://github.com/verl-project/verl/pull/195) +- [s3](https://github.com/pat-jj/s3) **Efficient Yet Effective** Search Agent Training via RL ![GitHub Repo stars](https://img.shields.io/github/stars/pat-jj/s3) +- [Rec-R1](https://arxiv.org/pdf/2503.24289): Bridging Generative Large Language Models and Recommendation Systems via Reinforcement Learning +- [Explore RL Data Scaling](https://arxiv.org/abs/2503.22230): Exploring Data Scaling Trends and Effects in Reinforcement Learning from Human Feedback +- [FIRE](https://arxiv.org/abs/2410.21236): Flaming-hot initiation with regular execution sampling for large language models +- [DQO](https://arxiv.org/abs/2410.09302): Enhancing multi-Step reasoning abilities of language models through direct Q-function optimization +- [ProRL](https://arxiv.org/abs/2505.24864): Prolonged Reinforcement Learning Expands Reasoning Boundaries in Large Language Models +- [cognition-engineering](https://github.com/gair-nlp/cognition-engineering): Test time scaling drives cognition engineering. ![GitHub Repo stars](https://img.shields.io/github/stars/gair-nlp/cognition-engineering) +- [Trust Region Preference Approximation](https://github.com/XueruiSu/Trust-Region-Preference-Approximation): A simple and stable **reinforcement learning algorithm** for LLM reasoning. ![GitHub Repo stars](https://img.shields.io/github/stars/XueruiSu/Trust-Region-Preference-Approximation) +- [AdaRFT](https://github.com/uscnlp-lime/verl): Efficient Reinforcement Finetuning via **Adaptive Curriculum Learning** ![GitHub Repo stars](https://img.shields.io/github/stars/uscnlp-lime/verl) +- [critic-rl](https://github.com/HKUNLP/critic-rl): LLM critics for code generation ![GitHub Repo stars](https://img.shields.io/github/stars/HKUNLP/critic-rl) +- [self-rewarding-reasoning-LLM](https://arxiv.org/pdf/2502.19613): self-rewarding and correction with **generative reward models** ![GitHub Repo stars](https://img.shields.io/github/stars/RLHFlow/Self-rewarding-reasoning-LLM) +- [DeepEnlighten](https://github.com/DolbyUUU/DeepEnlighten): Reproduce R1 with **social reasoning** tasks and analyze key findings ![GitHub Repo stars](https://img.shields.io/github/stars/DolbyUUU/DeepEnlighten) +- [MetaSpatial](https://github.com/PzySeere/MetaSpatial): Reinforcing **3D Spatial Reasoning** in **VLMs** for the **Metaverse** ![GitHub Repo stars](https://img.shields.io/github/stars/PzySeere/MetaSpatial) +- [PURE](https://github.com/CJReinforce/PURE): **Credit assignment** is the key to successful reinforcement fine-tuning using **process reward model** ![GitHub Repo stars](https://img.shields.io/github/stars/CJReinforce/PURE) +- [cognitive-behaviors](https://github.com/kanishkg/cognitive-behaviors): Cognitive Behaviors that Enable Self-Improving Reasoners, or, Four Habits of Highly Effective STaRs ![GitHub Repo stars](https://img.shields.io/github/stars/kanishkg/cognitive-behaviors) +- [deepscaler](https://github.com/agentica-project/rllm/tree/deepscaler): iterative context scaling with GRPO ![GitHub Repo stars](https://img.shields.io/github/stars/agentica-project/deepscaler) +- [DAPO](https://dapo-sia.github.io/): the fully open source SOTA RL algorithm that beats DeepSeek-R1-zero-32B ![GitHub Repo stars](https://img.shields.io/github/stars/verl-project/verl) +- [NoisyRollout](https://github.com/NUS-TRAIL/NoisyRollout): Reinforcing Visual Reasoning with Data Augmentation ![GitHub Repo stars](https://img.shields.io/github/stars/NUS-TRAIL/NoisyRollout) +- [SPEAR](https://github.com/TencentYoutuResearch/SPEAR): **Self-imitation** with **Progressive Exploration** for Agentic Reinforcement Learning (ICLR 2026) ![GitHub Repo stars](https://img.shields.io/github/stars/TencentYoutuResearch/SPEAR) +- [RuleReasoner](https://github.com/bigai-nlco/RuleReasoner): **RuleReasoner:** Reinforced Rule-based Reasoning via **Domain-aware Dynamic Sampling** (ICLR 2026) ![GitHub Repo stars](https://img.shields.io/github/stars/bigai-nlco/RuleReasoner) +- [MetaphorStar](https://metaphorstar.github.io/): **Image Metaphor** Understanding and Reasoning with End-to-End **Visual Reinforcement Learning** ![GitHub Repo stars](https://img.shields.io/github/stars/MING-ZCH/MetaphorStar) +- [DART-GUI](https://github.com/Computer-use-agents/dart-gui): a decoupled agentic RL framework for Computer Use Agents, achieving ~2× training speedup and ~5× environment utilization! ![GitHub Repo stars](https://img.shields.io/github/stars/Computer-use-agents/dart-gui) + +## Contribution Guide + +See [contributions guide](CONTRIBUTING.md) + +## About [ByteDance Seed Team](https://team.doubao.com/) + +Founded in 2023, ByteDance Seed Team is dedicated to crafting the industry's most advanced AI foundation models. The team aspires to become a world-class research team and make significant contributions to the advancement of science and society. You can get to know Bytedance Seed better through the following channels👇 + + + +We are HIRING! Send us an [email](mailto:the.verl.project@gmail.com) if you are interested in internship/FTE opportunities in RL for agents. diff --git a/verl/core.778099 b/verl/core.778099 new file mode 100644 index 0000000000000000000000000000000000000000..d588d25c3a9f965e2b18326862ae1179cfcc5ff4 --- /dev/null +++ b/verl/core.778099 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:281fa540822bd0f17a2b720cec4ff6013b97898f7654cf0d7b1a96a64570307f +size 150466560 diff --git a/verl/docker/Dockerfile.isaaclab230 b/verl/docker/Dockerfile.isaaclab230 new file mode 100644 index 0000000000000000000000000000000000000000..a1d0f6a54f68aa3d76a68c7b3d2c801ee9eafb68 --- /dev/null +++ b/verl/docker/Dockerfile.isaaclab230 @@ -0,0 +1,150 @@ + +#FROM nvcr.nju.edu.cn/nvidia/isaac-lab:2.3.0 +FROM isaac-lab-base:latest + +ENV ACCEPT_EULA=Y +ENTRYPOINT [] + +# desktop +RUN --mount=type=cache,target=/var/cache/apt \ + sed -i 's/archive.ubuntu.com/mirrors.ivolces.com/g' /etc/apt/sources.list && \ + sed -i 's/security.ubuntu.com/mirrors.ivolces.com/g' /etc/apt/sources.list && \ + apt-get update && \ + DEBIAN_FRONTEND=noninteractive apt-get install -y locales && \ + locale-gen en_US.UTF-8 && \ + update-locale LANG=en_US.UTF-8 LC_CTYPE=en_US.UTF-8 && \ + apt-get install -y wget curl \ + xfce4 \ + xfce4-goodies \ + xorg \ + dbus-x11 \ + x11-xserver-utils \ + tigervnc-standalone-server \ + tigervnc-common \ + tigervnc-tools \ + fonts-dejavu \ + fonts-liberation +# cuda 12.2 +RUN --mount=type=cache,target=/var/cache/apt \ + cd /tmp && \ + wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/3bf863cc.pub && \ + apt-key add 3bf863cc.pub && \ + echo "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64 /" > /etc/apt/sources.list.d/cuda.list && \ + apt-get update && \ + apt-get install -y libcusparselt0 libnccl2=2.27.3-1+cuda12.2 libglfw3 libgl1-mesa-glx libosmesa6 && \ + rm -f 3bf863cc.pub + +# libero +RUN --mount=type=cache,target=/root/.cache/pip \ + /workspace/isaaclab/_isaac_sim/python.sh -m pip install easydict==1.9 robosuite==1.4.0 bddl==1.0.1 future==0.18.2 cloudpickle==2.1.0 + +RUN --mount=type=cache,target=/root/.cache/pip \ + /workspace/isaaclab/_isaac_sim/python.sh -m pip install transformers[hf_xet] + +RUN --mount=type=cache,target=/root/.cache/pip \ + /workspace/isaaclab/_isaac_sim/python.sh -m pip install --upgrade numpy==1.26.4 ray[default] \ + accelerate codetiming datasets dill hydra-core pandas peft pyarrow>=19.0.0 pybind11 pylatexenc + +# openvla-oft +RUN --mount=type=cache,target=/root/.cache/pip \ + /workspace/isaaclab/_isaac_sim/python.sh -m pip install pre-commit torchdata packaging>=20.0 uvicorn fastapi latex2sympy2_extended math_verify tensorboard + + +# flash_attn +RUN cd /tmp && \ + wget -nv https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.0.post2/flash_attn-2.8.0.post2+cu12torch2.7cxx11abiFALSE-cp311-cp311-linux_x86_64.whl && \ + /workspace/isaaclab/_isaac_sim/python.sh -m pip install /tmp/flash_attn-2.8.0.post2+cu12torch2.7cxx11abiFALSE-cp311-cp311-linux_x86_64.whl && \ + rm -f /tmp/flash_attn-2.8.0.post2+cu12torch2.7cxx11abiFALSE-cp311-cp311-linux_x86_64.whl + +RUN --mount=type=cache,target=/root/.cache/pip \ + /workspace/isaaclab/_isaac_sim/python.sh -m pip install --upgrade protobuf==3.20.3 timm==0.9.16 + +RUN --mount=type=cache,target=/root/.cache/pip \ + /workspace/isaaclab/_isaac_sim/python.sh -m pip install orjson==3.11.3 pyvers==0.1.0 tensordict==0.10.0 --force --no-deps + + +RUN mkdir -p /root/.vnc && \ + cat <<'EOP' > /root/.vnc/xstartup +#!/bin/sh +unset SESSION_MANAGER +unset DBUS_SESSION_BUS_ADDRESS +[ -r \$HOME/.Xresources ] && xrdb \$HOME/.Xresources +xsetroot -solid grey +exec startxfce4 +EOP + +RUN cat <<'EOP' > /root/.vnc/config +geometry=1920x1080 +depth=24 +desktop=Isaac-Sim-Desktop +dpi=96 +localhost=no +EOP + +RUN cat <<'EOP' > /root/start_isaac_vnc.sh +#!/bin/bash +# 设置显示变量 +export DISPLAY=:1 + +# 检查VNC是否运行 +if ! pgrep -f "Xvnc.*:1" > /dev/null; then + echo "Starting VNC server..." + vncserver :1 -localhost no -geometry 1920x1080 -depth 24 -desktop "Isaac-Sim-Desktop" + sleep 3 +fi + +# 启动Isaac Sim +echo "Starting Isaac Sim..." +/workspace/isaaclab/_isaac_sim/isaac-sim.sh --allow-root +EOP + +RUN chmod +x /root/.vnc/xstartup && \ + chmod +x /root/start_isaac_vnc.sh + +RUN /workspace/isaaclab/_isaac_sim/isaac-sim.sh --allow-root --ext-precache-mode + +RUN cd /root && \ + git clone https://github.com/Lifelong-Robot-Learning/LIBERO.git && \ + cd LIBERO && \ + git apply <<'EOP' +diff --git a/setup.py b/setup.py +index 59d4900..dbe9811 100644 +--- a/setup.py ++++ b/setup.py +@@ -13,7 +13,8 @@ long_description = "".join(lines) + + setup( + name="libero", +- packages=[package for package in find_packages() if package.startswith("libero")], ++ #packages=[package for package in find_packages() if package.startswith("libero")], ++ packages=["libero"], + install_requires=[], + eager_resources=["*"], + include_package_data=True, +EOP + +RUN cd /root/LIBERO && \ + /workspace/isaaclab/_isaac_sim/python.sh -m pip install -e . + +# libero config +RUN mkdir -p /root/.libero && \ +cat <<'EOP' > /root/.libero/config.yaml +assets: /root/LIBERO/libero/libero/./assets +bddl_files: /root/LIBERO/libero/libero/./bddl_files +benchmark_root: /root/LIBERO/libero/libero +datasets: /root/LIBERO/libero/libero/../datasets +init_states: /root/LIBERO/libero/libero/./init_files +EOP + +# from https://github.com/nvidia-china-sae/RobotLearningLab +COPY RobotLearningLab/ /root/RobotLearningLab/ + +RUN cd /workspace/isaaclab/ && \ + rm -rf source && \ + ln -s /root/RobotLearningLab/source source && \ + /workspace/isaaclab/_isaac_sim/python.sh -m pip install -e ./source/isaaclab +# Ray cmd +RUN /workspace/isaaclab/_isaac_sim/python.sh -m pip install colorama && \ +cat <<'EOP' >> /root/.bashrc +alias ray='/workspace/isaaclab/_isaac_sim/python.sh /workspace/isaaclab/_isaac_sim/kit/python/lib/python3.11/site-packages/ray/scripts/scripts.py' +EOP \ No newline at end of file diff --git a/verl/docker/Dockerfile.stable.sglang b/verl/docker/Dockerfile.stable.sglang new file mode 100644 index 0000000000000000000000000000000000000000..3427b12503039a9265b247273d80432642276610 --- /dev/null +++ b/verl/docker/Dockerfile.stable.sglang @@ -0,0 +1,53 @@ +# sgl059 + +FROM lmsysorg/sglang:v0.5.9 + +ARG PIP_NO_CACHE_DIR=1 + +RUN pip install pybind11 + +RUN pip install nvidia-mathdx + +RUN MAX_JOBS=128 pip install -v --disable-pip-version-check --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" git+https://github.com/NVIDIA/apex.git + +RUN export NVTE_FRAMEWORK=pytorch && MAX_JOBS=128 NVTE_BUILD_THREADS_PER_JOB=4 pip3 install --resume-retries 999 --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.12 + +# RUN pip install --upgrade transformers tokenizers + +RUN pip install codetiming mathruler pylatexenc qwen_vl_utils cachetools pytest-asyncio + +RUN pip install --no-build-isolation flash_attn==2.8.3 + +RUN NSIGHT_VERSION=2025.6.1_2025.6.1.190-1_$(if [ "$(uname -m)" = "aarch64" ]; then echo "arm64"; else echo "amd64"; fi) && \ + wget https://developer.nvidia.com/downloads/assets/tools/secure/nsight-systems/2025_6/nsight-systems-${NSIGHT_VERSION}.deb && \ + apt-get update && apt-get install -y libxcb-cursor0 && \ + apt-get install -y ./nsight-systems-${NSIGHT_VERSION}.deb && \ + rm -rf /usr/local/cuda/bin/nsys && \ + ln -s /opt/nvidia/nsight-systems/2025.6.1/nsys /usr/local/cuda/bin/nsys && \ + rm -rf /usr/local/cuda/bin/nsys-ui && \ + ln -s /opt/nvidia/nsight-systems/2025.6.1/nsys-ui /usr/local/cuda/bin/nsys-ui && \ + rm nsight-systems-${NSIGHT_VERSION}.deb + +# sglang image has already installed DeepEP + +RUN pip3 install --no-deps trl==0.27.0 + +RUN pip3 install nvtx matplotlib liger_kernel + +RUN pip install -U git+https://github.com/ISEEKYAN/mbridge.git + +RUN pip install --no-deps git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.16.0 + +RUN pip install git+https://github.com/verl-project/verl.git@v0.7.0 && \ + pip uninstall -y verl + +RUN sed -i '/nvidia-cudnn-cu12/d' /usr/local/lib/python3.12/dist-packages/torch-2.9.1+cu129.dist-info/METADATA && \ + pip install --no-deps --force-reinstall nvidia-cudnn-cu12==9.16.0.29 + +# for packages compiled from source code +RUN apt-get update && \ + apt-get install -y --allow-downgrades --allow-change-held-packages \ + libcudnn9-cuda-12=9.16.0.29-1 \ + libcudnn9-dev-cuda-12=9.16.0.29-1 \ + libcudnn9-headers-cuda-12=9.16.0.29-1 && \ + rm -rf /var/lib/apt/lists/* \ No newline at end of file diff --git a/verl/docker/Dockerfile.stable.trtllm b/verl/docker/Dockerfile.stable.trtllm new file mode 100644 index 0000000000000000000000000000000000000000..40d57a9afefe6843ecc513a443f5f3681396c076 --- /dev/null +++ b/verl/docker/Dockerfile.stable.trtllm @@ -0,0 +1,59 @@ +# Base image from NGC TensorRT-LLM, which includes a pre-installed TensorRT-LLM. +# For available images, visit: https://nvidia.github.io/TensorRT-LLM/installation/containers.html +# Use TRTLLM_BASE_IMAGE to specify the base image (default: release:1.2.0rc6) +ARG TRTLLM_BASE_IMAGE=nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc14 +FROM ${TRTLLM_BASE_IMAGE} + +# Clear TORCH_CUDA_ARCH_LIST inherited from the base image so that +# FlashInfer's check_cuda_arch() queries the actual GPU at runtime +# instead of rejecting GPUs not in the build-time arch list. +ENV TORCH_CUDA_ARCH_LIST="" + +# ============================================================================== +# Install Megatron dependencies +# ============================================================================== +# DeepEP is required for IBGDA support. +# Clone and build gdrcopy and deepep-nvshmem dependencies. +WORKDIR /home/dpsk_a2a +RUN git clone -b v2.5.1 https://github.com/NVIDIA/gdrcopy.git && \ + pushd gdrcopy && \ + make prefix=/usr/local lib_install && \ + popd && rm -rf gdrcopy && \ + pip install nvidia-nvshmem-cu13==3.3.20 && \ + export NVSHMEM_DIR=/usr/local/lib/python3.12/dist-packages/nvidia/nvshmem && \ + export LD_LIBRARY_PATH="${NVSHMEM_DIR}/lib:$LD_LIBRARY_PATH" && \ + export PATH="${NVSHMEM_DIR}/bin:$PATH" && \ + pushd ${NVSHMEM_DIR}/lib && \ + ln -s libnvshmem_host.so.3 libnvshmem_host.so && \ + popd && \ + git clone -b hybrid-ep https://github.com/deepseek-ai/DeepEP.git && \ + pushd DeepEP && \ + export CPATH=/usr/local/cuda/targets/$(uname -m | sed 's/aarch64/sbsa-linux/;s/x86_64/x86_64-linux/')/include/cccl:$CPATH && \ + TORCH_CUDA_ARCH_LIST="9.0 10.0 12.0" python setup.py install && \ + popd && rm -rf deepep + +# Install Python dependencies +RUN pip3 install --no-cache-dir --no-deps trl==0.27.0 && \ + pip3 install --no-cache-dir nvtx matplotlib liger_kernel cachetools && \ + pip3 install --no-cache-dir cupy-cuda12x==14.0.1 && \ + pip install --no-cache-dir -U git+https://github.com/ISEEKYAN/mbridge.git@641a5a0 && \ + pip install --no-deps --no-cache-dir git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.16.0 + + +# ============================================================================== +# Install verl dependencies +# ============================================================================== +RUN pip install git+https://github.com/verl-project/verl.git@v0.7.1 +RUN pip uninstall -y verl +RUN pip install "verl[mcore] @ git+https://github.com/verl-project/verl.git@v0.7.1" +RUN pip uninstall -y verl + + +# Pin Ray to a version compatible with TRT-LLM 1.3.0rc14 +RUN pip install --no-cache-dir "ray[default]==2.54.1" + +# ============================================================================== +# Install a specific TensorRT-LLM on demand +# ============================================================================== +# Note: The NGC image already includes a pre-installed TensorRT-LLM, but you can install a specific version if needed. +# Refer to https://nvidia.github.io/TensorRT-LLM/installation/index.html for more details. diff --git a/verl/docker/Dockerfile.stable.vllm b/verl/docker/Dockerfile.stable.vllm new file mode 100644 index 0000000000000000000000000000000000000000..efa82b63039a12fe213b582cda97fcf3da5b12c3 --- /dev/null +++ b/verl/docker/Dockerfile.stable.vllm @@ -0,0 +1,121 @@ +# vllm: x86_64=0.18.0, aarch64=0.18.0 + +FROM nvidia/cuda:12.9.1-devel-ubuntu24.04 + +ARG DEBIAN_FRONTEND=noninteractive +ARG PIP_NO_CACHE_DIR=1 +ARG APT_MIRROR="" +# PEP 668: Ubuntu 24.04 blocks system-wide pip installs; override for Docker +ENV PIP_BREAK_SYSTEM_PACKAGES=1 + +RUN if [ -n "${APT_MIRROR}" ]; then \ + sed -i "s@http://.*archive.ubuntu.com@${APT_MIRROR}@g" /etc/apt/sources.list.d/ubuntu.sources; \ + fi + +RUN apt-get update && apt-get install -y \ + git \ + wget \ + cmake \ + build-essential \ + libibverbs-dev \ + libnuma-dev \ + librdmacm-dev \ + numactl \ + software-properties-common \ + vim \ + python3.12 \ + python3.12-dev \ + && rm -rf /var/lib/apt/lists/* + +RUN wget https://bootstrap.pypa.io/get-pip.py && \ + python3.12 get-pip.py && \ + rm get-pip.py + +RUN ln -sf /usr/bin/python3.12 /usr/bin/python3 && \ + ln -sf /usr/bin/python3.12 /usr/bin/python + +RUN pip install torch==2.10.0 torchvision==0.25.0 torchaudio==2.10.0 --index-url https://download.pytorch.org/whl/cu129 + +RUN pip install pybind11 wheel + +# ========================= +# Install cuDNN +# ========================= +RUN ARCH=$(if [ "$(uname -m)" = "aarch64" ]; then echo "sbsa"; else echo "x86_64"; fi) && \ + wget https://developer.download.nvidia.com/compute/cuda/repos/debian12/${ARCH}/cuda-keyring_1.1-1_all.deb && \ + dpkg -i cuda-keyring_1.1-1_all.deb && \ + apt-get update && \ + apt-get -y install cudnn && \ + rm cuda-keyring_1.1-1_all.deb && \ + rm -rf /var/lib/apt/lists/* + +RUN pip install nvidia-mathdx ninja + +RUN MAX_JOBS=256 pip install -v --disable-pip-version-check --no-build-isolation \ + --config-settings "--build-option=--cpp_ext" \ + --config-settings "--build-option=--cuda_ext" \ + git+https://github.com/NVIDIA/apex.git + +RUN export NVTE_FRAMEWORK=pytorch && \ + MAX_JOBS=256 NVTE_BUILD_THREADS_PER_JOB=4 \ + pip3 install --resume-retries 999 --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.12 + +RUN pip install codetiming mathruler pylatexenc qwen_vl_utils cachetools pytest-asyncio + +RUN export FLASH_ATTENTION_FORCE_BUILD="TRUE" && MAX_JOBS=32 pip install --no-build-isolation flash_attn==2.8.3 + +RUN NSIGHT_VERSION=2025.6.1_2025.6.1.190-1_$(if [ "$(uname -m)" = "aarch64" ]; then echo "arm64"; else echo "amd64"; fi) && \ + wget https://developer.nvidia.com/downloads/assets/tools/secure/nsight-systems/2025_6/nsight-systems-${NSIGHT_VERSION}.deb && \ + apt-get update && apt-get install -y libxcb-cursor0 && \ + apt-get install -y ./nsight-systems-${NSIGHT_VERSION}.deb && \ + rm -rf /usr/local/cuda/bin/nsys && \ + ln -s /opt/nvidia/nsight-systems/2025.6.1/nsys /usr/local/cuda/bin/nsys && \ + rm -rf /usr/local/cuda/bin/nsys-ui && \ + ln -s /opt/nvidia/nsight-systems/2025.6.1/nsys-ui /usr/local/cuda/bin/nsys-ui && \ + rm nsight-systems-${NSIGHT_VERSION}.deb && \ + rm -rf /var/lib/apt/lists/* + +# ========================= +# Install DeepEP +# ========================= +RUN cd /home && mkdir -p dpsk_a2a && cd dpsk_a2a && \ + git clone -b v2.5.1 https://github.com/NVIDIA/gdrcopy.git && \ + cd gdrcopy && \ + make prefix=/usr/local lib_install && \ + cd .. && rm -rf gdrcopy && \ + git clone -b hybrid-ep https://github.com/deepseek-ai/DeepEP.git && \ + export NVSHMEM_DIR=/usr/local/lib/python3.12/dist-packages/nvidia/nvshmem && \ + export LD_LIBRARY_PATH="${NVSHMEM_DIR}/lib:$LD_LIBRARY_PATH" && \ + export PATH="${NVSHMEM_DIR}/bin:$PATH" && \ + cd ${NVSHMEM_DIR}/lib && \ + ln -sf libnvshmem_host.so.3 libnvshmem_host.so && \ + cd /home/dpsk_a2a/DeepEP && \ + git checkout 3f601f7ac1c062c46502646ff04c535013bfca00 && \ + CUDA_TARGET=$(uname -m | sed 's/aarch64/sbsa-linux/;s/x86_64/x86_64-linux/') && \ + export CPATH=/usr/local/cuda/targets/${CUDA_TARGET}/include/cccl:$CPATH && \ + TORCH_CUDA_ARCH_LIST="9.0;10.0" python setup.py install + +RUN if [ "$(uname -m)" = "aarch64" ]; then apt-get remove -y python3-jwt; fi && \ + pip install vllm==0.18.0 + +RUN pip3 install --no-deps trl==0.27.0 + +RUN pip3 install nvtx matplotlib liger_kernel + +RUN pip install -U git+https://github.com/ISEEKYAN/mbridge.git@641a5a0 + +RUN pip install --no-deps git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.16.0 + +RUN pip install transformers==5.3.0 + +RUN pip install git+https://github.com/verl-project/verl.git@v0.7.1 && pip uninstall -y verl + +RUN apt-get update && apt-get install -y curl \ + && rm -rf /var/lib/apt/lists/* + +RUN apt-get update && \ + apt-get install -y --allow-downgrades --allow-change-held-packages \ + libcudnn9-cuda-12=9.16.0.29-1 \ + libcudnn9-dev-cuda-12=9.16.0.29-1 \ + libcudnn9-headers-cuda-12=9.16.0.29-1 && \ + rm -rf /var/lib/apt/lists/* diff --git a/verl/docker/README.md b/verl/docker/README.md new file mode 100644 index 0000000000000000000000000000000000000000..728479484b2cb910e12c8e368e6d6dd2eb669e48 --- /dev/null +++ b/verl/docker/README.md @@ -0,0 +1,93 @@ +# Dockerfiles of verl + +We provide pre-built Docker images for quick setup. And from this version, we utilize a new image release hierarchy for productivity and stability. + +Start from v0.6.0, we use vllm and sglang release image as our base image. + +Start from v0.7.0, since vllm/vllm-openai:v0.12.0 is a minimal image without some essential libraries, we use nvidia/cuda:12.9.1-devel-ubuntu22.04 as our base image for vllm. + +## Base Image + +- vLLM: https://hub.docker.com/r/nvidia/cuda +- SGLang: https://hub.docker.com/r/lmsysorg/sglang + +## Application Image + +Upon base image, the following packages are added: +- flash_attn +- Megatron-LM +- Apex +- TransformerEngine +- DeepEP + +Latest docker file: +- [Dockerfile.stable.vllm](https://github.com/verl-project/verl/blob/main/docker/Dockerfile.stable.vllm) +- [Dockerfile.stable.sglang](https://github.com/verl-project/verl/blob/main/docker/Dockerfile.stable.sglang) + +All pre-built images are available in dockerhub: https://hub.docker.com/r/verlai/verl. For example, `verlai/verl:sgl059.latest`, `verlai/verl:vllm017.latest`. + +You can find the latest images used for development and ci in our github workflows: +- [.github/workflows/vllm.yml](https://github.com/verl-project/verl/blob/main/.github/workflows/vllm.yml) +- [.github/workflows/sgl.yml](https://github.com/verl-project/verl/blob/main/.github/workflows/sgl.yml) + + +## Building Locally + +To build an image from source: + +```sh +docker build -f docker/Dockerfile.stable.vllm -t verl:vllm-local . +``` + +For users in China who need an apt mirror to speed up package downloads, pass `APT_MIRROR`: + +```sh +docker build -f docker/Dockerfile.stable.vllm \ + --build-arg APT_MIRROR=https://mirrors.tuna.tsinghua.edu.cn \ + -t verl:vllm-local . +``` + +### GB200 / aarch64 + +Pre-built images for GB200 (aarch64) are not yet published. Users should build locally on an aarch64 machine. Pre-built images will be added once available. + +```sh +docker build -f docker/Dockerfile.stable.vllm -t verl:vllm-arm64 . +``` + +## Installation from Docker + +After pulling the desired Docker image and installing desired inference and training frameworks, you can run it with the following steps: + +1. Launch the desired Docker image and attach into it: + +```sh +docker create --runtime=nvidia --gpus all --net=host --shm-size="10g" --cap-add=SYS_ADMIN -v .:/workspace/verl --name verl sleep infinity +docker start verl +docker exec -it verl bash +``` + +2. If you use the images provided, you only need to install verl itself without dependencies: + +```sh +# install the nightly version (recommended) +git clone https://github.com/verl-project/verl && cd verl +pip3 install --no-deps -e . +``` + +[Optional] If you hope to switch between different frameworks, you can install verl with the following command: + +```sh +# install the nightly version (recommended) +git clone https://github.com/verl-project/verl && cd verl +pip3 install -e .[vllm] +pip3 install -e .[sglang] +``` + +## Release History + +- 2026/03/10: update vllm stable image to vllm==0.17.0; update sglang stable image to sglang==0.5.9 +- 2026/01/17: update vllm stable image to torch==2.9.1, cudnn==9.16, deepep==1.2.1 +- 2025/12/23: update vllm stable image to vllm==0.12.0; update sglang stable image to sglang==0.5.6 +- 2025/11/18: update vllm stable image to vllm==0.11.1; update sglang stable image to sglang==0.5.5 + diff --git a/verl/docker/ascend/Dockerfile.ascend.sglang_8.3.rc1_a2 b/verl/docker/ascend/Dockerfile.ascend.sglang_8.3.rc1_a2 new file mode 100644 index 0000000000000000000000000000000000000000..1b0f46bed2afeb618db0a8402798e30e38351f97 --- /dev/null +++ b/verl/docker/ascend/Dockerfile.ascend.sglang_8.3.rc1_a2 @@ -0,0 +1,84 @@ +# Pull base image +FROM swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.3.rc1-910b-ubuntu22.04-py3.11 + +ARG ASCEND_CANN_PATH="/usr/local/Ascend" +ARG PIP_INDEX_URL="https://mirrors.aliyun.com/pypi/simple" +ARG PTA_BASE_VERSION="torch_npu-2.7.1.post2-cp311-cp311-manylinux_2_28" +ARG PTA_URL="https://gitcode.com/Ascend/pytorch/releases/download/v7.3.0-pytorch2.7.1" + +# Prepare required system dependencies +RUN apt-get update -y && \ + apt-get install -y --no-install-recommends gcc g++ cmake libnuma-dev wget git curl jq vim build-essential net-tools iputils-ping && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \ + pip config set global.index-url ${PIP_INDEX_URL} && \ + pip config set install.trusted-host mirrors.aliyun.com && \ + pip install --upgrade pip setuptools packaging && \ + pip cache purge + +# Prepare repositories with low update frequency +RUN ARCH=$(uname -m) && \ + echo "[LOG INFO] Detected architecture: $ARCH" && \ + # Set extra pip index for x86_64 platform + if [ "$ARCH" = "x86_64" ]; then \ + pip config set global.extra-index-url "https://download.pytorch.org/whl/cpu/"; \ + fi && \ + # Clone libs + git clone --depth 1 --branch v0.5.8 https://github.com/sgl-project/sglang.git && \ + git clone https://github.com/sgl-project/sgl-kernel-npu.git && cd sgl-kernel-npu && git checkout 46b73de && cd .. && \ + git clone https://gitcode.com/Ascend/MindSpeed.git && \ + cd MindSpeed && git checkout f2b0977e && cd .. + +# Install repositories with low update frequency +RUN cd sglang && \ + # Install sglang + mv python/pyproject.toml python/pyproject.toml.backup && \ + mv python/pyproject_other.toml python/pyproject.toml && \ + pip install -e "python[srt_npu]" && \ + pip install torch==2.7.1 torchvision==0.22.1 && \ + # Install torch_npu + ARCH=$(uname -m) && wget ${PTA_URL}/${PTA_BASE_VERSION}_${ARCH}.whl && pip install ${PTA_BASE_VERSION}_${ARCH}.whl && \ + echo "[LOG INFO] Torch_npu version is: ${PTA_BASE_VERSION}_${ARCH}.whl" && \ + cd .. + +# Install sgl-kernel-npu +RUN ARCH=$(uname -m) && \ + # Export and source env + export LD_LIBRARY_PATH=${ASCEND_CANN_PATH}/ascend-toolkit/8.3.RC1/${ARCH}-linux/devlib/linux/${ARCH}:$LD_LIBRARY_PATH && \ + source ${ASCEND_CANN_PATH}/ascend-toolkit/set_env.sh && \ + source ${ASCEND_CANN_PATH}/nnal/atb/set_env.sh && \ + pip install pybind11 && \ + cd sgl-kernel-npu && \ + bash build.sh && \ + pip install output/torch_memory_saver*.whl && \ + pip install output/sgl_kernel_npu*.whl && \ + # Deep_ep package is compiled for A3 by default; Recompile in deepep2 mode for A2, following https://github.com/sgl-project/sgl-kernel-npu/blob/main/python/deep_ep/README.md. + bash build.sh -a deepep2 && \ + pip install output/deep_ep*.whl && \ + cd "$(pip show deep-ep | grep -E '^Location:' | awk '{print $2}')" && ln -s deep_ep/deep_ep_cpp*.so && cd - && \ + cd .. + +# Install MindSpeed & Megatron +RUN pip install -e MindSpeed && \ + pip install git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.12.1 && \ + # Remove existing triton or triton-ascend installed by some third-party packages + pip uninstall -y triton timm && \ + # Install mbridge + pip install mbridge && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + +# Prepare and install verl (update frequently) +RUN git clone --recursive https://github.com/verl-project/verl.git && \ + cd verl && pip install -r requirements-npu.txt && pip install -v -e . && cd .. && \ + pip install ray==2.46.0 click==8.2.1 cachetools setuptools==80.10.2 nvtx && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + +# Show install results +RUN pip list + +# Setting Default Commands +CMD ["/bin/bash"] diff --git a/verl/docker/ascend/Dockerfile.ascend.sglang_8.3.rc1_a3 b/verl/docker/ascend/Dockerfile.ascend.sglang_8.3.rc1_a3 new file mode 100644 index 0000000000000000000000000000000000000000..dc9133764019db8bf7fd8fb669709ed457ddf3ae --- /dev/null +++ b/verl/docker/ascend/Dockerfile.ascend.sglang_8.3.rc1_a3 @@ -0,0 +1,82 @@ +# Pull base image +FROM swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.3.rc1-a3-ubuntu22.04-py3.11 + +ARG ASCEND_CANN_PATH="/usr/local/Ascend" +ARG PIP_INDEX_URL="https://mirrors.aliyun.com/pypi/simple" +ARG PTA_BASE_VERSION="torch_npu-2.7.1.post2-cp311-cp311-manylinux_2_28" +ARG PTA_URL="https://gitcode.com/Ascend/pytorch/releases/download/v7.3.0-pytorch2.7.1" + +# Prepare required system dependencies +RUN apt-get update -y && \ + apt-get install -y --no-install-recommends gcc g++ cmake libnuma-dev wget git curl jq vim build-essential net-tools iputils-ping && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \ + pip config set global.index-url ${PIP_INDEX_URL} && \ + pip config set install.trusted-host mirrors.aliyun.com && \ + pip install --upgrade pip setuptools packaging && \ + pip cache purge + +# Prepare repositories with low update frequency +RUN ARCH=$(uname -m) && \ + echo "[LOG INFO] Detected architecture: $ARCH" && \ + # Set extra pip index for x86_64 platform + if [ "$ARCH" = "x86_64" ]; then \ + pip config set global.extra-index-url "https://download.pytorch.org/whl/cpu/"; \ + fi && \ + # Clone libs + git clone --depth 1 --branch v0.5.8 https://github.com/sgl-project/sglang.git && \ + git clone https://github.com/sgl-project/sgl-kernel-npu.git && cd sgl-kernel-npu && git checkout 46b73de && cd .. && \ + git clone https://gitcode.com/Ascend/MindSpeed.git && \ + cd MindSpeed && git checkout f2b0977e && cd .. + +# Install repositories with low update frequency +RUN cd sglang && \ + # Install sglang + mv python/pyproject.toml python/pyproject.toml.backup && \ + mv python/pyproject_other.toml python/pyproject.toml && \ + pip install -e "python[srt_npu]" && \ + pip install torch==2.7.1 torchvision==0.22.1 && \ + # Install torch_npu + ARCH=$(uname -m) && wget ${PTA_URL}/${PTA_BASE_VERSION}_${ARCH}.whl && pip install ${PTA_BASE_VERSION}_${ARCH}.whl && \ + echo "[LOG INFO] Torch_npu version is: ${PTA_BASE_VERSION}_${ARCH}.whl" && \ + cd .. + +# Install sgl-kernel-npu +RUN ARCH=$(uname -m) && \ + # Export and source env + export LD_LIBRARY_PATH=${ASCEND_CANN_PATH}/ascend-toolkit/8.3.RC1/${ARCH}-linux/devlib/linux/${ARCH}:$LD_LIBRARY_PATH && \ + source ${ASCEND_CANN_PATH}/ascend-toolkit/set_env.sh && \ + source ${ASCEND_CANN_PATH}/nnal/atb/set_env.sh && \ + pip install pybind11 && \ + cd sgl-kernel-npu && \ + bash build.sh && \ + pip install output/torch_memory_saver*.whl && \ + pip install output/sgl_kernel_npu*.whl && \ + pip install output/deep_ep*.whl && \ + cd "$(pip show deep-ep | grep -E '^Location:' | awk '{print $2}')" && ln -s deep_ep/deep_ep_cpp*.so && cd - && \ + cd .. + +# Install MindSpeed & Megatron +RUN pip install -e MindSpeed && \ + pip install git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.12.1 && \ + # Remove existing triton or triton-ascend installed by some third-party packages + pip uninstall -y triton timm && \ + # Install mbridge + pip install mbridge && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + +# Prepare and install verl (update frequently) +RUN git clone --recursive https://github.com/verl-project/verl.git && \ + cd verl && pip install -r requirements-npu.txt && pip install -v -e . && cd .. && \ + pip install ray==2.46.0 click==8.2.1 cachetools setuptools==80.10.2 nvtx && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + +# Show install results +RUN pip list + +# Setting Default Commands +CMD ["/bin/bash"] diff --git a/verl/docker/ascend/Dockerfile.ascend.sglang_8.5.0_a2 b/verl/docker/ascend/Dockerfile.ascend.sglang_8.5.0_a2 new file mode 100644 index 0000000000000000000000000000000000000000..783f44982e4300455cd9722116aba2f774653ff7 --- /dev/null +++ b/verl/docker/ascend/Dockerfile.ascend.sglang_8.5.0_a2 @@ -0,0 +1,74 @@ +# Pull base image +FROM swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.5.0-910b-ubuntu22.04-py3.11 + +ARG PIP_INDEX_URL="https://mirrors.aliyun.com/pypi/simple" +ARG PTA_BASE_VERSION="torch_npu-2.8.0.post2-cp311-cp311-manylinux_2_28" +ARG PTA_URL="https://gitcode.com/Ascend/pytorch/releases/download/v7.3.0-pytorch2.8.0" + +# Prepare required system dependencies +RUN apt-get update -y && \ + apt-get install -y --no-install-recommends gcc g++ cmake libnuma-dev wget git curl jq vim build-essential net-tools iputils-ping unzip ca-certificates && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \ + pip config set global.index-url ${PIP_INDEX_URL} && \ + pip config set install.trusted-host mirrors.aliyun.com && \ + pip install --upgrade pip setuptools==80.10.2 packaging && \ + pip cache purge + +# Prepare repositories with low update frequency +RUN ARCH=$(uname -m) && \ + echo "[LOG INFO] Detected architecture: $ARCH" && \ + # Set extra pip index for x86_64 platform + if [ "$ARCH" = "x86_64" ]; then \ + pip config set global.extra-index-url "https://download.pytorch.org/whl/cpu/"; \ + fi && \ + # Clone libs + git clone https://github.com/sgl-project/sglang.git && cd sglang && git checkout v0.5.10 && cd .. && \ + git clone https://gitcode.com/Ascend/MindSpeed.git && \ + cd MindSpeed && git checkout 2.3.0_core_r0.12.1 && cd .. + +# Install repositories with low update frequency +RUN cd sglang && \ + # Install sglang + mv python/pyproject.toml python/pyproject.toml.backup && \ + mv python/pyproject_npu.toml python/pyproject.toml && \ + pip install torch==2.8.0 torchvision==0.23.0 && \ + # Install torch_npu + ARCH=$(uname -m) && wget ${PTA_URL}/${PTA_BASE_VERSION}_${ARCH}.whl && pip install ${PTA_BASE_VERSION}_${ARCH}.whl && \ + echo "[LOG INFO] Torch_npu version is: ${PTA_BASE_VERSION}_${ARCH}.whl" && \ + pip install -e python[all_npu] && \ + cd .. + +# Install sgl-kernel-npu +RUN ARCH=$(uname -m) && wget --no-check-certificate https://github.com/sgl-project/sgl-kernel-npu/releases/download/2026.02.01/sgl-kernel-npu-2026.02.01-torch2.8.0-py311-cann8.5.0-910b-${ARCH}.zip && \ + unzip sgl-kernel-npu*.zip && \ + pip install torch_memory_saver*.whl && \ + pip install sgl_kernel_npu*.whl && \ + pip install deep_ep*.whl + +# Install MindSpeed & Megatron +RUN pip install -e MindSpeed && \ + pip install git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.12.1 && \ + # Remove existing triton or triton-ascend installed by some third-party packages + pip uninstall -y triton timm && \ + # Install mbridge + pip install mbridge && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + +# Prepare and install verl (update frequently) +RUN git clone --recursive https://github.com/volcengine/verl.git && \ + cd verl/recipe && git checkout main && cd .. && \ + pip install -r requirements-npu.txt && pip install -v -e . && cd .. && \ + pip install click==8.2.1 cachetools nvtx opencv-python-headless==4.10.0.84 && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + +# Show install results +RUN pip list + +# Setting Default Commands +CMD ["/bin/bash"] + diff --git a/verl/docker/ascend/Dockerfile.ascend.sglang_8.5.0_a3 b/verl/docker/ascend/Dockerfile.ascend.sglang_8.5.0_a3 new file mode 100644 index 0000000000000000000000000000000000000000..ae1dc28b5c2e389ec9db27fede39d3eb66935458 --- /dev/null +++ b/verl/docker/ascend/Dockerfile.ascend.sglang_8.5.0_a3 @@ -0,0 +1,74 @@ +# Pull base image +FROM swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.5.0-a3-ubuntu22.04-py3.11 + +ARG PIP_INDEX_URL="https://mirrors.aliyun.com/pypi/simple" +ARG PTA_BASE_VERSION="torch_npu-2.8.0.post2-cp311-cp311-manylinux_2_28" +ARG PTA_URL="https://gitcode.com/Ascend/pytorch/releases/download/v7.3.0-pytorch2.8.0" + +# Prepare required system dependencies +RUN apt-get update -y && \ + apt-get install -y --no-install-recommends gcc g++ cmake libnuma-dev wget git curl jq vim build-essential net-tools iputils-ping unzip ca-certificates && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \ + pip config set global.index-url ${PIP_INDEX_URL} && \ + pip config set install.trusted-host mirrors.aliyun.com && \ + pip install --upgrade pip setuptools==80.10.2 packaging && \ + pip cache purge + +# Prepare repositories with low update frequency +RUN ARCH=$(uname -m) && \ + echo "[LOG INFO] Detected architecture: $ARCH" && \ + # Set extra pip index for x86_64 platform + if [ "$ARCH" = "x86_64" ]; then \ + pip config set global.extra-index-url "https://download.pytorch.org/whl/cpu/"; \ + fi && \ + # Clone libs + git clone https://github.com/sgl-project/sglang.git && cd sglang && git checkout v0.5.10 && cd .. && \ + git clone https://gitcode.com/Ascend/MindSpeed.git && \ + cd MindSpeed && git checkout 2.3.0_core_r0.12.1 && cd .. + +# Install repositories with low update frequency +RUN cd sglang && \ + # Install sglang + mv python/pyproject.toml python/pyproject.toml.backup && \ + mv python/pyproject_npu.toml python/pyproject.toml && \ + pip install torch==2.8.0 torchvision==0.23.0 && \ + # Install torch_npu + ARCH=$(uname -m) && wget ${PTA_URL}/${PTA_BASE_VERSION}_${ARCH}.whl && pip install ${PTA_BASE_VERSION}_${ARCH}.whl && \ + echo "[LOG INFO] Torch_npu version is: ${PTA_BASE_VERSION}_${ARCH}.whl" && \ + pip install -e python[all_npu] && \ + cd .. + +# Install sgl-kernel-npu +RUN ARCH=$(uname -m) && wget --no-check-certificate https://github.com/sgl-project/sgl-kernel-npu/releases/download/2026.02.01/sgl-kernel-npu-2026.02.01-torch2.8.0-py311-cann8.5.0-a3-${ARCH}.zip && \ + unzip sgl-kernel-npu*.zip && \ + pip install torch_memory_saver*.whl && \ + pip install sgl_kernel_npu*.whl && \ + pip install deep_ep*.whl + +# Install MindSpeed & Megatron +RUN pip install -e MindSpeed && \ + pip install git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.12.1 && \ + # Remove existing triton or triton-ascend installed by some third-party packages + pip uninstall -y triton timm && \ + # Install mbridge + pip install mbridge && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + +# Prepare and install verl (update frequently) +RUN git clone --recursive https://github.com/volcengine/verl.git && \ + cd verl/recipe && git checkout main && cd .. && \ + pip install -r requirements-npu.txt && pip install -v -e . && cd .. && \ + pip install click==8.2.1 cachetools nvtx opencv-python-headless==4.10.0.84 && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + +# Show install results +RUN pip list + +# Setting Default Commands +CMD ["/bin/bash"] + diff --git a/verl/docker/ascend/Dockerfile.ascend_8.2.rc1_a2 b/verl/docker/ascend/Dockerfile.ascend_8.2.rc1_a2 new file mode 100644 index 0000000000000000000000000000000000000000..c60df0c1712159e243d090ac35ea720186fb1c1b --- /dev/null +++ b/verl/docker/ascend/Dockerfile.ascend_8.2.rc1_a2 @@ -0,0 +1,61 @@ +# Pull base image +FROM swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.2.rc1-910b-ubuntu22.04-py3.11 + +# Prepare required system dependencies +RUN apt-get update -y && \ + apt-get install -y --no-install-recommends gcc g++ cmake libnuma-dev wget git curl jq vim build-essential && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \ + pip install --upgrade pip setuptools packaging && \ + pip cache purge + +# Prepare repositories with low update frequency +RUN ARCH=$(uname -m) && \ + # Set extra pip index for x86_64 platform + echo "[LOG INFO] Detected architecture: $ARCH" && \ + if [ "$ARCH" = "x86_64" ]; then \ + pip config set global.extra-index-url "https://download.pytorch.org/whl/cpu/"; \ + fi && \ + # Clone libs + git clone --depth 1 --branch v0.9.1 https://github.com/vllm-project/vllm && \ + git clone --depth 1 --branch v0.9.1 https://github.com/vllm-project/vllm-ascend.git && \ + git clone https://gitcode.com/Ascend/MindSpeed.git && \ + cd MindSpeed && git checkout f2b0977e && cd .. && \ + git clone --depth 1 --branch core_v0.12.1 https://github.com/NVIDIA/Megatron-LM.git + +# Install repositories with low update frequency +RUN ARCH=$(uname -m) && \ + # Export and source env + if [ "$ARCH" = "aarch64" ]; then \ + export LD_LIBRARY_PATH=/usr/local/Ascend/ascend-toolkit/8.2.RC1/aarch64-linux/devlib/linux/aarch64:$LD_LIBRARY_PATH; \ + elif [ "$ARCH" = "x86_64" ]; then \ + export LD_LIBRARY_PATH=/usr/local/Ascend/ascend-toolkit/8.2.RC1/x86_64-linux/devlib/linux/x86_64/:$LD_LIBRARY_PATH; \ + fi && \ + source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ + source /usr/local/Ascend/nnal/atb/set_env.sh && \ + # Install torch & torch_npu & torchvision + pip install torch==2.5.1 torch_npu==2.5.1 torchvision==0.20.1 && \ + # Install vllm + cd vllm && VLLM_TARGET_DEVICE=empty pip install -v -e . && cd .. && \ + # Install vllm-ascend + cd vllm-ascend && pip install -v -e . && cd .. && \ + # Install MindSpeed & Megatron + pip install -e MindSpeed && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + +ENV PYTHONPATH="/Megatron-LM${PYTHONPATH:+:${PYTHONPATH}}" + +# Prepare and install verl (update frequently) +RUN git clone --depth 1 https://github.com/verl-project/verl.git && \ + cd verl && pip install -r requirements-npu.txt && pip install -v -e . && cd .. && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + +# Show install results +RUN pip list + +# Setting Default Commands +CMD ["/bin/bash"] diff --git a/verl/docker/ascend/Dockerfile.ascend_8.2.rc1_a3 b/verl/docker/ascend/Dockerfile.ascend_8.2.rc1_a3 new file mode 100644 index 0000000000000000000000000000000000000000..15850d04577e7bd8d2c1dbfa1b2a8b4aaa343e4f --- /dev/null +++ b/verl/docker/ascend/Dockerfile.ascend_8.2.rc1_a3 @@ -0,0 +1,61 @@ +# Pull base image +FROM swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.2.rc1-a3-ubuntu22.04-py3.11 + +# Prepare required system dependencies +RUN apt-get update -y && \ + apt-get install -y --no-install-recommends gcc g++ cmake libnuma-dev wget git curl jq vim build-essential && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \ + pip install --upgrade pip setuptools packaging && \ + pip cache purge + +# Prepare repositories with low update frequency +RUN ARCH=$(uname -m) && \ + # Set extra pip index for x86_64 platform + echo "[LOG INFO] Detected architecture: $ARCH" && \ + if [ "$ARCH" = "x86_64" ]; then \ + pip config set global.extra-index-url "https://download.pytorch.org/whl/cpu/"; \ + fi && \ + # Clone libs + git clone --depth 1 --branch v0.9.1 https://github.com/vllm-project/vllm && \ + git clone --depth 1 --branch v0.9.1 https://github.com/vllm-project/vllm-ascend.git && \ + git clone https://gitcode.com/Ascend/MindSpeed.git && \ + cd MindSpeed && git checkout f2b0977e && cd .. && \ + git clone --depth 1 --branch core_v0.12.1 https://github.com/NVIDIA/Megatron-LM.git + +# Install repositories with low update frequency +RUN ARCH=$(uname -m) && \ + # Export and source env + if [ "$ARCH" = "aarch64" ]; then \ + export LD_LIBRARY_PATH=/usr/local/Ascend/ascend-toolkit/8.2.RC1/aarch64-linux/devlib/linux/aarch64:$LD_LIBRARY_PATH; \ + elif [ "$ARCH" = "x86_64" ]; then \ + export LD_LIBRARY_PATH=/usr/local/Ascend/ascend-toolkit/8.2.RC1/x86_64-linux/devlib/linux/x86_64/:$LD_LIBRARY_PATH; \ + fi && \ + source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ + source /usr/local/Ascend/nnal/atb/set_env.sh && \ + # Install torch & torch_npu & torchvision + pip install torch==2.5.1 torch_npu==2.5.1 torchvision==0.20.1 && \ + # Install vllm + cd vllm && VLLM_TARGET_DEVICE=empty pip install -v -e . && cd .. && \ + # Install vllm-ascend + cd vllm-ascend && pip install -v -e . && cd .. && \ + # Install MindSpeed & Megatron + pip install -e MindSpeed && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + +ENV PYTHONPATH="/Megatron-LM${PYTHONPATH:+:${PYTHONPATH}}" + +# Prepare and install verl (update frequently) +RUN git clone --depth 1 https://github.com/verl-project/verl.git && \ + cd verl && pip install -r requirements-npu.txt && pip install -v -e . && cd .. && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + +# Show install results +RUN pip list + +# Setting Default Commands +CMD ["/bin/bash"] diff --git a/verl/docker/ascend/Dockerfile.ascend_8.3.rc1_a2 b/verl/docker/ascend/Dockerfile.ascend_8.3.rc1_a2 new file mode 100644 index 0000000000000000000000000000000000000000..013df8ff3c78eb9cb7f842fe6657e682bd7c7f45 --- /dev/null +++ b/verl/docker/ascend/Dockerfile.ascend_8.3.rc1_a2 @@ -0,0 +1,65 @@ +# Pull base image +FROM swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.3.rc1-910b-ubuntu22.04-py3.11 + +# Prepare required system dependencies +RUN apt-get update -y && \ + apt-get install -y --no-install-recommends gcc g++ cmake libnuma-dev wget git curl jq vim build-essential && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \ + pip install --upgrade pip packaging setuptools==80.10.2 && \ + pip cache purge + +# Prepare repositories with low update frequency +RUN ARCH=$(uname -m) && \ + # Set extra pip index for x86_64 platform + echo "[LOG INFO] Detected architecture: $ARCH" && \ + if [ "$ARCH" = "x86_64" ]; then \ + pip config set global.extra-index-url "https://download.pytorch.org/whl/cpu/"; \ + fi && \ + # Clone libs + git clone --depth 1 --branch v0.11.0 https://github.com/vllm-project/vllm.git && \ + git clone --depth 1 --branch v0.11.0 https://github.com/vllm-project/vllm-ascend.git && \ + git clone https://gitcode.com/Ascend/MindSpeed.git && \ + cd MindSpeed && git checkout f2b0977e && cd .. && \ + git clone --depth 1 --branch core_v0.12.1 https://github.com/NVIDIA/Megatron-LM.git + +# Install repositories with low update frequency +RUN ARCH=$(uname -m) && \ + # Export and source env + if [ "$ARCH" = "aarch64" ]; then \ + export LD_LIBRARY_PATH=/usr/local/Ascend/ascend-toolkit/8.3.RC1/aarch64-linux/devlib/linux/aarch64:$LD_LIBRARY_PATH; \ + elif [ "$ARCH" = "x86_64" ]; then \ + export LD_LIBRARY_PATH=/usr/local/Ascend/ascend-toolkit/8.3.RC1/x86_64-linux/devlib/linux/x86_64/:$LD_LIBRARY_PATH; \ + fi && \ + source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ + source /usr/local/Ascend/nnal/atb/set_env.sh && \ + # Install torch & torch_npu & torchvision + pip install torch==2.7.1 torch_npu==2.7.1 torchvision==0.22.1 transformers==4.57.6 && \ + # Install vllm + cd vllm && VLLM_TARGET_DEVICE=empty pip install -v -e . && cd .. && \ + # Install vllm-ascend + cd vllm-ascend && pip install -v -e . && cd .. && \ + # Install MindSpeed & Megatron + pip install -e MindSpeed && \ + # Remove existing triton or triton-ascend installed by some third-party packages + pip uninstall -y triton triton-ascend && \ + # Install mbridge + pip install mbridge && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + +ENV PYTHONPATH="/Megatron-LM${PYTHONPATH:+:${PYTHONPATH}}" + +# Prepare and install verl (update frequently) +RUN git clone --depth 1 https://github.com/verl-project/verl.git && \ + cd verl && pip install -r requirements-npu.txt && pip install -v -e . && cd .. && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + +# Show install results +RUN pip list + +# Setting Default Commands +CMD ["/bin/bash"] diff --git a/verl/docker/ascend/Dockerfile.ascend_8.3.rc1_a3 b/verl/docker/ascend/Dockerfile.ascend_8.3.rc1_a3 new file mode 100644 index 0000000000000000000000000000000000000000..abf5cbfe91753acc731819cbd6d0891b7d1fdbaa --- /dev/null +++ b/verl/docker/ascend/Dockerfile.ascend_8.3.rc1_a3 @@ -0,0 +1,65 @@ +# Pull base image +FROM swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.3.rc1-a3-ubuntu22.04-py3.11 + +# Prepare required system dependencies +RUN apt-get update -y && \ + apt-get install -y --no-install-recommends gcc g++ cmake libnuma-dev wget git curl jq vim build-essential && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \ + pip install --upgrade pip packaging setuptools==80.10.2 && \ + pip cache purge + +# Prepare repositories with low update frequency +RUN ARCH=$(uname -m) && \ + # Set extra pip index for x86_64 platform + echo "[LOG INFO] Detected architecture: $ARCH" && \ + if [ "$ARCH" = "x86_64" ]; then \ + pip config set global.extra-index-url "https://download.pytorch.org/whl/cpu/"; \ + fi && \ + # Clone libs + git clone --depth 1 --branch v0.11.0 https://github.com/vllm-project/vllm.git && \ + git clone --depth 1 --branch v0.11.0 https://github.com/vllm-project/vllm-ascend.git && \ + git clone https://gitcode.com/Ascend/MindSpeed.git && \ + cd MindSpeed && git checkout f2b0977e && cd .. && \ + git clone --depth 1 --branch core_v0.12.1 https://github.com/NVIDIA/Megatron-LM.git + +# Install repositories with low update frequency +RUN ARCH=$(uname -m) && \ + # Export and source env + if [ "$ARCH" = "aarch64" ]; then \ + export LD_LIBRARY_PATH=/usr/local/Ascend/ascend-toolkit/8.3.RC1/aarch64-linux/devlib/linux/aarch64:$LD_LIBRARY_PATH; \ + elif [ "$ARCH" = "x86_64" ]; then \ + export LD_LIBRARY_PATH=/usr/local/Ascend/ascend-toolkit/8.3.RC1/x86_64-linux/devlib/linux/x86_64/:$LD_LIBRARY_PATH; \ + fi && \ + source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ + source /usr/local/Ascend/nnal/atb/set_env.sh && \ + # Install torch & torch_npu & torchvision + pip install torch==2.7.1 torch_npu==2.7.1 torchvision==0.22.1 transformers==4.57.6 && \ + # Install vllm + cd vllm && VLLM_TARGET_DEVICE=empty pip install -v -e . && cd .. && \ + # Install vllm-ascend + cd vllm-ascend && pip install -v -e . && cd .. && \ + # Install MindSpeed & Megatron + pip install -e MindSpeed && \ + # Remove existing triton or triton-ascend installed by some third-party packages + pip uninstall -y triton triton-ascend && \ + # Install mbridge + pip install mbridge && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + +ENV PYTHONPATH="/Megatron-LM${PYTHONPATH:+:${PYTHONPATH}}" + +# Prepare and install verl (update frequently) +RUN git clone --depth 1 https://github.com/verl-project/verl.git && \ + cd verl && pip install -r requirements-npu.txt && pip install -v -e . && cd .. && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + +# Show install results +RUN pip list + +# Setting Default Commands +CMD ["/bin/bash"] diff --git a/verl/docker/ascend/Dockerfile.ascend_8.5.0_a2 b/verl/docker/ascend/Dockerfile.ascend_8.5.0_a2 new file mode 100644 index 0000000000000000000000000000000000000000..0693ce1fccba517c66602ac606877c11d0c124f0 --- /dev/null +++ b/verl/docker/ascend/Dockerfile.ascend_8.5.0_a2 @@ -0,0 +1,69 @@ +# Pull base image +FROM swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.5.0-910b-ubuntu22.04-py3.11 + +ARG SOC_VERSION="ascend910b1" + +# Prepare required system dependencies +RUN apt-get update -y && \ + apt-get install -y --no-install-recommends gcc g++ cmake libnuma-dev wget git curl jq vim build-essential && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \ + pip install --upgrade pip packaging setuptools==80.10.2 && \ + pip cache purge + +# Prepare repositories with low update frequency +RUN ARCH=$(uname -m) && \ + # Set extra pip index for x86_64 platform + echo "[LOG INFO] Detected architecture: $ARCH" && \ + if [ "$ARCH" = "x86_64" ]; then \ + pip config set global.extra-index-url "https://download.pytorch.org/whl/cpu/"; \ + fi && \ + # Clone libs + git clone --depth 1 --branch v0.18.0 https://github.com/vllm-project/vllm.git && \ + git clone -b releases/v0.18.0 https://github.com/vllm-project/vllm-ascend.git && \ + git clone https://gitcode.com/Ascend/MindSpeed.git && \ + cd MindSpeed && git checkout 2.3.0_core_r0.12.1 && cd .. && \ + git clone --depth 1 --branch core_v0.12.1 https://github.com/NVIDIA/Megatron-LM.git + +# Install repositories with low update frequency +RUN ARCH=$(uname -m) && \ + # Export and source env + if [ "$ARCH" = "aarch64" ]; then \ + export LD_LIBRARY_PATH=/usr/local/Ascend/cann-8.5.0/aarch64-linux/devlib/linux/aarch64:$LD_LIBRARY_PATH; \ + elif [ "$ARCH" = "x86_64" ]; then \ + export LD_LIBRARY_PATH=/usr/local/Ascend/cann-8.5.0/x86_64-linux/devlib/linux/x86_64/:$LD_LIBRARY_PATH; \ + fi && \ + source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ + source /usr/local/Ascend/nnal/atb/set_env.sh && \ + # Install transformers + pip install transformers==4.57.6 && \ + # Install vllm + cd vllm && pip install -r requirements/build.txt && \ + VLLM_TARGET_DEVICE=empty pip install -v -e. && cd .. && \ + # Install vllm-ascend + cd vllm-ascend && pip install -r requirements.txt && \ + export COMPILE_CUSTOM_KERNELS=1 && pip install -v -e . && cd .. && \ + # Install MindSpeed & Megatron + pip install -e MindSpeed && \ + pip install -e Megatron-LM && \ + # Remove existing triton installed by some third-party packages + pip uninstall -y triton && \ + # Install mbridge + pip install mbridge && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + + +# Prepare and install verl (update frequently) +RUN git clone --recursive https://github.com/verl-project/verl.git && \ + cd verl && pip install -r requirements-npu.txt && pip install -v -e . && cd .. && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + +# Show install results +RUN pip list + +# Setting Default Commands +CMD ["/bin/bash"] \ No newline at end of file diff --git a/verl/docker/ascend/Dockerfile.ascend_8.5.0_a2_v0.7.1 b/verl/docker/ascend/Dockerfile.ascend_8.5.0_a2_v0.7.1 new file mode 100644 index 0000000000000000000000000000000000000000..9b1a474786a6822dd94c9011c71cac96e2f6a06a --- /dev/null +++ b/verl/docker/ascend/Dockerfile.ascend_8.5.0_a2_v0.7.1 @@ -0,0 +1,75 @@ +# Pull base image +FROM swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.5.0-910b-ubuntu22.04-py3.11 + +ARG SOC_VERSION="ascend910b1" + +# Prepare required system dependencies +RUN apt-get update -y && \ + apt-get install -y --no-install-recommends gcc g++ cmake libnuma-dev wget git curl jq vim build-essential && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \ + pip install --upgrade pip packaging setuptools==80.10.2 && \ + pip cache purge + +# Prepare repositories with low update frequency +RUN ARCH=$(uname -m) && \ + # Set extra pip index for x86_64 platform + echo "[LOG INFO] Detected architecture: $ARCH" && \ + if [ "$ARCH" = "x86_64" ]; then \ + pip config set global.extra-index-url "https://download.pytorch.org/whl/cpu/"; \ + fi && \ + # Clone libs + git clone --depth 1 --branch v0.13.0 https://github.com/vllm-project/vllm.git && \ + git clone -b releases/v0.13.0 https://github.com/vllm-project/vllm-ascend.git && \ + git clone https://gitcode.com/Ascend/MindSpeed.git && \ + cd MindSpeed && git checkout 2.3.0_core_r0.12.1 && cd .. && \ + git clone --depth 1 --branch core_v0.12.1 https://github.com/NVIDIA/Megatron-LM.git + +# Install repositories with low update frequency +RUN ARCH=$(uname -m) && \ + # Export and source env + if [ "$ARCH" = "aarch64" ]; then \ + export LD_LIBRARY_PATH=/usr/local/Ascend/cann-8.5.0/aarch64-linux/devlib/linux/aarch64:$LD_LIBRARY_PATH; \ + elif [ "$ARCH" = "x86_64" ]; then \ + export LD_LIBRARY_PATH=/usr/local/Ascend/cann-8.5.0/x86_64-linux/devlib/linux/x86_64/:$LD_LIBRARY_PATH; \ + fi && \ + source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ + source /usr/local/Ascend/nnal/atb/set_env.sh && \ + # Install transformers + pip install transformers==4.57.6 && \ + # Install vllm + cd vllm && pip install -r requirements/build.txt && \ + VLLM_TARGET_DEVICE=empty pip install -v -e. && cd .. && \ + # Install vllm-ascend + cd vllm-ascend && pip install -r requirements.txt && \ + export COMPILE_CUSTOM_KERNELS=1 && pip install -v -e . && cd .. && \ + # Install MindSpeed & Megatron + pip install -e MindSpeed && \ + pip install -e Megatron-LM && \ + # Remove existing triton installed by some third-party packages + pip uninstall -y triton && \ + # Install mbridge + pip install mbridge==0.15.1 torch_npu==2.8.0 && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + + +# Prepare and install verl (recipe update frequently, verl is release/v0.7.1) +RUN git clone -b release/v0.7.1 https://github.com/verl-project/verl.git && cd verl && \ + git submodule update --init --recursive --remote && \ + pip install -r requirements-npu.txt && pip install -v -e . && cd .. && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + +# get mbridge path and cuda to npu for deepseek model +RUN MBRIDGE_PATH=$(pip show mbridge | grep Location | awk '{print $2}') && \ + TARGET_FILE="${MBRIDGE_PATH}/mbridge/models/ext/deepseek_v3/dequant_fp8_safetensor_io.py" && \ + sed -i '34s/cuda/npu/;51s/cuda/npu/' "$TARGET_FILE" + +# Show install results +RUN pip list + +# Setting Default Commands +CMD ["/bin/bash"] diff --git a/verl/docker/ascend/Dockerfile.ascend_8.5.0_a3 b/verl/docker/ascend/Dockerfile.ascend_8.5.0_a3 new file mode 100644 index 0000000000000000000000000000000000000000..1d3b45f308d108956c4d3c6508b5ba7052175a85 --- /dev/null +++ b/verl/docker/ascend/Dockerfile.ascend_8.5.0_a3 @@ -0,0 +1,74 @@ +# Pull base image +FROM swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.5.0-a3-ubuntu22.04-py3.11 + +ARG SOC_VERSION="ascend910_9392" + +ENV GIT_SSL_NO_VERIFY=true +RUN pip config --user set global.index https://mirrors.huaweicloud.com/repository/pypi && \ + pip config --user set global.index-url https://mirrors.huaweicloud.com/repository/pypi/simple && \ + pip config --user set global.trusted-host mirrors.huaweicloud.com + +# Prepare required system dependencies +RUN apt-get update -y && \ + apt-get install -y --no-install-recommends gcc g++ cmake libnuma-dev wget git curl jq vim build-essential && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \ + pip install --upgrade pip packaging setuptools==80.10.2 && \ + pip cache purge + +# Prepare repositories with low update frequency +RUN ARCH=$(uname -m) && \ + # Set extra pip index for x86_64 platform + echo "[LOG INFO] Detected architecture: $ARCH" && \ + if [ "$ARCH" = "x86_64" ]; then \ + pip config set global.extra-index-url "https://download.pytorch.org/whl/cpu/"; \ + fi && \ + # Clone libs + git clone --depth 1 --branch v0.18.0 https://github.com/vllm-project/vllm.git && \ + git clone -b releases/v0.18.0 https://github.com/vllm-project/vllm-ascend.git && \ + git clone https://gitcode.com/Ascend/MindSpeed.git && \ + cd MindSpeed && git checkout 2.3.0_core_r0.12.1 && cd .. && \ + git clone --depth 1 --branch core_v0.12.1 https://github.com/NVIDIA/Megatron-LM.git + +# Install repositories with low update frequency +RUN ARCH=$(uname -m) && \ + # Export and source env + if [ "$ARCH" = "aarch64" ]; then \ + export LD_LIBRARY_PATH=/usr/local/Ascend/cann-8.5.0/aarch64-linux/devlib/linux/aarch64:$LD_LIBRARY_PATH; \ + elif [ "$ARCH" = "x86_64" ]; then \ + export LD_LIBRARY_PATH=/usr/local/Ascend/cann-8.5.0/x86_64-linux/devlib/linux/x86_64/:$LD_LIBRARY_PATH; \ + fi && \ + source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ + source /usr/local/Ascend/nnal/atb/set_env.sh && \ + # Install transformers + pip install transformers==4.57.6 && \ + # Install vllm + cd vllm && pip install -r requirements/build.txt && \ + VLLM_TARGET_DEVICE=empty pip install -v -e. && cd .. && \ + # Install vllm-ascend + cd vllm-ascend && pip install -r requirements.txt && \ + export COMPILE_CUSTOM_KERNELS=1 && pip install -v -e . && cd .. && \ + # Install MindSpeed & Megatron + pip install -e MindSpeed && \ + pip install -e Megatron-LM && \ + # Remove existing triton installed by some third-party packages + pip uninstall -y triton && \ + # Install mbridge + pip install mbridge && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + + +# Prepare and install verl (update frequently) +RUN git clone --recursive https://github.com/verl-project/verl.git && \ + cd verl && pip install -r requirements-npu.txt && pip install -v -e . && cd .. && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + +# Show install results +RUN pip list + +# Setting Default Commands +CMD ["/bin/bash"] diff --git a/verl/docker/ascend/Dockerfile.ascend_8.5.0_a3_v0.7.1 b/verl/docker/ascend/Dockerfile.ascend_8.5.0_a3_v0.7.1 new file mode 100644 index 0000000000000000000000000000000000000000..966b7e4ba4883b1df9e1e9dcac745517f0686b3f --- /dev/null +++ b/verl/docker/ascend/Dockerfile.ascend_8.5.0_a3_v0.7.1 @@ -0,0 +1,75 @@ +# Pull base image +FROM swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.5.0-a3-ubuntu22.04-py3.11 + +ARG SOC_VERSION="ascend910_9392" + +# Prepare required system dependencies +RUN apt-get update -y && \ + apt-get install -y --no-install-recommends gcc g++ cmake libnuma-dev wget git curl jq vim build-essential && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \ + pip install --upgrade pip packaging setuptools==80.10.2 && \ + pip cache purge + +# Prepare repositories with low update frequency +RUN ARCH=$(uname -m) && \ + # Set extra pip index for x86_64 platform + echo "[LOG INFO] Detected architecture: $ARCH" && \ + if [ "$ARCH" = "x86_64" ]; then \ + pip config set global.extra-index-url "https://download.pytorch.org/whl/cpu/"; \ + fi && \ + # Clone libs + git clone --depth 1 --branch v0.13.0 https://github.com/vllm-project/vllm.git && \ + git clone -b releases/v0.13.0 https://github.com/vllm-project/vllm-ascend.git && \ + git clone https://gitcode.com/Ascend/MindSpeed.git && \ + cd MindSpeed && git checkout 2.3.0_core_r0.12.1 && cd .. && \ + git clone --depth 1 --branch core_v0.12.1 https://github.com/NVIDIA/Megatron-LM.git + +# Install repositories with low update frequency +RUN ARCH=$(uname -m) && \ + # Export and source env + if [ "$ARCH" = "aarch64" ]; then \ + export LD_LIBRARY_PATH=/usr/local/Ascend/cann-8.5.0/aarch64-linux/devlib/linux/aarch64:$LD_LIBRARY_PATH; \ + elif [ "$ARCH" = "x86_64" ]; then \ + export LD_LIBRARY_PATH=/usr/local/Ascend/cann-8.5.0/x86_64-linux/devlib/linux/x86_64/:$LD_LIBRARY_PATH; \ + fi && \ + source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ + source /usr/local/Ascend/nnal/atb/set_env.sh && \ + # Install transformers + pip install transformers==4.57.6 && \ + # Install vllm + cd vllm && pip install -r requirements/build.txt && \ + VLLM_TARGET_DEVICE=empty pip install -v -e. && cd .. && \ + # Install vllm-ascend + cd vllm-ascend && pip install -r requirements.txt && \ + export COMPILE_CUSTOM_KERNELS=1 && pip install -v -e . && cd .. && \ + # Install MindSpeed & Megatron + pip install -e MindSpeed && \ + pip install -e Megatron-LM && \ + # Remove existing triton installed by some third-party packages + pip uninstall -y triton && \ + # Install mbridge + pip install mbridge torch_npu==2.8.0 && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + + +# Prepare and install verl (recipe update frequently, verl is release/v0.7.1) +RUN git clone -b release/v0.7.1 https://github.com/verl-project/verl.git && cd verl && \ + git submodule update --init --recursive --remote && \ + pip install -r requirements-npu.txt && pip install -v -e . && cd .. && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + +# get mbridge path and cuda to npu for deepseek model +RUN MBRIDGE_PATH=$(pip show mbridge | grep Location | awk '{print $2}') && \ + TARGET_FILE="${MBRIDGE_PATH}/mbridge/models/ext/deepseek_v3/dequant_fp8_safetensor_io.py" && \ + sed -i '34s/cuda/npu/;51s/cuda/npu/' "$TARGET_FILE" + +# Show install results +RUN pip list + +# Setting Default Commands +CMD ["/bin/bash"] diff --git a/verl/docker/ascend/Dockerfile.ascend_8.5.2_a2_qwen3-5 b/verl/docker/ascend/Dockerfile.ascend_8.5.2_a2_qwen3-5 new file mode 100644 index 0000000000000000000000000000000000000000..28110c68c82907d38eb0d52f9df562b342f56ff9 --- /dev/null +++ b/verl/docker/ascend/Dockerfile.ascend_8.5.2_a2_qwen3-5 @@ -0,0 +1,105 @@ +# Pull base image +FROM swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.5.1-910b-ubuntu22.04-py3.11 + +ARG SOC_VERSION="ascend910b1" + +# Prepare required system dependencies +RUN apt-get update -y && \ + apt-get install -y --no-install-recommends gcc g++ cmake libnuma-dev wget git curl jq vim build-essential && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \ + pip install --upgrade pip packaging setuptools==80.10.2 && \ + pip cache purge + +# 更新cann版本 +RUN set -e && \ + cd /tmp && \ + ARCH=$(uname -m) && \ + if [ "$ARCH" = "aarch64" ]; then \ + wget -v --no-check-certificate https://ascend-repo.obs.cn-east-2.myhuaweicloud.com/CANN/CANN%208.5.2/Ascend-cann-toolkit_8.5.2_linux-aarch64.run && \ + bash Ascend-cann-toolkit_8.5.2_linux-aarch64.run --quiet --install && \ + rm -f Ascend-cann-toolkit_8.5.2_linux-aarch64.run && \ + + wget -v --no-check-certificate https://ascend-repo.obs.cn-east-2.myhuaweicloud.com/CANN/CANN%208.5.2/Ascend-cann-910b-ops_8.5.2_linux-aarch64.run && \ + bash Ascend-cann-910b-ops_8.5.2_linux-aarch64.run --quiet --install && \ + rm -f Ascend-cann-910b-ops_8.5.2_linux-aarch64.run && \ + + source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ + wget -v --no-check-certificate https://ascend-repo.obs.cn-east-2.myhuaweicloud.com/CANN/CANN%208.5.2/Ascend-cann-nnal_8.5.2_linux-aarch64.run && \ + bash Ascend-cann-nnal_8.5.2_linux-aarch64.run --quiet --install && \ + rm -f Ascend-cann-nnal_8.5.2_linux-aarch64.run ; \ + elif [ "$ARCH" = "x86_64" ]; then \ + wget -v --no-check-certificate https://ascend-repo.obs.cn-east-2.myhuaweicloud.com/CANN/CANN%208.5.2/Ascend-cann-toolkit_8.5.2_linux-x86_64.run && \ + bash Ascend-cann-toolkit_8.5.2_linux-x86_64.run --quiet --install && \ + rm -f Ascend-cann-toolkit_8.5.2_linux-x86_64.run && \ + + wget -v --no-check-certificate https://ascend-repo.obs.cn-east-2.myhuaweicloud.com/CANN/CANN%208.5.2/Ascend-cann-910b-ops_8.5.2_linux-x86_64.run && \ + bash Ascend-cann-910b-ops_8.5.2_linux-x86_64.run --quiet --install && \ + rm -f Ascend-cann-910b-ops_8.5.2_linux-x86_64.run && \ + + source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ + wget -v --no-check-certificate https://ascend-repo.obs.cn-east-2.myhuaweicloud.com/CANN/CANN%208.5.2/Ascend-cann-nnal_8.5.2_linux-x86_64.run && \ + bash Ascend-cann-nnal_8.5.2_linux-x86_64.run --quiet --install && \ + rm -f Ascend-cann-nnal_8.5.2_linux-x86_64.run ; \ + fi && \ + + source /usr/local/Ascend/cann/set_env.sh && \ + rm -rf /tmp/* + +# Prepare repositories with low update frequency +RUN ARCH=$(uname -m) && \ + # Set extra pip index for x86_64 platform + echo "[LOG INFO] Detected architecture: $ARCH" && \ + if [ "$ARCH" = "x86_64" ]; then \ + pip config set global.extra-index-url "https://download.pytorch.org/whl/cpu/"; \ + fi && \ + # Clone libs + git clone --depth 1 --branch v0.18.0 https://github.com/vllm-project/vllm.git && \ + git clone https://github.com/vllm-project/vllm-ascend.git && \ + cd vllm-ascend && git checkout 54879467c41784a446aa5b486a391d9bfbf488fa && cd .. && \ + git clone https://github.com/huggingface/transformers.git && \ + cd transformers && git checkout cc7ab9be508ce6ed3637bba9e50367b29b742dc6 && cd .. + +# 避免vllm-ascend 安装失败 +ENV MAX_JOBS=1 +ENV MAKEFLAGS="-j1" +ENV CMAKE_BUILD_PARALLEL_LEVEL=1 + +# Install repositories with low update frequency +RUN ARCH=$(uname -m) && \ + # Export and source env + if [ "$ARCH" = "aarch64" ]; then \ + export LD_LIBRARY_PATH=/usr/local/Ascend/cann-8.5.2/aarch64-linux/devlib/linux/aarch64:$LD_LIBRARY_PATH; \ + elif [ "$ARCH" = "x86_64" ]; then \ + export LD_LIBRARY_PATH=/usr/local/Ascend/cann-8.5.2/x86_64-linux/devlib/linux/x86_64/:$LD_LIBRARY_PATH; \ + fi && \ + source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ + source /usr/local/Ascend/nnal/atb/set_env.sh && \ + + # Install vllm + cd vllm && pip install -r requirements/build.txt && \ + VLLM_TARGET_DEVICE=empty pip install -v -e. && cd .. && \ + # Install vllm-ascend + cd vllm-ascend && pip install -r requirements.txt && \ + export COMPILE_CUSTOM_KERNELS=1 && pip install -v -e . --no-build-isolation && cd .. && \ + # Install transformers + cd transformers && pip install -e . && cd .. && \ + # Install torch & torch_npu & torchvision + pip install torch==2.9.0 torch_npu==2.9.0 torchvision==0.24.0 accelerate==1.13.0 mathruler && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + + +# Prepare and install verl (update frequently) +RUN git clone https://github.com/verl-project/verl.git && \ + cd verl && git checkout 4045d67063052dcb800c918c107b8d5a87046006 && pip install -r requirements-npu.txt && pip install -v -e . && cd .. && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + +# Show install results +RUN pip list + +# Setting Default Commands +CMD ["/bin/bash"] \ No newline at end of file diff --git a/verl/docker/ascend/Dockerfile.ascend_8.5.2_a3_qwen3-5 b/verl/docker/ascend/Dockerfile.ascend_8.5.2_a3_qwen3-5 new file mode 100644 index 0000000000000000000000000000000000000000..98ee06dafad93b589f80e77f04454b844c9463bc --- /dev/null +++ b/verl/docker/ascend/Dockerfile.ascend_8.5.2_a3_qwen3-5 @@ -0,0 +1,104 @@ +# Pull base image +FROM swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.5.1-a3-ubuntu22.04-py3.11 + +ARG SOC_VERSION="ascend910_9392" + +RUN set -e && \ + cd /tmp && \ + ARCH=$(uname -m) && \ + if [ "$ARCH" = "aarch64" ]; then \ + wget -v --no-check-certificate https://ascend-repo.obs.cn-east-2.myhuaweicloud.com/CANN/CANN%208.5.2/Ascend-cann-toolkit_8.5.2_linux-aarch64.run && \ + bash Ascend-cann-toolkit_8.5.2_linux-aarch64.run --quiet --install && \ + rm -f Ascend-cann-toolkit_8.5.2_linux-aarch64.run && \ + + wget -v --no-check-certificate https://ascend-repo.obs.cn-east-2.myhuaweicloud.com/CANN/CANN%208.5.2/Ascend-cann-A3-ops_8.5.2_linux-aarch64.run && \ + bash Ascend-cann-A3-ops_8.5.2_linux-aarch64.run --quiet --install && \ + rm -f Ascend-cann-A3-ops_8.5.2_linux-aarch64.run && \ + + source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ + wget -v --no-check-certificate https://ascend-repo.obs.cn-east-2.myhuaweicloud.com/CANN/CANN%208.5.2/Ascend-cann-nnal_8.5.2_linux-aarch64.run && \ + bash Ascend-cann-nnal_8.5.2_linux-aarch64.run --quiet --install && \ + rm -f Ascend-cann-nnal_8.5.2_linux-aarch64.run ; \ + elif [ "$ARCH" = "x86_64" ]; then \ + wget -v --no-check-certificate https://ascend-repo.obs.cn-east-2.myhuaweicloud.com/CANN/CANN%208.5.2/Ascend-cann-toolkit_8.5.2_linux-x86_64.run && \ + bash Ascend-cann-toolkit_8.5.2_linux-x86_64.run --quiet --install && \ + rm -f Ascend-cann-toolkit_8.5.2_linux-x86_64.run && \ + + wget -v --no-check-certificate https://ascend-repo.obs.cn-east-2.myhuaweicloud.com/CANN/CANN%208.5.2/Ascend-cann-A3-ops_8.5.2_linux-x86_64.run && \ + bash Ascend-cann-A3-ops_8.5.2_linux-x86_64.run --quiet --install && \ + rm -f Ascend-cann-A3-ops_8.5.2_linux-x86_64.run && \ + + source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ + wget -v --no-check-certificate https://ascend-repo.obs.cn-east-2.myhuaweicloud.com/CANN/CANN%208.5.2/Ascend-cann-nnal_8.5.2_linux-x86_64.run && \ + bash Ascend-cann-nnal_8.5.2_linux-x86_64.run --quiet --install && \ + rm -f Ascend-cann-nnal_8.5.2_linux-x86_64.run ; \ + fi && \ + + source /usr/local/Ascend/cann/set_env.sh && \ + rm -rf /tmp/* + +# Prepare required system dependencies +RUN apt-get update -y && \ + apt-get install -y --no-install-recommends gcc g++ cmake libnuma-dev wget git curl jq vim build-essential && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \ + pip install --upgrade pip packaging setuptools==80.10.2 && \ + pip cache purge + +# Prepare repositories with low update frequency +RUN ARCH=$(uname -m) && \ + # Set extra pip index for x86_64 platform + echo "[LOG INFO] Detected architecture: $ARCH" && \ + if [ "$ARCH" = "x86_64" ]; then \ + pip config set global.extra-index-url "https://download.pytorch.org/whl/cpu/"; \ + fi && \ + # Clone libs + git clone --depth 1 --branch v0.18.0 https://github.com/vllm-project/vllm.git && \ + git clone https://github.com/vllm-project/vllm-ascend.git && \ + cd vllm-ascend && git checkout 54879467c41784a446aa5b486a391d9bfbf488fa && cd .. && \ + git clone https://github.com/huggingface/transformers.git && \ + cd transformers && git checkout cc7ab9be508ce6ed3637bba9e50367b29b742dc6 && cd .. + +# 避免vllm-ascend 安装失败 +ENV MAX_JOBS=1 +ENV MAKEFLAGS="-j1" +ENV CMAKE_BUILD_PARALLEL_LEVEL=1 + +# Install repositories with low update frequency +RUN ARCH=$(uname -m) && \ + # Export and source env + if [ "$ARCH" = "aarch64" ]; then \ + export LD_LIBRARY_PATH=/usr/local/Ascend/cann-8.5.2/aarch64-linux/devlib/linux/aarch64:$LD_LIBRARY_PATH; \ + elif [ "$ARCH" = "x86_64" ]; then \ + export LD_LIBRARY_PATH=/usr/local/Ascend/cann-8.5.2/x86_64-linux/devlib/linux/x86_64/:$LD_LIBRARY_PATH; \ + fi && \ + source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ + source /usr/local/Ascend/nnal/atb/set_env.sh && \ + + # Install vllm + cd vllm && pip install -r requirements/build.txt && \ + VLLM_TARGET_DEVICE=empty pip install -v -e. && cd .. && \ + # Install vllm-ascend + cd vllm-ascend && pip install -r requirements.txt && \ + export COMPILE_CUSTOM_KERNELS=1 && pip install -v -e . --no-build-isolation && cd .. && \ + # Install transformers + cd transformers && pip install -e . && cd .. && \ + # Install torch & torch_npu & torchvision + pip install torch==2.9.0 torch_npu==2.9.0 torchvision==0.24.0 accelerate==1.13.0 mathruler && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + + +# Prepare and install verl (update frequently) +RUN git clone https://github.com/verl-project/verl.git && \ + cd verl && git checkout 4045d67063052dcb800c918c107b8d5a87046006 && pip install -r requirements-npu.txt && pip install -v -e . && cd .. && \ + # Clear extra files + rm -rf /tmp/* /var/tmp/* && \ + pip cache purge + +# Show install results +RUN pip list + +# Setting Default Commands +CMD ["/bin/bash"] diff --git a/verl/docker/aws/Dockerfile.extention.awsefa b/verl/docker/aws/Dockerfile.extention.awsefa new file mode 100644 index 0000000000000000000000000000000000000000..e221e9f1e60bd36869733773e6b26e46924fc2d6 --- /dev/null +++ b/verl/docker/aws/Dockerfile.extention.awsefa @@ -0,0 +1,55 @@ +# Base Image support aws EFA +# Build Image with frameworks based on this +FROM verlai/verl:app-verl0.6-transformers4.56.1-sglang0.5.2-mcore0.13.0-te2.2 + +# For aws instances with EFA net interface (Sagemaker AI Pod) +# install EFA driver: +######## AWS EFA ############ +ENV NCCL_VERSION=2.25.1-1 +ENV DEBIAN_FRONTEND=noninteractive +ENV EFA_INSTALLER_VERSION=1.40.0 +ENV AWS_OFI_NCCL_VERSION=1.14.2 +ENV FI_EFA_SET_CUDA_SYNC_MEMOPS=0 +ENV FI_PROVIDER=efa + +RUN apt update && apt install -y linux-image-generic libhwloc-dev + +RUN cd /tmp && \ + curl -O https://efa-installer.amazonaws.com/aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz && \ + tar -xf aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz && \ + cd aws-efa-installer && \ + ./efa_installer.sh -y -g --skip-kmod --skip-limit-conf --no-verify && \ + ldconfig && \ + rm -rf /tmp/aws-efa-installer /var/lib/apt/lists/* + +# NCCL EFA Plugin +RUN cd /tmp && \ + curl -LO https://github.com/aws/aws-ofi-nccl/archive/refs/tags/v${AWS_OFI_NCCL_VERSION}.tar.gz && \ + tar -xzf /tmp/v${AWS_OFI_NCCL_VERSION}.tar.gz && \ + rm /tmp/v${AWS_OFI_NCCL_VERSION}.tar.gz && \ + mv aws-ofi-nccl-${AWS_OFI_NCCL_VERSION} aws-ofi-nccl && \ + cd /tmp/aws-ofi-nccl && \ + ./autogen.sh && \ + ./configure --prefix=/opt/amazon/efa \ + --with-libfabric=/opt/amazon/efa \ + --with-cuda=/usr/local/cuda \ + --enable-platform-aws \ + --with-mpi=/opt/amazon/openmpi && \ + make -j$(nproc) install && \ + rm -rf /tmp/aws-ofi/nccl + +# NCCL +RUN echo "/usr/local/lib" >> /etc/ld.so.conf.d/local.conf && \ + echo "/opt/amazon/openmpi/lib" >> /etc/ld.so.conf.d/efa.conf && \ + ldconfig + +ENV OMPI_MCA_pml=^cm,ucx \ + OMPI_MCA_btl=tcp,self \ + OMPI_MCA_btl_tcp_if_exclude=lo,docker0,veth_def_agent \ + OPAL_PREFIX=/opt/amazon/openmpi \ + NCCL_SOCKET_IFNAME=^docker,lo,veth_def_agent \ + FI_EFA_USE_HUGE_PAGE=0 + +# docker build -t verl:awsefa --label "commit=$(git rev-parse --short HEAD)" . +# on aws: +# docker run --ipc=host --privileged --name verldev --gpus all --network=host --shm-size=1800gb -itd verl:awsefa diff --git a/verl/docker/aws/Dockerfile.ngc.vllm0.8.sagemaker b/verl/docker/aws/Dockerfile.ngc.vllm0.8.sagemaker new file mode 100644 index 0000000000000000000000000000000000000000..2746fa9f8e115c0c6ad43627b2d4200df75d597b --- /dev/null +++ b/verl/docker/aws/Dockerfile.ngc.vllm0.8.sagemaker @@ -0,0 +1,46 @@ +# Using a pre-built image from AWS DLC which contains the current version of python (3.10) and supported cuda version (12.1) +FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/huggingface-pytorch-training:2.1.0-transformers4.36.0-gpu-py310-cu121-ubuntu20.04 + +# uninstall nv-pytorch fork +RUN pip3 uninstall -y pytorch-quantization \ + pytorch-triton torch torch-tensorrt torchvision \ + xgboost transformer_engine flash_attn apex megatron-core + +# Define environments +ENV MAX_JOBS=32 +ENV VLLM_WORKER_MULTIPROC_METHOD=spawn +ENV DEBIAN_FRONTEND=noninteractive +ENV NODE_OPTIONS="" +ENV HF_HUB_ENABLE_HF_TRANSFER="1" + +# Install systemctl +RUN apt-get update && \ + apt-get install -y -o Dpkg::Options::="--force-confdef" systemd && \ + apt-get clean + +# Install tini +RUN apt-get update && \ + apt-get install -y tini && \ + apt-get clean + +# Install torch-2.6.0 + vllm-0.8.2 +RUN pip install --no-cache-dir vllm==0.8.2 torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 tensordict torchdata==0.11.0 \ + transformers>=4.49.0 accelerate datasets peft hf-transfer \ + ray[default] codetiming hydra-core pandas pyarrow>=15.0.0 pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler \ + pytest pre-commit py-spy pyext ruff tensorboard + +# Install flash_attn-2.7.4.post1 +RUN pip uninstall -y transformer-engine flash-attn && \ + pip install flash-attn==2.7.4.post1 --no-build-isolation + +# Fix cv2 +RUN pip uninstall -y pynvml nvidia-ml-py && \ + pip install --no-cache-dir nvidia-ml-py>=12.560.30 opencv-python-headless==4.8.0.74 fastapi==0.115.6 && \ + pip install --no-cache-dir --upgrade optree>=0.13.0 + +# Install verl +RUN pip install --no-cache-dir verl[vllm] -U + +# Reset pip config +RUN pip config unset global.index-url && \ + pip config unset global.extra-index-url diff --git a/verl/docker/rocm/Apptainerfile.rocm b/verl/docker/rocm/Apptainerfile.rocm new file mode 100644 index 0000000000000000000000000000000000000000..a040e89b237d6a5a0209f5bae809d862e871e57c --- /dev/null +++ b/verl/docker/rocm/Apptainerfile.rocm @@ -0,0 +1,57 @@ +Bootstrap: docker + +# Support - Traing: fsdp; Inference: vllm +# FROM: rocm/vllm:rocm6.2_mi300_ubuntu20.04_py3.9_vllm_0.6.4 +# Support - Traing: fsdp; Inference: vllm, sglang +FROM lmsysorg/sglang:v0.4.5-rocm630 + +%environment + export PYTORCH_ROCM_ARCH="gfx90a;gfx942" + + export HIPCC_COMPILE_FLAGS_APPEND="--amdgpu-target=gfx90a;gfx942 -D__HIP_PLATFORM_AMD__" + export CFLAGS="-D__HIP_PLATFORM_AMD__" + export CXXFLAGS="-D__HIP_PLATFORM_AMD__" + +%post + # Create source directory + mkdir -p /opt/src + + # Uninstall and reinstall vllm + pip uninstall -y vllm + cd /opt/src + git clone -b v0.6.3 https://github.com/vllm-project/vllm.git + cd vllm + MAX_JOBS=$(nproc) python3 setup.py install + cd /opt + rm -rf /opt/src/vllm + + # Install dependencies + pip install "tensordict<0.6" --no-deps + pip install accelerate \ + codetiming \ + datasets \ + dill \ + hydra-core \ + liger-kernel \ + numpy \ + pandas \ + peft \ + "pyarrow>=15.0.0" \ + pylatexenc \ + "ray[data,train,tune,serve]" \ + torchdata \ + transformers \ + wandb \ + orjson \ + pybind11 + + # Clone and install verl from GitHub + cd /opt + git clone https://github.com/verl-project/verl.git + cd verl + # Uncomment to use a specific version + # git checkout v0.3.0.post0 + pip install -e . --no-deps + + # Install torch_memory_saver + pip install git+https://github.com/ExtremeViscent/torch_memory_saver.git --no-deps \ No newline at end of file diff --git a/verl/docker/rocm/Dockerfile.rocm b/verl/docker/rocm/Dockerfile.rocm new file mode 100644 index 0000000000000000000000000000000000000000..83b9f6a283669ecdfe44fa379c687fbf855311c2 --- /dev/null +++ b/verl/docker/rocm/Dockerfile.rocm @@ -0,0 +1,172 @@ +FROM ubuntu:22.04 AS ubuntu + +# +# Install basic packages from OS distro +# +FROM ubuntu AS base + +ENV DEBIAN_FRONTEND=noninteractive +ARG PYTHON_VERSION=3.12 + +RUN --mount=target=/var/lib/apt/lists,type=cache,sharing=locked \ + --mount=target=/var/cache/apt,type=cache,sharing=locked \ + apt update && \ + apt install -y git software-properties-common curl rsync dialog gfortran wget sqlite3 + +RUN --mount=target=/var/lib/apt/lists,type=cache,sharing=locked \ + --mount=target=/var/cache/apt,type=cache,sharing=locked \ + if ! python3 --version | grep -q ${PYTHON_VERSION} ; then \ + add-apt-repository -y ppa:deadsnakes/ppa && apt update && \ + apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \ + python${PYTHON_VERSION}-lib2to3 python-is-python3 ; fi + +RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 && \ + update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} && \ + ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config && \ + curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} + + + +RUN wget -nv -O /tmp/cmake-3.26.4-linux-x86_64.tar.gz https://cmake.org/files/v3.26/cmake-3.26.4-linux-x86_64.tar.gz && \ + tar zfx /tmp/cmake-3.26.4-linux-x86_64.tar.gz -C /opt/ && \ + mv /opt/cmake-3.26.4-linux-x86_64 /opt/cmake-3.26.4 && \ + rm -f /tmp/cmake-3.26.4-linux-x86_64.tar.gz + +ENV PATH=/opt/cmake-3.26.4/bin:$PATH + +# +# Install ROCm rpm packages +# +FROM base AS rocm_deb + +ARG ROCM_VERSION=7.0.2 +ARG AMDGPU_VERSION=7.0.2 +ARG GFX_ARCH=gfx942|gfx950 + +RUN curl -sL https://repo.radeon.com/rocm/rocm.gpg.key | apt-key add - \ + && printf "deb [arch=amd64] https://repo.radeon.com/rocm/apt/$ROCM_VERSION/ jammy main\n" | tee /etc/apt/sources.list.d/rocm.list \ + && printf "deb [arch=amd64] https://repo.radeon.com/amdgpu/$AMDGPU_VERSION/ubuntu jammy main\n" | tee /etc/apt/sources.list.d/amdgpu.list \ + && printf "Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600\n" | tee /etc/apt/preferences.d/rocm-pin-600 \ + && apt-get update \ + && DEBIAN_FRONTEND=noninteractive apt-get install -y rocm && \ + find /opt/rocm/lib -type f -name '*gfx*' | grep -Ev "${GFX_ARCH}" | xargs rm -f && \ + find /opt/rocm/lib/hipblaslt/library -type f -name '*gfx*' | grep -Ev "${GFX_ARCH}" | xargs rm -f && \ + find /opt/rocm/lib/rocblas/library -type f -name '*gfx*' | grep -Ev "${GFX_ARCH}" | xargs rm -f && \ + find /opt/rocm/share/miopen/db -type f -name '*gfx*' | grep -Ev "${GFX_ARCH}" | xargs rm -f && \ + +ENV ROCM_HOME=/opt/rocm +ENV CPLUS_INCLUDE_PATH=/opt/rocm/include: +ENV LD_LIBRARY_PATH=/opt/rocm/lib:$LD_LIBRARY_PATH +ENV PATH=/opt/rocm/bin:/opt/rocm/llvm/bin:$PATH + +# +# Install pytorch packages +# +FROM rocm_deb AS rocm_torch + + +RUN --mount=type=cache,target=/root/.cache/pip \ + pip install --upgrade pip "setuptools<80" wheel numpy einops ninja && \ + pip install /opt/rocm/share/amd_smi + +RUN --mount=type=cache,target=/root/.cache/pip \ + cd /tmp && \ + wget -nv https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0.2/torch-2.9.1.dev20251204+rocm7.0.2.lw.git351ff442-cp312-cp312-linux_x86_64.whl -O torch-2.9.1.dev20251204+rocm7.0.2.lw.git351ff442-cp312-cp312-linux_x86_64.whl && \ + wget -nv https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0.2/apex-1.9.0a0+rocm7.0.2.git07c3ee53-cp312-cp312-linux_x86_64.whl -O apex-1.9.0a0+rocm7.0.2.git07c3ee53-cp312-cp312-linux_x86_64.whl && \ + wget -nv https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0.2/torchaudio-2.9.0+rocm7.0.2.gite3c6ee2b-cp312-cp312-linux_x86_64.whl -O torchaudio-2.9.0+rocm7.0.2.gite3c6ee2b-cp312-cp312-linux_x86_64.whl && \ + wget -nv https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0.2/torchvision-0.24.0+rocm7.0.2.gitb919bd0c-cp312-cp312-linux_x86_64.whl -O torchvision-0.24.0+rocm7.0.2.gitb919bd0c-cp312-cp312-linux_x86_64.whl && \ + wget -nv https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0.2/triton-3.5.1+rocm7.0.2.gita272dfa8-cp312-cp312-linux_x86_64.whl -O triton-3.5.1+rocm7.0.2.gita272dfa8-cp312-cp312-linux_x86_64.whl && \ + pip3 install *.whl && \ + rm -f *.whl + +ARG GPU_ARCH="gfx942;gfx950" +ENV PYTORCH_ROCM_ARCH=${GPU_ARCH} + +WORKDIR /root + +FROM rocm_torch AS new_toolset + +RUN --mount=type=cache,target=/root/.cache/pip \ + pip install numpy einops packaging psutil ninja + +# +# Build Flash Attention wheel +# +FROM new_toolset AS fa_build + +ARG FA_REPO="https://github.com/ROCm/flash-attention" +ARG FA_TAG="83f9e450cd10e20701fb109db9c7703d376f282b" + +RUN git clone ${FA_REPO} \ + && cd flash-attention \ + && git checkout ${FA_TAG} \ + && git submodule init \ + && git submodule update + +RUN cd flash-attention \ + && GPU_ARCHS=${GPU_ARCH} BUILD_TARGET=rocm MAX_JOBS=$(nproc) python3 setup.py bdist_wheel \ + && mkdir /install && cp dist/*.whl /install + +# +# Install Flash Attention +# +FROM new_toolset AS install_fa + +RUN --mount=type=bind,from=fa_build,source=/install,target=/tmp/install \ + --mount=type=cache,target=/root/.cache/pip \ + pip install /tmp/install/*.whl + +# +# Install TE +# +FROM install_fa AS te + +ENV NVTE_USE_HIPBLASLT=1 +ENV NVTE_USE_ROCM=1 +ENV NVTE_FRAMEWORK=pytorch +ENV NVTE_ROCM_ARCH=${GPU_ARCH} +ENV NVTE_USE_CAST_TRANSPOSE_TRITON=0 +ENV NVTE_CK_USES_BWD_V3=1 +ENV NVTE_CK_V3_BF16_CVT=2 + +ARG TE_TAG="386bd316" +RUN pip install pybind11 && git clone https://github.com/ROCm/TransformerEngine.git && \ + cd TransformerEngine && git checkout ${TE_TAG} && git submodule update --init --recursive && \ + GPU_ARCHS=${GPU_ARCH} MAX_JOBS=256 python3 setup.py install && \ + cd .. && rm -rf TransformerEngine +# TE patch +RUN FILE=$(find /usr/local/lib/python3.12/dist-packages/ -path "*transformer_engine*/transformer_engine/pytorch/attention/dot_product_attention/utils.py") && \ + sed -i '121s/2\.8\.3/2.8.4/' "$FILE" +# +# Install vllm +# +FROM te AS install_vllm + +ARG VLLM_TAG="1ff9d3353" +RUN pip install setuptools_scm && \ + mkdir /workspace && cd /workspace && \ + ln -sf /opt/rocm/lib/libamdhip64.so /usr/lib/libamdhip64.so && \ + git clone https://github.com/vllm-project/vllm && \ + cd vllm && git checkout ${VLLM_TAG} && pip install -r requirements/rocm.txt && \ + MAX_JOBS=64 python3 setup.py develop --no-deps +RUN pip uninstall tilelang -y && pip install xgrammar==0.1.32 + +RUN apt-get update && apt install vim -y + +FROM install_vllm AS install_verl +RUN cd /workspace && git clone https://github.com/verl-project/verl.git && \ + cd verl && pip install . +RUN pip install cupy-rocm-7-0 + + +ENV MIOPEN_DEBUG_CONV_DIRECT=0 +ARG AITER_TAG="45c428e54" +RUN cd /workspace && git clone --recursive https://github.com/ROCm/aiter.git && cd aiter && git checkout ${AITER_TAG} && python3 setup.py develop + +# for qwen3.5 +RUN pip install -U git+https://github.com/ISEEKYAN/mbridge.git@v0.15.1 && \ + pip install transformers --upgrade && \ + pip install megatron-core==0.16.0 && \ + pip install mathruler && \ + pip install qwen_vl_utils && \ + pip install flash-linear-attention diff --git a/verl/docker/rocm/Dockerfile.rocm6 b/verl/docker/rocm/Dockerfile.rocm6 new file mode 100644 index 0000000000000000000000000000000000000000..b1345db540503dee1cbc54c678c41174d9b491c2 --- /dev/null +++ b/verl/docker/rocm/Dockerfile.rocm6 @@ -0,0 +1,322 @@ +# FROM "compute-artifactory.amd.com:5000/rocm-plus-docker/framework/compute-rocm-rel-6.4:94_ubuntu22.04_py3.10_pytorch_release-2.7_575e247" +# FROM "rlfoundation.azurecr.io/rocm6.3.4:vllm-0.8.5-numa-patch-ubuntu-22.04" +FROM "rlsys/rocm-6.3.4-patch:rocm6.3.4-numa-patch_ubuntu-22.04" + +SHELL ["/bin/bash", "-ceuxo", "pipefail"] + +ENV MAX_JOBS=512 + +ENV PATH="/usr/local/python3.12/bin:$PATH" +RUN ln -sf /usr/bin/python3.12 /usr/bin/python && \ + ln -sf /usr/bin/pip3.12 /usr/bin/pip + +############################################ +############################################ +RUN apt-get update +RUN apt-get install -y pkg-config liblzma-dev +############################################ +############################################ + + +########################################### +##########Install TransformerEngine######## +########################################### +WORKDIR /workspace/ +# transformer-engine install +# https://github.com/ROCm/TransformerEngine + +RUN rm -rf TransformerEngine +RUN git clone --recursive https://github.com/ROCm/TransformerEngine.git +WORKDIR /workspace/TransformerEngine +RUN git checkout 236178e5 +# git checkout bb061ade +# git checkout 864405c + +ENV NVTE_FRAMEWORK=pytorch +ENV NVTE_ROCM_ARCH=gfx942 +ENV NVTE_USE_HIPBLASLT=1 +ENV NVTE_USE_ROCM=1 + +# export CMAKE_PREFIX_PATH="/opt/rocm:/opt/rocm/hip:/usr/local:/usr:${CMAKE_PREFIX_PATH:-}" +ENV CMAKE_PREFIX_PATH="/opt/rocm:/opt/rocm/hip:/usr/local:/usr" + + +# ENV NVTE_BUILD_MAX_JOBS=$(MAX_JOBS) + +RUN MAX_JOBS=$(MAX_JOBS) pip install . -vvv + +WORKDIR /workspace/ +########################################### +########################################### +########################################### + + + + + +#################################################################################### +################Install vllm - sglang require vllm 0.6.7 dependency################# +#################################################################################### +#### Require vllm 0.6.7 - checkout 113274a0 +WORKDIR /workspace/ +RUN rm -rf vllm +RUN pip uninstall -y vllm +# Refer to here (down-grade vllm to 0.6.3): https://docs.vllm.ai/en/v0.6.3/getting_started/amd-installation.html +RUN git clone https://github.com/ROCm/vllm.git +# git clone https://github.com/vllm-project/vllm.git +WORKDIR /workspace/vllm +RUN git checkout 113274a0 +ENV PYTORCH_ROCM_ARCH="gfx90a;gfx942" +#ENV MAX_JOBS=512 +ENV MAX_JOBS=${MAX_JOBS} +RUN pip install "boto3>=1.26.0" +RUN pip install setuptools_scm +# will add src into py. You can delete the repo +RUN python3 setup.py install +WORKDIR /workspace/ +#################################################################################### +#################################################################################### +#################################################################################### + + + +########################################### +############For hack docker################ +########################################### +RUN pip install setuptools==75.8.0 +########################################### +########################################### +########################################### + + + +########################################### +############build sgalng################### +########################################### +# Set environment variables +ENV BASE_DIR=/sgl-workspace +ENV BUILD_TYPE=all +ENV SGL_REPO=https://github.com/sgl-project/sglang +ENV SGL_BRANCH=v0.4.6.post5 +ENV TRITON_REPO=https://github.com/ROCm/triton.git +ENV TRITON_COMMIT=improve_fa_decode_3.0.0 +ENV AITER_REPO=https://github.com/ROCm/aiter.git +ENV AITER_COMMIT=v0.1.2 +# v0.1.2 version - commit id: 9d11f47 +# ENV AITER_COMMIT=9d11f47 + +ENV HIP_FORCE_DEV_KERNARG=1 +ENV HSA_NO_SCRATCH_RECLAIM=1 +ENV SGLANG_SET_CPU_AFFINITY=1 +ENV SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1 +ENV NCCL_MIN_NCHANNELS=112 +ENV MOE_PADDING=1 +ENV VLLM_FP8_PADDING=1 +ENV VLLM_FP8_ACT_PADDING=1 +ENV VLLM_FP8_WEIGHT_PADDING=1 +ENV VLLM_FP8_REDUCE_CONV=1 +ENV TORCHINDUCTOR_MAX_AUTOTUNE=1 +ENV TORCHINDUCTOR_MAX_AUTOTUNE_POINTWISE=1 +ENV HIPCC_COMPILE_FLAGS_APPEND="--offload-arch=gfx942" +ENV AMDGPU_TARGETS=gfx942 +ENV ROCM_ARCH=gfx942 +ENV PYTORCH_ROCM_ARCH="gfx90a;gfx942" + +# Switch to working directory +WORKDIR /sgl-workspace + +# Clean and create directory +RUN rm -rf /sgl-workspace && mkdir -p /sgl-workspace + +# Clone and build sglang +RUN git clone ${SGL_REPO} \ + && cd sglang \ + && git checkout ${SGL_BRANCH} || echo "Using default branch" \ + && cd sgl-kernel \ + && rm -f pyproject.toml \ + && mv pyproject_rocm.toml pyproject.toml \ + && python setup_rocm.py install \ + && cd .. \ + && if [ "$BUILD_TYPE" = "srt" ]; then \ + python -m pip --no-cache-dir install -e "python[srt_hip]"; \ + else \ + python -m pip --no-cache-dir install -e "python[all_hip]"; \ + fi \ + && cd /sgl-workspace \ + && cp -r /sgl-workspace/sglang /sglang \ + && python -m pip cache purge + +# Install common Python packages +RUN pip install IPython orjson python-multipart torchao pybind11 + +# Rebuild Triton +RUN pip uninstall -y triton || true \ + && git clone ${TRITON_REPO} \ + && cd triton \ + && git checkout ${TRITON_COMMIT} \ + && cd python \ + && python3 setup.py install \ + && cd /sgl-workspace + + +# ENV HIPCC_COMPILE_FLAGS_APPEND="--offload-arch=gfx942 --amdgpu-lower-module-lds-strategy=1" +# ENV HIPCC_COMPILE_FLAGS_APPEND="--offload-arch=gfx942" + +# Build aiter +#version: Commit 9d11f47 + # && git checkout ${AITER_COMMIT} \ +RUN pip uninstall -y aiter || true +RUN git clone ${AITER_REPO} \ + && cd aiter \ + && git checkout ${AITER_COMMIT} \ + && git submodule sync \ + && git submodule update --init --recursive \ + && PREBUILD_KERNELS=1 GPU_ARCHS=gfx942 python3 setup.py install \ + && cd /sgl-workspace + # && PREBUILD_KERNELS=1 GPU_ARCHS=gfx942 python3 setup.py develop \ + # && PREBUILD_KERNELS=1 GPU_ARCHS=gfx942 python3 setup.py develop \ + +# Copy MI300X config +RUN find /sgl-workspace/sglang/python/sglang/srt/layers/quantization/configs/ \ + /sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs/ \ + -type f -name '*MI300X*' | \ + xargs -I {} sh -c 'vf_config=$(echo "$1" | sed "s/MI300X/MI300X_VF/"); cp "$1" "$vf_config"' -- {} + +# Environment setup complete. +RUN echo "Environment setup complete." + +WORKDIR /workspace/ +########################################### +########################################### +########################################### + + + + + + +########################################### +###############vllm v0.8.5################# +########################################### +# ENV GITHUB_USERNAME=yushengsu-thu +# ENV GITHUB_MAIL=yushengsu@gmail.com + +# RUN git config --global user.name "${GITHUB_USERNAME}" \ +# && git config --global user.email "${GITHUB_MAIL}" + +WORKDIR /workspace/ + +ENV VLLM_TARGET_DEVICE=rocm +ENV ROCM_PATH=/opt/rocm +ENV SETUPTOOLS_SCM_PRETEND_VERSION=0.8.5.dev + +# Find the repo path in: DockerFile/Dockerfile.rocm_yang +# RUN git clone https://github.com/RLFoundation/vllm-patch.git +RUN pip uninstall -y vllm || true +RUN rm -rf vllm-patch +RUN git clone https://github.com/RLFoundation/vllm-patch.git \ + && cd vllm-patch \ + && git checkout v0.8.5-sleep-numa \ + && rm -rf build/ dist/ *.egg-info \ + && ln -sf /opt/rocm/lib/libamdhip64.so /usr/lib/libamdhip64.so \ + && SETUPTOOLS_SCM_PRETEND_VERSION=0.8.5.dev PYTORCH_ROCM_ARCH="gfx90a;gfx942" MAX_JOBS=${MAX_JOBS} python3 setup.py install + # RUN SETUPTOOLS_SCM_PRETEND_VERSION=0.8.5.dev PYTORCH_ROCM_ARCH="gfx90a;gfx942" MAX_JOBS=${MAX_JOBS} python3 setup.py develop + +WORKDIR /workspace/ +########################################### +########################################### +########################################### + + + + +######################################### +#### Install megatron-core############### +######################################### +RUN pip uninstall -y megatron-core && \ + git clone https://github.com/yushengsu-thu/Megatron-LM-amd_version.git && \ + cd Megatron-LM-amd_version && \ + pip install -vvv -e . && \ + cd /workspace/ +######################################### +######################################### +######################################### + + + + +####################################### +################apex################### +####################################### +WORKDIR /workspace/ +RUN pip uninstall -y apex && \ + git clone https://github.com/ROCm/apex.git && \ + cd apex && \ + python setup.py install && \ + cd /workspace/ +####################################### +####################################### +####################################### + + + + +################################################################################ +###########################Add torch_memory_saver############################### +################################################################################ +# Set environment variables +ENV HIPCC_COMPILE_FLAGS_APPEND="--amdgpu-target=gfx90a;gfx942 -D__HIP_PLATFORM_AMD__" +ENV CFLAGS="-D__HIP_PLATFORM_AMD__" +ENV CXXFLAGS="-D__HIP_PLATFORM_AMD__" +RUN pip install "git+https://github.com/YangWang92/torch_memory_saver_numa.git@numa" +################################################################################ +################################################################################ +################################################################################ + + + +######################################## +######Install ray####################### +######################################## +# need to add this patch: https://github.com/ray-project/ray/pull/53531/files +RUN pip uninstall ray -y +RUN pip install "ray[data,train,tune,serve]>=2.47.0" +######################################## +######################################## +######################################## + + + +########################################## +#######Install other dependencies######### +########################################## +RUN pip install "tensordict==0.6.2" --no-deps && \ + pip install accelerate \ + codetiming \ + datasets \ + dill \ + hydra-core \ + liger-kernel \ + numpy \ + pandas \ + peft \ + "pyarrow>=15.0.0" \ + pylatexenc \ + torchdata \ + wandb \ + orjson \ + pybind11 + +WORKDIR /workspace/ +RUN git clone https://github.com/verl-project/verl.git && \ + cd verl && \ + pip install -e . +########################################## +########################################## +########################################## + + + +WORKDIR /workspace/ + +CMD ["/usr/bin/bash"] diff --git a/verl/docker/rocm/Dockerfile.rocm7 b/verl/docker/rocm/Dockerfile.rocm7 new file mode 100644 index 0000000000000000000000000000000000000000..9b2148be74ca5210f55bbfcee27f406b71bf7784 --- /dev/null +++ b/verl/docker/rocm/Dockerfile.rocm7 @@ -0,0 +1,151 @@ +# default base image +ARG REMOTE_VLLM="1" +ARG COMMON_WORKDIR=/app +ARG BASE_IMAGE=rocm/vllm-dev:base + +FROM ${BASE_IMAGE} AS base + +ARG ARG_PYTORCH_ROCM_ARCH +ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} + +# Install some basic utilities +RUN apt-get update -q -y && apt-get install -q -y \ + sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \ + apt-transport-https ca-certificates wget curl +# Remove sccache +RUN python3 -m pip install --upgrade pip +RUN apt-get purge -y sccache; python3 -m pip uninstall -y sccache; rm -f "$(which sccache)" +ARG COMMON_WORKDIR +WORKDIR ${COMMON_WORKDIR} + + +# ----------------------- +# vLLM fetch stages +FROM base AS fetch_vllm_0 +ONBUILD COPY ./ vllm/ +FROM base AS fetch_vllm_1 +#ARG VLLM_REPO="https://github.com/ROCm/vllm.git" +#ARG VLLM_BRANCH="main" +ARG VLLM_REPO=https://github.com/HollowMan6/vllm.git +ARG VLLM_BRANCH="sleep_amd" +ONBUILD RUN git clone ${VLLM_REPO} \ + && cd vllm \ + && git checkout ${VLLM_BRANCH} +FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm + +# ----------------------- +# vLLM build stages +FROM fetch_vllm AS build_vllm +# Build vLLM +RUN cd vllm \ + && python3 -m pip install -r requirements/rocm.txt \ + && python3 setup.py clean --all \ + && ln -sf /opt/rocm/lib/libamdhip64.so /usr/lib/libamdhip64.so \ + && VLLM_TARGET_DEVICE=rocm ROCM_PATH=/opt/rocm/ VLLM_GPU_LANG=HIP SETUPTOOLS_SCM_PRETEND_VERSION=0.11.0.dev python3 setup.py bdist_wheel --dist-dir=dist + #&& python3 setup.py bdist_wheel --dist-dir=dist +FROM scratch AS export_vllm +ARG COMMON_WORKDIR +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/dist/*.whl / +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/requirements /requirements +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tests /tests +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/examples /examples +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite + +# ----------------------- +# Test vLLM image +FROM base AS test + +RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/* + +# Install vLLM +#RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ +COPY --from=export_vllm /*.whl /install +COPY --from=export_vllm /requirements /install/requirements +COPY --from=export_vllm /benchmarks /install/benchmarks +COPY --from=export_vllm /tests /install/tests +COPY --from=export_vllm /examples /install/examples +COPY --from=export_vllm /.buildkite /install/.buildkite + +RUN cd /install \ + && pip install -U -r requirements/rocm.txt \ + && pip install -U -r requirements/rocm-test.txt \ + && pip uninstall -y vllm \ + && pip install *.whl + +WORKDIR /vllm-workspace +ARG COMMON_WORKDIR +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace + +# install development dependencies (for testing) +RUN cd /vllm-workspace \ + && rm -rf vllm \ + && python3 -m pip install -e tests/vllm_test_utils \ + && python3 -m pip install lm-eval[api]==0.4.4 \ + && python3 -m pip install pytest-shard + +# ----------------------- +# Final vLLM image +FROM base AS final + +RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/* +# Error related to odd state for numpy 1.20.3 where there is no METADATA etc, but an extra LICENSES_bundled.txt. +# Manually remove it so that later steps of numpy upgrade can continue +RUN case "$(which python3)" in \ + *"/opt/conda/envs/py_3.9"*) \ + rm -rf /opt/conda/envs/py_3.9/lib/python3.9/site-packages/numpy-1.20.3.dist-info/;; \ + *) ;; esac + +RUN python3 -m pip install --upgrade huggingface-hub[cli] + +# Install vLLM +RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ + cd /install \ + && pip install -U -r requirements/rocm.txt \ + && pip uninstall -y vllm \ + && pip install *.whl + +ARG COMMON_WORKDIR + +# Copy over the benchmark scripts as well +COPY --from=export_vllm /benchmarks ${COMMON_WORKDIR}/vllm/benchmarks +COPY --from=export_vllm /examples ${COMMON_WORKDIR}/vllm/examples + +ENV RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES=1 +ENV TOKENIZERS_PARALLELISM=false + +# ENV that can improve safe tensor loading, and end-to-end time +ENV SAFETENSORS_FAST_GPU=1 + +# Performance environment variable. +ENV HIP_FORCE_DEV_KERNARG=1 + +# ----------------------- +# Install verl +ARG VERL_REPO=https://github.com/verl-project/verl.git +ARG VERL_BRANCH=main +RUN pip install "tensordict==0.6.2" --no-deps && \ + pip install accelerate \ + codetiming \ + datasets \ + dill \ + hydra-core \ + liger-kernel \ + numpy \ + pandas \ + peft \ + "pyarrow>=15.0.0" \ + pylatexenc \ + torchdata \ + wandb \ + orjson \ + pybind11 + +WORKDIR /workspace/ +RUN git clone ${VERL_REPO} && \ + cd verl && \ + git checkout ${VERL_BRANCH} && \ + pip install -e . + +CMD ["/bin/bash"] + diff --git a/verl/docker/rocm/Dockerfile.rocm_verl-0.3.0.post1 b/verl/docker/rocm/Dockerfile.rocm_verl-0.3.0.post1 new file mode 100644 index 0000000000000000000000000000000000000000..caef6b4f2e8e6317e2b18466927e05d17e94d750 --- /dev/null +++ b/verl/docker/rocm/Dockerfile.rocm_verl-0.3.0.post1 @@ -0,0 +1,58 @@ +# Build the docker in the repo dir: +# docker build -f docker/Dockerfile.rocm -t verl-rocm:03.04.2015 . +# docker images # you can find your built docker + + +# Support - Traing: fsdp; Inference: vllm +# FROM rocm/vllm:rocm6.2_mi300_ubuntu20.04_py3.9_vllm_0.6.4 +# Support - Traing: fsdp; Inference: vllm, sglang +FROM lmsysorg/sglang:v0.4.6.post5-rocm630 + +# Set working directory +# WORKDIR $PWD/app + +# Set environment variables +ENV PYTORCH_ROCM_ARCH="gfx90a;gfx942" + +ENV HIPCC_COMPILE_FLAGS_APPEND="--amdgpu-target=gfx90a;gfx942 -D__HIP_PLATFORM_AMD__" +ENV CFLAGS="-D__HIP_PLATFORM_AMD__" +ENV CXXFLAGS="-D__HIP_PLATFORM_AMD__" + +# Install vllm +RUN pip uninstall -y vllm && \ + rm -rf vllm && \ + git clone -b v0.6.3 https://github.com/vllm-project/vllm.git && \ + cd vllm && \ + MAX_JOBS=$(nproc) python3 setup.py install && \ + cd .. && \ + rm -rf vllm + +# Copy the entire project directory +COPY . . + +# Install dependencies +RUN pip install "tensordict==0.6.2" --no-deps && \ + pip install accelerate \ + codetiming \ + datasets \ + dill \ + hydra-core \ + liger-kernel \ + numpy \ + pandas \ + peft \ + "pyarrow>=15.0.0" \ + pylatexenc \ + "ray[data,train,tune,serve]<2.45.0" \ + torchdata \ + transformers \ + wandb \ + orjson \ + pybind11 + +RUN git clone https://github.com/verl-project/verl.git && \ + cd verl && \ + pip install -e . + +# Install torch_memory_saver +RUN pip install git+https://github.com/ExtremeViscent/torch_memory_saver.git --no-deps diff --git a/verl/docker/rocm/Dockerfile.rocm_verl-0.4.1 b/verl/docker/rocm/Dockerfile.rocm_verl-0.4.1 new file mode 100644 index 0000000000000000000000000000000000000000..afcab9686b248660d0a0721abec7f9a1ef8ebc7a --- /dev/null +++ b/verl/docker/rocm/Dockerfile.rocm_verl-0.4.1 @@ -0,0 +1,323 @@ +# FROM "compute-artifactory.amd.com:5000/rocm-plus-docker/framework/compute-rocm-rel-6.4:94_ubuntu22.04_py3.10_pytorch_release-2.7_575e247" +# FROM "rlfoundation.azurecr.io/rocm6.3.4:vllm-0.8.5-numa-patch-ubuntu-22.04" +FROM "rlsys/rocm-6.3.4-patch:rocm6.3.4-numa-patch_ubuntu-22.04" + +SHELL ["/bin/bash", "-ceuxo", "pipefail"] + +ENV MAX_JOBS=512 + +ENV PATH="/usr/local/python3.12/bin:$PATH" +RUN ln -sf /usr/bin/python3.12 /usr/bin/python && \ + ln -sf /usr/bin/pip3.12 /usr/bin/pip + +############################################ +############################################ +RUN apt-get update +RUN apt-get install -y pkg-config liblzma-dev +############################################ +############################################ + + +########################################### +##########Install TransformerEngine######## +########################################### +WORKDIR /workspace/ +# transformer-engine install +# https://github.com/ROCm/TransformerEngine + +RUN rm -rf TransformerEngine +RUN git clone --recursive https://github.com/ROCm/TransformerEngine.git +WORKDIR /workspace/TransformerEngine +RUN git checkout 236178e5 +# git checkout bb061ade +# git checkout 864405c + +ENV NVTE_FRAMEWORK=pytorch +ENV NVTE_ROCM_ARCH=gfx942 +ENV NVTE_USE_HIPBLASLT=1 +ENV NVTE_USE_ROCM=1 + +# export CMAKE_PREFIX_PATH="/opt/rocm:/opt/rocm/hip:/usr/local:/usr:${CMAKE_PREFIX_PATH:-}" +ENV CMAKE_PREFIX_PATH="/opt/rocm:/opt/rocm/hip:/usr/local:/usr" + + +# ENV NVTE_BUILD_MAX_JOBS=$(MAX_JOBS) + +RUN MAX_JOBS=$(MAX_JOBS) pip install . -vvv + +WORKDIR /workspace/ +########################################### +########################################### +########################################### + + + + + +#################################################################################### +################Install vllm - sglang require vllm 0.6.7 dependency################# +#################################################################################### +#### Require vllm 0.6.7 - checkout 113274a0 +WORKDIR /workspace/ +RUN rm -rf vllm +RUN pip uninstall -y vllm +# Refer to here (down-grade vllm to 0.6.3): https://docs.vllm.ai/en/v0.6.3/getting_started/amd-installation.html +RUN git clone https://github.com/ROCm/vllm.git +# git clone https://github.com/vllm-project/vllm.git +WORKDIR /workspace/vllm +RUN git checkout 113274a0 +ENV PYTORCH_ROCM_ARCH="gfx90a;gfx942" +#ENV MAX_JOBS=512 +ENV MAX_JOBS=${MAX_JOBS} +RUN pip install "boto3>=1.26.0" +RUN pip install setuptools_scm +# will add src into py. You can delete the repo +RUN python3 setup.py install +WORKDIR /workspace/ +#################################################################################### +#################################################################################### +#################################################################################### + + + +########################################### +############For hack docker################ +########################################### +RUN pip install setuptools==75.8.0 +########################################### +########################################### +########################################### + + + +########################################### +############build sgalng################### +########################################### +# Set environment variables +ENV BASE_DIR=/sgl-workspace +ENV BUILD_TYPE=all +ENV SGL_REPO=https://github.com/sgl-project/sglang +ENV SGL_BRANCH=v0.4.6.post5 +ENV TRITON_REPO=https://github.com/ROCm/triton.git +ENV TRITON_COMMIT=improve_fa_decode_3.0.0 +ENV AITER_REPO=https://github.com/ROCm/aiter.git +ENV AITER_COMMIT=v0.1.2 +# v0.1.2 version - commit id: 9d11f47 +# ENV AITER_COMMIT=9d11f47 + +ENV HIP_FORCE_DEV_KERNARG=1 +ENV HSA_NO_SCRATCH_RECLAIM=1 +ENV SGLANG_SET_CPU_AFFINITY=1 +ENV SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1 +ENV NCCL_MIN_NCHANNELS=112 +ENV MOE_PADDING=1 +ENV VLLM_FP8_PADDING=1 +ENV VLLM_FP8_ACT_PADDING=1 +ENV VLLM_FP8_WEIGHT_PADDING=1 +ENV VLLM_FP8_REDUCE_CONV=1 +ENV TORCHINDUCTOR_MAX_AUTOTUNE=1 +ENV TORCHINDUCTOR_MAX_AUTOTUNE_POINTWISE=1 +ENV HIPCC_COMPILE_FLAGS_APPEND="--offload-arch=gfx942" +ENV AMDGPU_TARGETS=gfx942 +ENV ROCM_ARCH=gfx942 +ENV PYTORCH_ROCM_ARCH="gfx90a;gfx942" + +# Switch to working directory +WORKDIR /sgl-workspace + +# Clean and create directory +RUN rm -rf /sgl-workspace && mkdir -p /sgl-workspace + +# Clone and build sglang +RUN git clone ${SGL_REPO} \ + && cd sglang \ + && git checkout ${SGL_BRANCH} || echo "Using default branch" \ + && cd sgl-kernel \ + && rm -f pyproject.toml \ + && mv pyproject_rocm.toml pyproject.toml \ + && python setup_rocm.py install \ + && cd .. \ + && if [ "$BUILD_TYPE" = "srt" ]; then \ + python -m pip --no-cache-dir install -e "python[srt_hip]"; \ + else \ + python -m pip --no-cache-dir install -e "python[all_hip]"; \ + fi \ + && cd /sgl-workspace \ + && cp -r /sgl-workspace/sglang /sglang \ + && python -m pip cache purge + +# Install common Python packages +RUN pip install IPython orjson python-multipart torchao pybind11 + +# Rebuild Triton +RUN pip uninstall -y triton || true \ + && git clone ${TRITON_REPO} \ + && cd triton \ + && git checkout ${TRITON_COMMIT} \ + && cd python \ + && python3 setup.py install \ + && cd /sgl-workspace + + +# ENV HIPCC_COMPILE_FLAGS_APPEND="--offload-arch=gfx942 --amdgpu-lower-module-lds-strategy=1" +# ENV HIPCC_COMPILE_FLAGS_APPEND="--offload-arch=gfx942" + +# Build aiter +#version: Commit 9d11f47 + # && git checkout ${AITER_COMMIT} \ +RUN pip uninstall -y aiter || true +RUN git clone ${AITER_REPO} \ + && cd aiter \ + && git checkout ${AITER_COMMIT} \ + && git submodule sync \ + && git submodule update --init --recursive \ + && PREBUILD_KERNELS=1 GPU_ARCHS=gfx942 python3 setup.py install \ + && cd /sgl-workspace + # && PREBUILD_KERNELS=1 GPU_ARCHS=gfx942 python3 setup.py develop \ + # && PREBUILD_KERNELS=1 GPU_ARCHS=gfx942 python3 setup.py develop \ + +# Copy MI300X config +RUN find /sgl-workspace/sglang/python/sglang/srt/layers/quantization/configs/ \ + /sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs/ \ + -type f -name '*MI300X*' | \ + xargs -I {} sh -c 'vf_config=$(echo "$1" | sed "s/MI300X/MI300X_VF/"); cp "$1" "$vf_config"' -- {} + +# Environment setup complete. +RUN echo "Environment setup complete." + +WORKDIR /workspace/ +########################################### +########################################### +########################################### + + + + + + +########################################### +###############vllm v0.8.5################# +########################################### +# ENV GITHUB_USERNAME=yushengsu-thu +# ENV GITHUB_MAIL=yushengsu@gmail.com + +# RUN git config --global user.name "${GITHUB_USERNAME}" \ +# && git config --global user.email "${GITHUB_MAIL}" + +WORKDIR /workspace/ + +ENV VLLM_TARGET_DEVICE=rocm +ENV ROCM_PATH=/opt/rocm +ENV SETUPTOOLS_SCM_PRETEND_VERSION=0.8.5.dev + +# Find the repo path in: DockerFile/Dockerfile.rocm_yang +# RUN git clone https://github.com/RLFoundation/vllm-patch.git +RUN pip uninstall -y vllm || true +RUN rm -rf vllm-patch +RUN git clone https://github.com/RLFoundation/vllm-patch.git \ + && cd vllm-patch \ + && git checkout v0.8.5-sleep-numa \ + && rm -rf build/ dist/ *.egg-info \ + && ln -sf /opt/rocm/lib/libamdhip64.so /usr/lib/libamdhip64.so \ + && SETUPTOOLS_SCM_PRETEND_VERSION=0.8.5.dev PYTORCH_ROCM_ARCH="gfx90a;gfx942" MAX_JOBS=${MAX_JOBS} python3 setup.py install + # RUN SETUPTOOLS_SCM_PRETEND_VERSION=0.8.5.dev PYTORCH_ROCM_ARCH="gfx90a;gfx942" MAX_JOBS=${MAX_JOBS} python3 setup.py develop + +WORKDIR /workspace/ +########################################### +########################################### +########################################### + + + + +######################################### +#### Install megatron-core############### +######################################### +RUN pip uninstall -y megatron-core && \ + git clone https://github.com/yushengsu-thu/Megatron-LM-amd_version.git && \ + cd Megatron-LM-amd_version && \ + pip install -vvv -e . && \ + cd /workspace/ +######################################### +######################################### +######################################### + + + + +####################################### +################apex################### +####################################### +WORKDIR /workspace/ +RUN pip uninstall -y apex && \ + git clone https://github.com/ROCm/apex.git && \ + cd apex && \ + python setup.py install && \ + cd /workspace/ +####################################### +####################################### +####################################### + + + + +################################################################################ +###########################Add torch_memory_saver############################### +################################################################################ +# Set environment variables +ENV HIPCC_COMPILE_FLAGS_APPEND="--amdgpu-target=gfx90a;gfx942 -D__HIP_PLATFORM_AMD__" +ENV CFLAGS="-D__HIP_PLATFORM_AMD__" +ENV CXXFLAGS="-D__HIP_PLATFORM_AMD__" +RUN pip install "git+https://github.com/YangWang92/torch_memory_saver_numa.git@numa" +################################################################################ +################################################################################ +################################################################################ + + + +######################################## +######Install ray####################### +######################################## +# need to add this patch: https://github.com/ray-project/ray/pull/53531/files +RUN pip uninstall ray -y +RUN pip install "ray[data,train,tune,serve]>=2.47.0" +######################################## +######################################## +######################################## + + + +########################################## +#######Install other dependencies######### +########################################## +RUN pip install "tensordict==0.6.2" --no-deps && \ + pip install accelerate \ + codetiming \ + datasets \ + dill \ + hydra-core \ + liger-kernel \ + numpy \ + pandas \ + peft \ + "pyarrow>=15.0.0" \ + pylatexenc \ + torchdata \ + wandb \ + orjson \ + pybind11 + +WORKDIR /workspace/ +RUN git clone https://github.com/verl-project/verl.git && \ + cd verl && \ + pip install -e . +########################################## +########################################## +########################################## + + + +WORKDIR /workspace/ + +CMD ["/usr/bin/bash"] +CMD ["/usr/bin/bash"] diff --git a/verl/docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.sglang.vllm.mcore0.12 b/verl/docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.sglang.vllm.mcore0.12 new file mode 100644 index 0000000000000000000000000000000000000000..68c3e1a043b3955f3773ca1665de5d8d1dd4fc17 --- /dev/null +++ b/verl/docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.sglang.vllm.mcore0.12 @@ -0,0 +1,41 @@ +# Start from the verl base image +# Dockerfile.base +FROM verlai/verl:base-verl0.4-cu124-cudnn9.8-torch2.6-fa2.7.4 + +# Define environments +ENV MAX_JOBS=32 +ENV VLLM_WORKER_MULTIPROC_METHOD=spawn +ENV DEBIAN_FRONTEND=noninteractive +ENV NODE_OPTIONS="" +ENV PIP_ROOT_USER_ACTION=ignore +ENV HF_HUB_ENABLE_HF_TRANSFER="1" + +# Install sglang-0.4.6.post5 and torch-memory-saver +RUN pip install --resume-retries 999 "sglang[all]==0.4.6.post5" --no-cache-dir --find-links https://flashinfer.ai/whl/cu124/torch2.6/flashinfer-python && pip install torch-memory-saver --no-cache-dir + +# Some sglang operations in 0.4.6.post5 require vllm +# [Warning] vllm can have some packages not compatible with sglang, for example, flashinfer +RUN pip install --resume-retries 999 --no-cache-dir vllm==0.8.5.post1 + +# Fix packages +RUN pip install --no-cache-dir "tensordict==0.6.2" "transformers[hf_xet]>=4.51.0" accelerate datasets peft hf-transfer \ + "numpy<2.0.0" "pyarrow>=19.0.1" pandas \ + ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \ + pytest py-spy pyext pre-commit ruff + +RUN pip uninstall -y pynvml nvidia-ml-py && \ + pip install --resume-retries 999 --no-cache-dir --upgrade "nvidia-ml-py>=12.560.30" "fastapi[standard]>=0.115.0" "optree>=0.13.0" "pydantic>=2.9" "grpcio>=1.62.1" + +RUN pip install --resume-retries 999 --no-cache-dir nvidia-cudnn-cu12==9.8.0.87 + +# Install TransformerEngine +RUN export NVTE_FRAMEWORK=pytorch && pip3 install --resume-retries 999 --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@v2.2.1 + +# Install Megatron-LM +RUN pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.12.2 + +# Fix for transformers 4.53.0 +RUN pip3 install --no-cache-dir "transformers[hf_xet]<4.52.0" + +# Install mbridge +RUN pip3 install --no-cache-dir mbridge \ No newline at end of file diff --git a/verl/docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.sglang.vllm.mcore0.12.deepep b/verl/docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.sglang.vllm.mcore0.12.deepep new file mode 100644 index 0000000000000000000000000000000000000000..f08aeee9c7389f036e4b8ca032209495d9f5bffe --- /dev/null +++ b/verl/docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.sglang.vllm.mcore0.12.deepep @@ -0,0 +1,82 @@ +# Start from the verl base image +# Dockerfile.base +FROM verlai/verl:base-verl0.4-cu124-cudnn9.8-torch2.6-fa2.7.4 + +# Define environments +ENV MAX_JOBS=32 +ENV VLLM_WORKER_MULTIPROC_METHOD=spawn +ENV DEBIAN_FRONTEND=noninteractive +ENV NODE_OPTIONS="" +ENV PIP_ROOT_USER_ACTION=ignore +ENV HF_HUB_ENABLE_HF_TRANSFER="1" + +# Install sglang-0.4.6.post5 and torch-memory-saver +RUN pip install --resume-retries 999 "sglang[all]==0.4.6.post5" --no-cache-dir --find-links https://flashinfer.ai/whl/cu124/torch2.6/flashinfer-python && pip install torch-memory-saver --no-cache-dir + +# Some sglang operations in 0.4.6.post5 require vllm +# [Warning] vllm can have some packages not compatible with sglang, for example, flashinfer +RUN pip install --resume-retries 999 --no-cache-dir vllm==0.8.5.post1 + +# Fix packages +RUN pip install --no-cache-dir "tensordict==0.6.2" "transformers[hf_xet]>=4.51.0" accelerate datasets peft hf-transfer \ + "numpy<2.0.0" "pyarrow>=19.0.1" pandas \ + ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \ + pytest py-spy pyext pre-commit ruff + +RUN pip uninstall -y pynvml nvidia-ml-py && \ + pip install --resume-retries 999 --no-cache-dir --upgrade "nvidia-ml-py>=12.560.30" "fastapi[standard]>=0.115.0" "optree>=0.13.0" "pydantic>=2.9" "grpcio>=1.62.1" + +RUN pip install --resume-retries 999 --no-cache-dir nvidia-cudnn-cu12==9.8.0.87 + +# Install TransformerEngine +RUN export NVTE_FRAMEWORK=pytorch && pip3 install --resume-retries 999 --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@v2.2.1 + +# Install Megatron-LM +RUN pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.12.2 + +# Fix for transformers 4.53.0 +RUN pip3 install --no-cache-dir "transformers[hf_xet]<4.52.0" + +# Install mbridge +RUN pip3 install --no-cache-dir mbridge + +# Install DeepEP +## the dependency of IBGDA +RUN ln -s /usr/lib/x86_64-linux-gnu/libmlx5.so.1 /usr/lib/x86_64-linux-gnu/libmlx5.so + +## Clone and build deepep and deepep-nvshmem +RUN git clone -b v2.3.1 https://github.com/NVIDIA/gdrcopy.git && \ + git clone https://github.com/deepseek-ai/DeepEP.git && \ + cd DeepEP && git checkout a84a248 + +# Prepare nvshmem +RUN wget https://developer.nvidia.com/downloads/assets/secure/nvshmem/nvshmem_src_3.2.5-1.txz && \ + tar -xvf nvshmem_src_3.2.5-1.txz && mv nvshmem_src deepep-nvshmem && \ + cd deepep-nvshmem && git apply ../DeepEP/third-party/nvshmem.patch + +ENV CUDA_HOME=/usr/local/cuda +### Set MPI environment variables. Having errors when not set. +ENV CPATH=/usr/local/mpi/include:$CPATH +ENV LD_LIBRARY_PATH=/usr/local/mpi/lib:$LD_LIBRARY_PATH +ENV LD_LIBRARY_PATH=/usr/local/x86_64-linux-gnu:$LD_LIBRARY_PATH +ENV GDRCOPY_HOME=/workspace/gdrcopy + +## Build deepep-nvshmem +RUN cd deepep-nvshmem && \ + NVSHMEM_SHMEM_SUPPORT=0 \ + NVSHMEM_UCX_SUPPORT=0 \ + NVSHMEM_USE_NCCL=0 \ + NVSHMEM_MPI_SUPPORT=0 \ + NVSHMEM_IBGDA_SUPPORT=1 \ + NVSHMEM_PMIX_SUPPORT=0 \ + NVSHMEM_TIMEOUT_DEVICE_POLLING=0 \ + NVSHMEM_USE_GDRCOPY=1 \ + cmake -G Ninja -S . -B build/ -DCMAKE_INSTALL_PREFIX=/workspace/deepep-nvshmem/install && cmake --build build/ --target install + +ENV NVSHMEM_DIR=/workspace/deepep-nvshmem/install +ENV LD_LIBRARY_PATH=$NVSHMEM_DIR/lib:$LD_LIBRARY_PATH +ENV PATH=$NVSHMEM_DIR/bin:$PATH + +## Build deepep +RUN cd DeepEP && \ + python setup.py install \ No newline at end of file diff --git a/verl/docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.sglang.vllm.mcore0.13.preview b/verl/docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.sglang.vllm.mcore0.13.preview new file mode 100644 index 0000000000000000000000000000000000000000..0e0bdd43fec1eac3e4a85de5416e34fad4ac7c69 --- /dev/null +++ b/verl/docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.sglang.vllm.mcore0.13.preview @@ -0,0 +1,82 @@ +# Start from the verl base image +# Dockerfile.base +FROM verlai/verl:base-verl0.4-cu124-cudnn9.8-torch2.6-fa2.7.4 + +# Define environments +ENV MAX_JOBS=32 +ENV VLLM_WORKER_MULTIPROC_METHOD=spawn +ENV DEBIAN_FRONTEND=noninteractive +ENV NODE_OPTIONS="" +ENV PIP_ROOT_USER_ACTION=ignore +ENV HF_HUB_ENABLE_HF_TRANSFER="1" + +# Install sglang-0.4.6.post5 and torch-memory-saver +RUN pip install --resume-retries 999 "sglang[all]==0.4.6.post5" --no-cache-dir --find-links https://flashinfer.ai/whl/cu124/torch2.6/flashinfer-python && pip install torch-memory-saver --no-cache-dir + +# Some sglang operations in 0.4.6.post5 require vllm +# [Warning] vllm can have some packages not compatible with sglang, for example, flashinfer +RUN pip install --resume-retries 999 --no-cache-dir vllm==0.8.5.post1 + +# Fix packages +RUN pip install --no-cache-dir "tensordict==0.6.2" "transformers[hf_xet]>=4.51.0" accelerate datasets peft hf-transfer \ + "numpy<2.0.0" "pyarrow>=19.0.1" pandas \ + ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \ + pytest py-spy pyext pre-commit ruff + +RUN pip uninstall -y pynvml nvidia-ml-py && \ + pip install --resume-retries 999 --no-cache-dir --upgrade "nvidia-ml-py>=12.560.30" "fastapi[standard]>=0.115.0" "optree>=0.13.0" "pydantic>=2.9" "grpcio>=1.62.1" + +RUN pip install --resume-retries 999 --no-cache-dir nvidia-cudnn-cu12==9.8.0.87 + +# Install TransformerEngine +RUN export NVTE_FRAMEWORK=pytorch && pip3 install --resume-retries 999 --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.5 + +# Install Megatron-LM +RUN pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/Megatron-LM.git@core_r0.13.0 + +# Fix for transformers 4.53.0 +RUN pip3 install --no-cache-dir "transformers[hf_xet]<4.52.0" + +# Install mbridge +RUN pip3 install --no-cache-dir mbridge + +# Install DeepEP +## the dependency of IBGDA +RUN ln -s /usr/lib/x86_64-linux-gnu/libmlx5.so.1 /usr/lib/x86_64-linux-gnu/libmlx5.so + +## Clone and build deepep and deepep-nvshmem +RUN git clone -b v2.3.1 https://github.com/NVIDIA/gdrcopy.git && \ + git clone https://github.com/deepseek-ai/DeepEP.git && \ + cd DeepEP && git checkout a84a248 + +# Prepare nvshmem +RUN wget https://developer.nvidia.com/downloads/assets/secure/nvshmem/nvshmem_src_3.2.5-1.txz && \ + tar -xvf nvshmem_src_3.2.5-1.txz && mv nvshmem_src deepep-nvshmem && \ + cd deepep-nvshmem && git apply ../DeepEP/third-party/nvshmem.patch + +ENV CUDA_HOME=/usr/local/cuda +### Set MPI environment variables. Having errors when not set. +ENV CPATH=/usr/local/mpi/include:$CPATH +ENV LD_LIBRARY_PATH=/usr/local/mpi/lib:$LD_LIBRARY_PATH +ENV LD_LIBRARY_PATH=/usr/local/x86_64-linux-gnu:$LD_LIBRARY_PATH +ENV GDRCOPY_HOME=/workspace/gdrcopy + +## Build deepep-nvshmem +RUN cd deepep-nvshmem && \ + NVSHMEM_SHMEM_SUPPORT=0 \ + NVSHMEM_UCX_SUPPORT=0 \ + NVSHMEM_USE_NCCL=0 \ + NVSHMEM_MPI_SUPPORT=0 \ + NVSHMEM_IBGDA_SUPPORT=1 \ + NVSHMEM_PMIX_SUPPORT=0 \ + NVSHMEM_TIMEOUT_DEVICE_POLLING=0 \ + NVSHMEM_USE_GDRCOPY=1 \ + cmake -G Ninja -S . -B build/ -DCMAKE_INSTALL_PREFIX=/workspace/deepep-nvshmem/install && cmake --build build/ --target install + +ENV NVSHMEM_DIR=/workspace/deepep-nvshmem/install +ENV LD_LIBRARY_PATH=$NVSHMEM_DIR/lib:$LD_LIBRARY_PATH +ENV PATH=$NVSHMEM_DIR/bin:$PATH + +## Build deepep +RUN cd DeepEP && \ + python setup.py install \ No newline at end of file diff --git a/verl/docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.vllm.mcore0.12 b/verl/docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.vllm.mcore0.12 new file mode 100644 index 0000000000000000000000000000000000000000..616c701c523a76eb11cfe8364a40ca7329023f73 --- /dev/null +++ b/verl/docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.vllm.mcore0.12 @@ -0,0 +1,47 @@ +# Start from the verl base image +# Dockerfile.base +FROM verlai/verl:base-verl0.4-cu124-cudnn9.8-torch2.6-fa2.7.4 + +# Define environments +ENV MAX_JOBS=32 +ENV VLLM_WORKER_MULTIPROC_METHOD=spawn +ENV DEBIAN_FRONTEND=noninteractive +ENV NODE_OPTIONS="" +ENV PIP_ROOT_USER_ACTION=ignore +ENV HF_HUB_ENABLE_HF_TRANSFER="1" + +# Install torch-2.6.0+cu124 + vllm-0.8.5.post1 +# torch-2.6.0+cu124: cxx11abi=False +# torch-2.6.0+cu126: cxx11abi=True +# see https://github.com/flashinfer-ai/flashinfer/issues/911 +RUN pip install --resume-retries 999 --no-cache-dir vllm==0.8.5.post1 + +# Install flashinfer-0.2.2.post1+cu126 (cxx11abi=True) +# vllm-0.8.3 does not support flashinfer>=0.2.3 +# see https://github.com/vllm-project/vllm/pull/15777 +RUN aria2c --max-tries=9999 https://github.com/flashinfer-ai/flashinfer/releases/download/v0.2.2.post1/flashinfer_python-0.2.2.post1+cu124torch2.6-cp38-abi3-linux_x86_64.whl && \ + pip install --no-cache-dir flashinfer_python-0.2.2.post1+cu124torch2.6-cp38-abi3-linux_x86_64.whl && \ + rm flashinfer_python-0.2.2.post1+cu124torch2.6-cp38-abi3-linux_x86_64.whl + +# Fix packages +RUN pip install --no-cache-dir "tensordict==0.6.2" "transformers[hf_xet]>=4.51.0" accelerate datasets peft hf-transfer \ + "numpy<2.0.0" "pyarrow>=19.0.1" pandas \ + ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \ + pytest py-spy pyext pre-commit ruff + +RUN pip uninstall -y pynvml nvidia-ml-py && \ + pip install --resume-retries 999 --no-cache-dir --upgrade "nvidia-ml-py>=12.560.30" "fastapi[standard]>=0.115.0" "optree>=0.13.0" "pydantic>=2.9" "grpcio>=1.62.1" + +RUN pip install --resume-retries 999 --no-cache-dir nvidia-cudnn-cu12==9.8.0.87 + +# Install TransformerEngine +RUN export NVTE_FRAMEWORK=pytorch && pip3 install --resume-retries 999 --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@v2.2.1 + +# Install Megatron-LM +RUN pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.12.2 + +# Fix for transformers 4.53.0 +RUN pip3 install --no-cache-dir "transformers[hf_xet]<4.52.0" + +# Install mbridge +RUN pip3 install --no-cache-dir mbridge \ No newline at end of file diff --git a/verl/docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.vllm.mcore0.12.deepep b/verl/docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.vllm.mcore0.12.deepep new file mode 100644 index 0000000000000000000000000000000000000000..a2be5e998dd32a2bcfbdcc63bedd2dd1ae634943 --- /dev/null +++ b/verl/docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.vllm.mcore0.12.deepep @@ -0,0 +1,88 @@ +# Start from the verl base image +# Dockerfile.base +FROM verlai/verl:base-verl0.4-cu124-cudnn9.8-torch2.6-fa2.7.4 + +# Define environments +ENV MAX_JOBS=32 +ENV VLLM_WORKER_MULTIPROC_METHOD=spawn +ENV DEBIAN_FRONTEND=noninteractive +ENV NODE_OPTIONS="" +ENV PIP_ROOT_USER_ACTION=ignore +ENV HF_HUB_ENABLE_HF_TRANSFER="1" + +# Install torch-2.6.0+cu124 + vllm-0.8.5.post1 +# torch-2.6.0+cu124: cxx11abi=False +# torch-2.6.0+cu126: cxx11abi=True +# see https://github.com/flashinfer-ai/flashinfer/issues/911 +RUN pip install --resume-retries 999 --no-cache-dir vllm==0.8.5.post1 + +# Install flashinfer-0.2.2.post1+cu126 (cxx11abi=True) +# vllm-0.8.3 does not support flashinfer>=0.2.3 +# see https://github.com/vllm-project/vllm/pull/15777 +RUN aria2c --max-tries=9999 https://github.com/flashinfer-ai/flashinfer/releases/download/v0.2.2.post1/flashinfer_python-0.2.2.post1+cu124torch2.6-cp38-abi3-linux_x86_64.whl && \ + pip install --no-cache-dir flashinfer_python-0.2.2.post1+cu124torch2.6-cp38-abi3-linux_x86_64.whl && \ + rm flashinfer_python-0.2.2.post1+cu124torch2.6-cp38-abi3-linux_x86_64.whl + +# Fix packages +RUN pip install --no-cache-dir "tensordict==0.6.2" "transformers[hf_xet]>=4.51.0" accelerate datasets peft hf-transfer \ + "numpy<2.0.0" "pyarrow>=19.0.1" pandas \ + ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \ + pytest py-spy pyext pre-commit ruff + +RUN pip uninstall -y pynvml nvidia-ml-py && \ + pip install --resume-retries 999 --no-cache-dir --upgrade "nvidia-ml-py>=12.560.30" "fastapi[standard]>=0.115.0" "optree>=0.13.0" "pydantic>=2.9" "grpcio>=1.62.1" + +RUN pip install --resume-retries 999 --no-cache-dir nvidia-cudnn-cu12==9.8.0.87 + +# Install TransformerEngine +RUN export NVTE_FRAMEWORK=pytorch && pip3 install --resume-retries 999 --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@v2.2.1 + +# Install Megatron-LM +RUN pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.12.2 + +# Fix for transformers 4.53.0 +RUN pip3 install --no-cache-dir "transformers[hf_xet]<4.52.0" + +# Install mbridge +RUN pip3 install --no-cache-dir mbridge + +# Install DeepEP +## the dependency of IBGDA +RUN ln -s /usr/lib/x86_64-linux-gnu/libmlx5.so.1 /usr/lib/x86_64-linux-gnu/libmlx5.so + +## Clone and build deepep and deepep-nvshmem +RUN git clone -b v2.3.1 https://github.com/NVIDIA/gdrcopy.git && \ + git clone https://github.com/deepseek-ai/DeepEP.git && \ + cd DeepEP && git checkout a84a248 + +# Prepare nvshmem +RUN wget https://developer.nvidia.com/downloads/assets/secure/nvshmem/nvshmem_src_3.2.5-1.txz && \ + tar -xvf nvshmem_src_3.2.5-1.txz && mv nvshmem_src deepep-nvshmem && \ + cd deepep-nvshmem && git apply ../DeepEP/third-party/nvshmem.patch + +ENV CUDA_HOME=/usr/local/cuda +### Set MPI environment variables. Having errors when not set. +ENV CPATH=/usr/local/mpi/include:$CPATH +ENV LD_LIBRARY_PATH=/usr/local/mpi/lib:$LD_LIBRARY_PATH +ENV LD_LIBRARY_PATH=/usr/local/x86_64-linux-gnu:$LD_LIBRARY_PATH +ENV GDRCOPY_HOME=/workspace/gdrcopy + +## Build deepep-nvshmem +RUN cd deepep-nvshmem && \ + NVSHMEM_SHMEM_SUPPORT=0 \ + NVSHMEM_UCX_SUPPORT=0 \ + NVSHMEM_USE_NCCL=0 \ + NVSHMEM_MPI_SUPPORT=0 \ + NVSHMEM_IBGDA_SUPPORT=1 \ + NVSHMEM_PMIX_SUPPORT=0 \ + NVSHMEM_TIMEOUT_DEVICE_POLLING=0 \ + NVSHMEM_USE_GDRCOPY=1 \ + cmake -G Ninja -S . -B build/ -DCMAKE_INSTALL_PREFIX=/workspace/deepep-nvshmem/install && cmake --build build/ --target install + +ENV NVSHMEM_DIR=/workspace/deepep-nvshmem/install +ENV LD_LIBRARY_PATH=$NVSHMEM_DIR/lib:$LD_LIBRARY_PATH +ENV PATH=$NVSHMEM_DIR/bin:$PATH + +## Build deepep +RUN cd DeepEP && \ + python setup.py install \ No newline at end of file diff --git a/verl/docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.vllm.mcore0.13.preview b/verl/docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.vllm.mcore0.13.preview new file mode 100644 index 0000000000000000000000000000000000000000..be6183a8479fe3e379f7ddaa06034c3e1ff64278 --- /dev/null +++ b/verl/docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.vllm.mcore0.13.preview @@ -0,0 +1,85 @@ +# Start from the verl base image +# Dockerfile.base +FROM verlai/verl:base-verl0.4-cu124-cudnn9.8-torch2.6-fa2.7.4 + +# Define environments +ENV MAX_JOBS=32 +ENV VLLM_WORKER_MULTIPROC_METHOD=spawn +ENV DEBIAN_FRONTEND=noninteractive +ENV NODE_OPTIONS="" +ENV PIP_ROOT_USER_ACTION=ignore +ENV HF_HUB_ENABLE_HF_TRANSFER="1" + +# Install torch-2.6.0+cu124 + vllm-0.8.5.post1 +# torch-2.6.0+cu124: cxx11abi=False +# torch-2.6.0+cu126: cxx11abi=True +# see https://github.com/flashinfer-ai/flashinfer/issues/911 +RUN pip install --resume-retries 999 --no-cache-dir vllm==0.8.5.post1 + +# Install flashinfer-0.2.2.post1+cu126 (cxx11abi=True) +# vllm-0.8.3 does not support flashinfer>=0.2.3 +# see https://github.com/vllm-project/vllm/pull/15777 +RUN aria2c --max-tries=9999 https://github.com/flashinfer-ai/flashinfer/releases/download/v0.2.2.post1/flashinfer_python-0.2.2.post1+cu124torch2.6-cp38-abi3-linux_x86_64.whl && \ + pip install --no-cache-dir flashinfer_python-0.2.2.post1+cu124torch2.6-cp38-abi3-linux_x86_64.whl && \ + rm flashinfer_python-0.2.2.post1+cu124torch2.6-cp38-abi3-linux_x86_64.whl + +# Fix packages +RUN pip install --no-cache-dir "tensordict==0.6.2" "transformers[hf_xet]>=4.51.0" accelerate datasets peft hf-transfer \ + "numpy<2.0.0" "pyarrow>=19.0.1" pandas \ + ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \ + pytest py-spy pyext pre-commit ruff + +RUN pip uninstall -y pynvml nvidia-ml-py && \ + pip install --resume-retries 999 --no-cache-dir --upgrade "nvidia-ml-py>=12.560.30" "fastapi[standard]>=0.115.0" "optree>=0.13.0" "pydantic>=2.9" "grpcio>=1.62.1" + +RUN pip install --resume-retries 999 --no-cache-dir nvidia-cudnn-cu12==9.8.0.87 + +# Install TransformerEngine +RUN export NVTE_FRAMEWORK=pytorch && pip3 install --resume-retries 999 --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.5 + +# Install Megatron-LM +RUN pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.12.2 + +# Install mbridge +RUN pip3 install --no-cache-dir mbridge + +# Install DeepEP +## the dependency of IBGDA +RUN ln -s /usr/lib/x86_64-linux-gnu/libmlx5.so.1 /usr/lib/x86_64-linux-gnu/libmlx5.so + +## Clone and build deepep and deepep-nvshmem +RUN git clone -b v2.3.1 https://github.com/NVIDIA/gdrcopy.git && \ + git clone https://github.com/deepseek-ai/DeepEP.git && \ + cd DeepEP && git checkout a84a248 + +# Prepare nvshmem +RUN wget https://developer.nvidia.com/downloads/assets/secure/nvshmem/nvshmem_src_3.2.5-1.txz && \ + tar -xvf nvshmem_src_3.2.5-1.txz && mv nvshmem_src deepep-nvshmem && \ + cd deepep-nvshmem && git apply ../DeepEP/third-party/nvshmem.patch + +ENV CUDA_HOME=/usr/local/cuda +### Set MPI environment variables. Having errors when not set. +ENV CPATH=/usr/local/mpi/include:$CPATH +ENV LD_LIBRARY_PATH=/usr/local/mpi/lib:$LD_LIBRARY_PATH +ENV LD_LIBRARY_PATH=/usr/local/x86_64-linux-gnu:$LD_LIBRARY_PATH +ENV GDRCOPY_HOME=/workspace/gdrcopy + +## Build deepep-nvshmem +RUN cd deepep-nvshmem && \ + NVSHMEM_SHMEM_SUPPORT=0 \ + NVSHMEM_UCX_SUPPORT=0 \ + NVSHMEM_USE_NCCL=0 \ + NVSHMEM_MPI_SUPPORT=0 \ + NVSHMEM_IBGDA_SUPPORT=1 \ + NVSHMEM_PMIX_SUPPORT=0 \ + NVSHMEM_TIMEOUT_DEVICE_POLLING=0 \ + NVSHMEM_USE_GDRCOPY=1 \ + cmake -G Ninja -S . -B build/ -DCMAKE_INSTALL_PREFIX=/workspace/deepep-nvshmem/install && cmake --build build/ --target install + +ENV NVSHMEM_DIR=/workspace/deepep-nvshmem/install +ENV LD_LIBRARY_PATH=$NVSHMEM_DIR/lib:$LD_LIBRARY_PATH +ENV PATH=$NVSHMEM_DIR/bin:$PATH + +## Build deepep +RUN cd DeepEP && \ + python setup.py install \ No newline at end of file diff --git a/verl/docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.base b/verl/docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.base new file mode 100644 index 0000000000000000000000000000000000000000..25b1d9431e237287bfd720788974894d1c07873e --- /dev/null +++ b/verl/docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.base @@ -0,0 +1,113 @@ +# Base Docker Image of verl, with CUDA/Torch/FlashAttn/Apex/TransformerEngine, without other frameworks +# Target: verlai/verl:base-v2-cu124-cudnn9.8-torch2.6-fa2.8.0-te2.3 +# Start from the NVIDIA official image (ubuntu-22.04 + cuda-12.6 + python-3.10) +# https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-08.html +FROM nvcr.io/nvidia/pytorch:24.08-py3 + +# Define environments +ENV MAX_JOBS=16 +ENV VLLM_WORKER_MULTIPROC_METHOD=spawn +ENV DEBIAN_FRONTEND=noninteractive +ENV NODE_OPTIONS="" +ENV PIP_ROOT_USER_ACTION=ignore +ENV HF_HUB_ENABLE_HF_TRANSFER="1" + +# Define installation arguments +ARG APT_SOURCE=https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ +ARG PIP_INDEX=https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple + +# Set apt source +RUN cp /etc/apt/sources.list /etc/apt/sources.list.bak && \ + { \ + echo "deb ${APT_SOURCE} jammy main restricted universe multiverse"; \ + echo "deb ${APT_SOURCE} jammy-updates main restricted universe multiverse"; \ + echo "deb ${APT_SOURCE} jammy-backports main restricted universe multiverse"; \ + echo "deb ${APT_SOURCE} jammy-security main restricted universe multiverse"; \ + } > /etc/apt/sources.list + +# Install systemctl +RUN apt-get update && \ + apt-get install -y -o Dpkg::Options::="--force-confdef" systemd && \ + apt-get clean + +# Install tini +RUN apt-get update && \ + apt-get install -y tini aria2 && \ + apt-get clean + +# Change pip source +RUN pip config set global.index-url "${PIP_INDEX}" && \ + pip config set global.extra-index-url "${PIP_INDEX}" && \ + python -m pip install --upgrade pip + +# Uninstall nv-pytorch fork +RUN pip uninstall -y torch torchvision torchaudio \ + pytorch-quantization pytorch-triton torch-tensorrt \ + xgboost transformer_engine flash_attn apex megatron-core grpcio + +# Reinstall CUDA 12.4 +RUN aria2c https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-ubuntu2204.pin && \ + mv cuda-ubuntu2204.pin /etc/apt/preferences.d/cuda-repository-pin-600 + +RUN aria2c --always-resume=true --max-tries=99999 https://developer.download.nvidia.com/compute/cuda/12.4.1/local_installers/cuda-repo-ubuntu2204-12-4-local_12.4.1-550.54.15-1_amd64.deb && \ + dpkg -i cuda-repo-ubuntu2204-12-4-local_12.4.1-550.54.15-1_amd64.deb && \ + cp /var/cuda-repo-ubuntu2204-12-4-local/cuda-*-keyring.gpg /usr/share/keyrings/ && \ + apt-get update && \ + apt-get -y install cuda-toolkit-12-4 && \ + rm cuda-repo-ubuntu2204-12-4-local_12.4.1-550.54.15-1_amd64.deb && \ + update-alternatives --set cuda /usr/local/cuda-12.4 && \ + rm -rf /usr/local/cuda-12.6 + +RUN pip install --resume-retries 999 --no-cache-dir torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 + +RUN pip install --resume-retries 999 --no-cache-dir "tensordict==0.6.2" torchdata "transformers[hf_xet]>=4.51.0" accelerate datasets peft hf-transfer \ + "numpy<2.0.0" "pyarrow>=19.0.1" pandas \ + ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \ + pytest py-spy pyext pre-commit ruff + +# Install flash-attn-2.7.4.post1 (cxx11abi=False) +RUN wget -nv https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl && \ + pip install --no-cache-dir flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl + +# Fix packages +RUN pip uninstall -y pynvml nvidia-ml-py && \ + pip install --no-cache-dir --upgrade "nvidia-ml-py>=12.560.30" "fastapi[standard]>=0.115.0" "optree>=0.13.0" "pydantic>=2.9" "grpcio>=1.62.1" + +# Install cudnn +RUN aria2c --max-tries=9999 https://developer.download.nvidia.com/compute/cudnn/9.8.0/local_installers/cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb && \ + dpkg -i cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb && \ + cp /var/cudnn-local-repo-ubuntu2204-9.8.0/cudnn-*-keyring.gpg /usr/share/keyrings/ && \ + apt-get update && \ + apt-get -y install cudnn-cuda-12 && \ + rm cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb + +# Install Apex +RUN git clone https://github.com/NVIDIA/apex.git && \ + cd apex && \ + pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./ + +# Profiling tools +RUN aria2c --always-resume=true --max-tries=99999 https://developer.nvidia.com/downloads/assets/tools/secure/nsight-systems/2025_3/nsight-systems-2025.3.1_2025.3.1.90-1_amd64.deb && \ + apt-get update && apt-get install -y libxcb-cursor0 && \ + dpkg -i ./nsight-systems-2025.3.1_2025.3.1.90-1_amd64.deb && \ + rm -rf /usr/local/cuda/bin/nsys && \ + ln -s /opt/nvidia/nsight-systems/2025.3.1/target-linux-x64/nsys /usr/local/cuda/bin/nsys && \ + rm -rf /usr/local/cuda/bin/nsys-ui && \ + ln -s /opt/nvidia/nsight-systems/2025.3.1/target-linux-x64/nsys-ui /usr/local/cuda/bin/nsys-ui && \ + rm nsight-systems-2025.3.1_2025.3.1.90-1_amd64.deb + +# Fix opencv +RUN pip install --resume-retries 999 --no-cache-dir opencv-python + +RUN pip install --resume-retries 999 --no-cache-dir opencv-fixer && \ + python -c "from opencv_fixer import AutoFix; AutoFix()" + +RUN pip install --resume-retries 999 --no-cache-dir cuda-bindings + +# Reset pip config +RUN pip config unset global.index-url && \ + pip config unset global.extra-index-url + +RUN apt-get update && \ + apt-get install -y libfreeimage3 libfreeimage-dev zlib1g htop + diff --git a/verl/docker/verl0.4-cu124-torch2.6-fa2.7.4/README.md b/verl/docker/verl0.4-cu124-torch2.6-fa2.7.4/README.md new file mode 100644 index 0000000000000000000000000000000000000000..d242e06a570cb7f2b29bc2f26f17159a6c5ae443 --- /dev/null +++ b/verl/docker/verl0.4-cu124-torch2.6-fa2.7.4/README.md @@ -0,0 +1,31 @@ +# verl image with verl v0.4.x + +## Important packages version + +```txt +cuda==12.4 +cudnn==9.8.0 +torch==2.6.0 +flash_attn=2.7.4 +sglang==0.4.6.post5 +vllm==0.8.5.post1 +nvidia-cudnn-cu12==9.8.0.87 +transformer_engine==2.3 +megatron.core==core_v0.12.2 +# Preview +transformer_engine==2.5 +megatron.core==core_r0.13.0 +``` + +## Target + +- Base image: + - `verlai/verl:base-verl0.4-cu124-cudnn9.8-torch2.6-fa2.7.4` +- App image: + - `verlai/verl:app-verl0.4-sglang0.4.6.post5-vllm0.8.5-mcore0.12.2-te2.2`: SGLang requires vLLM in 0.4.6.post5 version, vLLM can have some package conflicts with SGLang + - `verlai/verl:app-verl0.4-sglang0.4.6.post5-vllm0.8.5-mcore0.12.2-te2.2-deepep`: Built with deepep + - `verlai/verl:app-verl0.4-vllm0.8.5-mcore0.12.2-te2.2` + - `verlai/verl:app-verl0.4-vllm0.8.5-mcore0.12.2-te2.2-deepep`: Built with deepep +- Preview image: + - `verlai/verl:app-verl0.4-sglang0.4.6.post5-vllm0.8.5-mcore0.13.0-te2.2-preview` + - `verlai/verl:app-verl0.4-vllm0.8.5-mcore0.13.0-te2.2-preview` \ No newline at end of file diff --git a/verl/docker/verl0.5-cu126-torch2.7-fa2.7.4/Dockerfile.app.sglang0.4.10.post2.mcore0.13 b/verl/docker/verl0.5-cu126-torch2.7-fa2.7.4/Dockerfile.app.sglang0.4.10.post2.mcore0.13 new file mode 100644 index 0000000000000000000000000000000000000000..d1be8fb450bf0fe62fcf3d2cc3a140867ce99827 --- /dev/null +++ b/verl/docker/verl0.5-cu126-torch2.7-fa2.7.4/Dockerfile.app.sglang0.4.10.post2.mcore0.13 @@ -0,0 +1,37 @@ +# Start from the verl base image +# Dockerfile.base +FROM verlai/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.1-fa2.7.4 + +# Define environments +ENV MAX_JOBS=8 +ENV VLLM_WORKER_MULTIPROC_METHOD=spawn +ENV DEBIAN_FRONTEND=noninteractive +ENV NODE_OPTIONS="" +ENV PIP_ROOT_USER_ACTION=ignore +ENV HF_HUB_ENABLE_HF_TRANSFER="1" + +# Install sglang-0.4.10 +# Install FlashInfer Python package +RUN pip install --upgrade pip setuptools packaging +RUN pip install --resume-retries 999 --no-cache-dir --no-build-isolation flashinfer-python==0.2.9rc1 +RUN pip install --resume-retries 999 --no-cache-dir --no-build-isolation "sglang[all]==0.4.10.post2" + +# Fix packages +RUN pip install --no-cache-dir "tensordict==0.6.2" "transformers[hf_xet]==4.55.4" accelerate datasets peft hf-transfer \ + "numpy<2.0.0" "pyarrow>=19.0.1" pandas \ + ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \ + pytest py-spy pyext pre-commit ruff + +RUN pip uninstall -y pynvml nvidia-ml-py && \ + pip install --resume-retries 999 --no-cache-dir --upgrade "nvidia-ml-py>=12.560.30" "fastapi[standard]>=0.115.0" "optree>=0.13.0" "pydantic>=2.9" "grpcio>=1.62.1" + +RUN pip install --resume-retries 999 --no-cache-dir nvidia-cudnn-cu12==9.8.0.87 + +# Install TransformerEngine +RUN export NVTE_FRAMEWORK=pytorch && pip3 install --resume-retries 999 --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@v2.2.1 + +# Install Megatron-LM +RUN pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.13.0 + +# Install mbridge +RUN pip3 install --no-cache-dir mbridge diff --git a/verl/docker/verl0.5-cu126-torch2.7-fa2.7.4/Dockerfile.app.sglang0.4.9.post6.mcore0.13 b/verl/docker/verl0.5-cu126-torch2.7-fa2.7.4/Dockerfile.app.sglang0.4.9.post6.mcore0.13 new file mode 100644 index 0000000000000000000000000000000000000000..d79201a92eec3d778e32a023cfa3a3b32bf3ac55 --- /dev/null +++ b/verl/docker/verl0.5-cu126-torch2.7-fa2.7.4/Dockerfile.app.sglang0.4.9.post6.mcore0.13 @@ -0,0 +1,37 @@ +# Start from the verl base image +# Dockerfile.base +FROM verlai/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.1-fa2.7.4 + +# Define environments +ENV MAX_JOBS=8 +ENV VLLM_WORKER_MULTIPROC_METHOD=spawn +ENV DEBIAN_FRONTEND=noninteractive +ENV NODE_OPTIONS="" +ENV PIP_ROOT_USER_ACTION=ignore +ENV HF_HUB_ENABLE_HF_TRANSFER="1" + +# Install sglang-0.4.10 +# Install FlashInfer Python package +RUN pip install --upgrade pip setuptools packaging +RUN pip install --resume-retries 999 --no-cache-dir --no-build-isolation flashinfer-python==0.2.9rc1 +RUN pip install --resume-retries 999 --no-cache-dir --no-build-isolation "sglang[all]==0.4.9.post6" + +# Fix packages +RUN pip install --no-cache-dir "tensordict==0.6.2" "transformers[hf_xet]==4.55.4" accelerate datasets peft hf-transfer \ + "numpy<2.0.0" "pyarrow>=19.0.1" pandas \ + ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \ + pytest py-spy pyext pre-commit ruff + +RUN pip uninstall -y pynvml nvidia-ml-py && \ + pip install --resume-retries 999 --no-cache-dir --upgrade "nvidia-ml-py>=12.560.30" "fastapi[standard]>=0.115.0" "optree>=0.13.0" "pydantic>=2.9" "grpcio>=1.62.1" + +RUN pip install --resume-retries 999 --no-cache-dir nvidia-cudnn-cu12==9.8.0.87 + +# Install TransformerEngine +RUN export NVTE_FRAMEWORK=pytorch && pip3 install --resume-retries 999 --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@v2.2.1 + +# Install Megatron-LM +RUN pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.13.0 + +# Install mbridge +RUN pip3 install --no-cache-dir mbridge \ No newline at end of file diff --git a/verl/docker/verl0.5-cu126-torch2.7-fa2.7.4/Dockerfile.app.vllm.mcore0.13 b/verl/docker/verl0.5-cu126-torch2.7-fa2.7.4/Dockerfile.app.vllm.mcore0.13 new file mode 100644 index 0000000000000000000000000000000000000000..9d73e0ffeeb1c8d17c6d9022d5e5411dc5d70cc1 --- /dev/null +++ b/verl/docker/verl0.5-cu126-torch2.7-fa2.7.4/Dockerfile.app.vllm.mcore0.13 @@ -0,0 +1,38 @@ +# Start from the verl base image +# Dockerfile.base +FROM verlai/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.1-fa2.7.4 + +# Define environments +ENV MAX_JOBS=32 +ENV VLLM_WORKER_MULTIPROC_METHOD=spawn +ENV DEBIAN_FRONTEND=noninteractive +ENV NODE_OPTIONS="" +ENV PIP_ROOT_USER_ACTION=ignore +ENV HF_HUB_ENABLE_HF_TRANSFER="1" + +# Install torch-2.7.1+cu126 + vllm-0.10.0 +RUN pip install --resume-retries 999 --no-cache-dir vllm==0.10.0 + +# Fix packages +# transformers 4.54.0 still not support +RUN pip install --no-cache-dir "tensordict==0.6.2" "transformers[hf_xet]>=4.55.4" accelerate datasets peft hf-transfer \ + "numpy<2.0.0" "pyarrow>=19.0.1" pandas \ + ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \ + pytest py-spy pyext pre-commit ruff + +RUN pip uninstall -y pynvml nvidia-ml-py && \ + pip install --resume-retries 999 --no-cache-dir --upgrade "nvidia-ml-py>=12.560.30" "fastapi[standard]>=0.115.0" "optree>=0.13.0" "pydantic>=2.9" "grpcio>=1.62.1" + +RUN pip install --resume-retries 999 --no-cache-dir nvidia-cudnn-cu12==9.8.0.87 + +# Install TransformerEngine +RUN export NVTE_FRAMEWORK=pytorch && pip3 install --resume-retries 999 --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@v2.2.1 + +# Install Megatron-LM +RUN pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.13.0 + +# Install mbridge +RUN pip3 install --no-cache-dir mbridge + +# Fix qwen vl +RUN pip3 install --no-cache-dir --no-deps trl \ No newline at end of file diff --git a/verl/docker/verl0.5-cu126-torch2.7-fa2.7.4/Dockerfile.app.vllm.mcore0.15 b/verl/docker/verl0.5-cu126-torch2.7-fa2.7.4/Dockerfile.app.vllm.mcore0.15 new file mode 100644 index 0000000000000000000000000000000000000000..296fd3b74f641efce2c90d6aec81fdcb81b8a867 --- /dev/null +++ b/verl/docker/verl0.5-cu126-torch2.7-fa2.7.4/Dockerfile.app.vllm.mcore0.15 @@ -0,0 +1,39 @@ +# Start from the verl base image +# Dockerfile.base +FROM iseekyan/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.1-fa2.7.4-h100 + +# Define environments +ENV MAX_JOBS=32 +ENV VLLM_WORKER_MULTIPROC_METHOD=spawn +ENV DEBIAN_FRONTEND=noninteractive +ENV NODE_OPTIONS="" +ENV PIP_ROOT_USER_ACTION=ignore +ENV HF_HUB_ENABLE_HF_TRANSFER="1" + +# Install torch-2.7.1+cu126 + vllm-0.10.0 +RUN pip install --resume-retries 999 --no-cache-dir vllm==0.10.0 + +# Fix packages +# transformers 4.54.0 still not support +RUN pip install --no-cache-dir "tensordict==0.6.2" "transformers[hf_xet]>=4.55.4" accelerate datasets peft hf-transfer \ + "numpy<2.0.0" "pyarrow>=19.0.1" pandas \ + ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \ + pytest py-spy pyext pre-commit ruff + +RUN pip uninstall -y pynvml nvidia-ml-py && \ + pip install --resume-retries 999 --no-cache-dir --upgrade "nvidia-ml-py>=12.560.30" "fastapi[standard]>=0.115.0" "optree>=0.13.0" "pydantic>=2.9" "grpcio>=1.62.1" + +RUN pip install --resume-retries 999 --no-cache-dir nvidia-cudnn-cu12==9.8.0.87 + +# Install TransformerEngine +RUN export NVTE_FRAMEWORK=pytorch && pip3 install --resume-retries 999 --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.7 +RUN pip install onnxscript + +# Install Megatron-LM +RUN pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.15.0rc4 + +# Install mbridge +RUN pip3 install --no-cache-dir mbridge==v0.15.0 + +# Fix qwen vl +RUN pip3 install --no-cache-dir --no-deps trl \ No newline at end of file diff --git a/verl/docker/verl0.5-cu126-torch2.7-fa2.7.4/Dockerfile.base.torch2.7.1 b/verl/docker/verl0.5-cu126-torch2.7-fa2.7.4/Dockerfile.base.torch2.7.1 new file mode 100644 index 0000000000000000000000000000000000000000..6d85cf512016e4cc0988da3d099230d5d0c2b5d7 --- /dev/null +++ b/verl/docker/verl0.5-cu126-torch2.7-fa2.7.4/Dockerfile.base.torch2.7.1 @@ -0,0 +1,133 @@ +# Base Docker Image of verl, with CUDA/Torch/FlashAttn/Apex/TransformerEngine, without other frameworks +# Target: verlai/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.1-fa2.8.0-fi0.2.6 +# Start from the NVIDIA official image (ubuntu-22.04 + cuda-12.6 + python-3.10) +# https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-08.html +FROM nvcr.io/nvidia/pytorch:24.08-py3 + +# Define environments +ENV MAX_JOBS=16 +ENV VLLM_WORKER_MULTIPROC_METHOD=spawn +ENV DEBIAN_FRONTEND=noninteractive +ENV NODE_OPTIONS="" +ENV PIP_ROOT_USER_ACTION=ignore +ENV HF_HUB_ENABLE_HF_TRANSFER="1" + +# Define installation arguments +ARG APT_SOURCE=https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ +ARG PIP_INDEX=https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple + +# Set apt source +RUN cp /etc/apt/sources.list /etc/apt/sources.list.bak && \ + { \ + echo "deb ${APT_SOURCE} jammy main restricted universe multiverse"; \ + echo "deb ${APT_SOURCE} jammy-updates main restricted universe multiverse"; \ + echo "deb ${APT_SOURCE} jammy-backports main restricted universe multiverse"; \ + echo "deb ${APT_SOURCE} jammy-security main restricted universe multiverse"; \ + } > /etc/apt/sources.list + +# Install systemctl +RUN apt-get update && \ + apt-get install -y -o Dpkg::Options::="--force-confdef" systemd && \ + apt-get clean + +# Install tini +RUN apt-get update && \ + apt-get install -y tini aria2 libfreeimage3 libfreeimage-dev zlib1g htop && \ + apt-get clean + +# Change pip source +RUN pip config set global.index-url "${PIP_INDEX}" && \ + pip config set global.extra-index-url "${PIP_INDEX}" && \ + python -m pip install --upgrade pip + +# Uninstall nv-pytorch fork +RUN pip uninstall -y torch torchvision torchaudio \ + pytorch-quantization pytorch-triton torch-tensorrt \ + xgboost transformer_engine flash_attn apex megatron-core grpcio + +RUN pip install --resume-retries 999 --no-cache-dir torch==2.7.1 torchvision==0.22.1 torchaudio==2.7.1 + +# Install flash-attn-2.7.4.post1, although built with torch2.6, it is compatible with torch2.7 +# https://github.com/Dao-AILab/flash-attention/issues/1644#issuecomment-2899396361 +RUN ABI_FLAG=$(python -c "import torch; print('TRUE' if torch._C._GLIBCXX_USE_CXX11_ABI else 'FALSE')") && \ + URL="https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abi${ABI_FLAG}-cp310-cp310-linux_x86_64.whl" && \ + FILE="flash_attn-2.7.4.post1+cu12torch2.6cxx11abi${ABI_FLAG}-cp310-cp310-linux_x86_64.whl" && \ + wget -nv "${URL}" && \ + pip install --no-cache-dir "${FILE}" + +# Fix packages +RUN pip uninstall -y pynvml nvidia-ml-py && \ + pip install --no-cache-dir --upgrade "nvidia-ml-py>=12.560.30" "fastapi[standard]>=0.115.0" "optree>=0.13.0" "pydantic>=2.9" "grpcio>=1.62.1" + +# Install cudnn +RUN aria2c --max-tries=9999 https://developer.download.nvidia.com/compute/cudnn/9.8.0/local_installers/cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb && \ + dpkg -i cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb && \ + cp /var/cudnn-local-repo-ubuntu2204-9.8.0/cudnn-*-keyring.gpg /usr/share/keyrings/ && \ + apt-get update && \ + apt-get -y install cudnn-cuda-12 && \ + rm cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb + +# Install Apex +RUN pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" --resume-retries 999 git+https://github.com/NVIDIA/apex.git + +# Profiling tools +RUN aria2c --always-resume=true --max-tries=99999 https://developer.nvidia.com/downloads/assets/tools/secure/nsight-systems/2025_3/nsight-systems-2025.3.1_2025.3.1.90-1_amd64.deb && \ + apt-get update && apt-get install -y libxcb-cursor0 + +RUN apt-get install -y ./nsight-systems-2025.3.1_2025.3.1.90-1_amd64.deb && \ + rm -rf /usr/local/cuda/bin/nsys && \ + ln -s /opt/nvidia/nsight-systems/2025.3.1/target-linux-x64/nsys /usr/local/cuda/bin/nsys && \ + rm -rf /usr/local/cuda/bin/nsys-ui && \ + ln -s /opt/nvidia/nsight-systems/2025.3.1/target-linux-x64/nsys-ui /usr/local/cuda/bin/nsys-ui && \ + rm nsight-systems-2025.3.1_2025.3.1.90-1_amd64.deb + +RUN pip install --resume-retries 999 --no-cache-dir "tensordict==0.6.2" torchdata "transformers[hf_xet]>=4.52.3" accelerate datasets peft hf-transfer \ + "numpy<2.0.0" "pyarrow>=19.0.1" pandas cuda-bindings \ + ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \ + pytest py-spy pyext pre-commit ruff + +# Install DeepEP +## the dependency of IBGDA +RUN ln -s /usr/lib/x86_64-linux-gnu/libmlx5.so.1 /usr/lib/x86_64-linux-gnu/libmlx5.so + +## Clone and build deepep and deepep-nvshmem +RUN git clone -b v2.3.1 https://github.com/NVIDIA/gdrcopy.git && \ + git clone https://github.com/deepseek-ai/DeepEP.git && \ + cd DeepEP && git checkout a84a248 + +# Prepare nvshmem +RUN wget https://developer.nvidia.com/downloads/assets/secure/nvshmem/nvshmem_src_3.2.5-1.txz && \ + tar -xvf nvshmem_src_3.2.5-1.txz && mv nvshmem_src deepep-nvshmem && \ + cd deepep-nvshmem && git apply ../DeepEP/third-party/nvshmem.patch + +ENV CUDA_HOME=/usr/local/cuda +### Set MPI environment variables. Having errors when not set. +ENV CPATH=/usr/local/mpi/include:$CPATH +ENV LD_LIBRARY_PATH=/usr/local/mpi/lib:$LD_LIBRARY_PATH +ENV LD_LIBRARY_PATH=/usr/local/x86_64-linux-gnu:$LD_LIBRARY_PATH +ENV GDRCOPY_HOME=/workspace/gdrcopy + +## Build deepep-nvshmem +RUN cd deepep-nvshmem && \ + NVSHMEM_SHMEM_SUPPORT=0 \ + NVSHMEM_UCX_SUPPORT=0 \ + NVSHMEM_USE_NCCL=0 \ + NVSHMEM_MPI_SUPPORT=0 \ + NVSHMEM_IBGDA_SUPPORT=1 \ + NVSHMEM_PMIX_SUPPORT=0 \ + NVSHMEM_TIMEOUT_DEVICE_POLLING=0 \ + NVSHMEM_USE_GDRCOPY=1 \ + cmake -G Ninja -S . -B build/ -DCMAKE_INSTALL_PREFIX=/workspace/deepep-nvshmem/install && cmake --build build/ --target install + +ENV NVSHMEM_DIR=/workspace/deepep-nvshmem/install +ENV LD_LIBRARY_PATH=$NVSHMEM_DIR/lib:$LD_LIBRARY_PATH +ENV PATH=$NVSHMEM_DIR/bin:$PATH + +## Build deepep +RUN cd DeepEP && \ + python setup.py install + +# Reset pip config +RUN pip config unset global.index-url && \ + pip config unset global.extra-index-url + diff --git a/verl/docker/verl0.5-cu126-torch2.7-fa2.7.4/README.md b/verl/docker/verl0.5-cu126-torch2.7-fa2.7.4/README.md new file mode 100644 index 0000000000000000000000000000000000000000..79527b76c241d6aa28c54b0e11b262863fdc1c86 --- /dev/null +++ b/verl/docker/verl0.5-cu126-torch2.7-fa2.7.4/README.md @@ -0,0 +1,27 @@ +# verl image with verl v0.5 + +## Important packages version + +```txt +cuda==12.6 +cudnn==9.8.0 +torch==2.7.1 +flash_attn=2.7.4.post1 +sglang==0.4.9.post6 +vllm==0.8.5.post1 +nvidia-cudnn-cu12==9.8.0.87 +transformer_engine==2.3 +megatron.core==core_v0.12.2 +# Preview +transformer_engine==2.5 +megatron.core==core_r0.13.0 +``` + +## Target + +- Base image: + - `verlai/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.1-fa2.7.4`: We offer a base image with deep ep built in, for vllm/sglang +- App image: + - `verlai/verl:app-verl0.5-transformers4.55.4-vllm0.10.0-mcore0.13.0-te2.2` + - `verlai/verl:app-verl0.5-transformers4.55.4-sglang0.4.10.post2-mcore0.13.0-te2.2` + - `iseekyan/verl:app-verl0.5-transformers4.55.4-vllm0.10.0-mcore0.15.0-te2.7` diff --git a/verl/docker/verl0.5-cu126-torch2.7.1-fa2.8.0/Dockerfile.app.sglang.mcore0.12 b/verl/docker/verl0.5-cu126-torch2.7.1-fa2.8.0/Dockerfile.app.sglang.mcore0.12 new file mode 100644 index 0000000000000000000000000000000000000000..205814a234bf4d1162a2fa89f1cd051d6463fb47 --- /dev/null +++ b/verl/docker/verl0.5-cu126-torch2.7.1-fa2.8.0/Dockerfile.app.sglang.mcore0.12 @@ -0,0 +1,37 @@ +# Start from the verl base image +# Dockerfile.base +FROM verlai/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.1-fa2.8.0 + +# Define environments +ENV MAX_JOBS=8 +ENV VLLM_WORKER_MULTIPROC_METHOD=spawn +ENV DEBIAN_FRONTEND=noninteractive +ENV NODE_OPTIONS="" +ENV PIP_ROOT_USER_ACTION=ignore +ENV HF_HUB_ENABLE_HF_TRANSFER="1" + +# Install sglang-0.4.8 and torch-memory-saver +# Install FlashInfer Python package +RUN pip install --upgrade pip setuptools packaging +RUN pip install --resume-retries 999 --no-cache-dir --no-build-isolation flashinfer-python==0.2.6.post1 +RUN pip install --resume-retries 999 --no-cache-dir "sglang[all]==0.4.8" && pip install torch-memory-saver --no-cache-dir + +# Fix packages +RUN pip install --no-cache-dir "tensordict==0.6.2" "transformers[hf_xet]>=4.51.0" accelerate datasets peft hf-transfer \ + "numpy<2.0.0" "pyarrow>=19.0.1" pandas \ + ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \ + pytest py-spy pyext pre-commit ruff + +RUN pip uninstall -y pynvml nvidia-ml-py && \ + pip install --resume-retries 999 --no-cache-dir --upgrade "nvidia-ml-py>=12.560.30" "fastapi[standard]>=0.115.0" "optree>=0.13.0" "pydantic>=2.9" "grpcio>=1.62.1" + +RUN pip install --resume-retries 999 --no-cache-dir nvidia-cudnn-cu12==9.8.0.87 + +# Install TransformerEngine +RUN export NVTE_FRAMEWORK=pytorch && pip3 install --resume-retries 999 --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@v2.3 + +# Install Megatron-LM +RUN pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.12.2 + +# Install mbridge +RUN pip3 install --no-cache-dir mbridge \ No newline at end of file diff --git a/verl/docker/verl0.5-cu126-torch2.7.1-fa2.8.0/Dockerfile.app.sglang.mcore0.13.preview b/verl/docker/verl0.5-cu126-torch2.7.1-fa2.8.0/Dockerfile.app.sglang.mcore0.13.preview new file mode 100644 index 0000000000000000000000000000000000000000..5704813d9eee80a7ca790b204a836e81d58cbf1a --- /dev/null +++ b/verl/docker/verl0.5-cu126-torch2.7.1-fa2.8.0/Dockerfile.app.sglang.mcore0.13.preview @@ -0,0 +1,37 @@ +# Start from the verl base image +# Dockerfile.base +FROM verlai/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.1-fa2.8.0 + +# Define environments +ENV MAX_JOBS=8 +ENV VLLM_WORKER_MULTIPROC_METHOD=spawn +ENV DEBIAN_FRONTEND=noninteractive +ENV NODE_OPTIONS="" +ENV PIP_ROOT_USER_ACTION=ignore +ENV HF_HUB_ENABLE_HF_TRANSFER="1" + +# Install sglang-0.4.8 and torch-memory-saver +# Install FlashInfer Python package +RUN pip install --upgrade pip setuptools packaging +RUN pip install --resume-retries 999 --no-cache-dir --no-build-isolation flashinfer-python==0.2.6.post1 +RUN pip install --resume-retries 999 --no-cache-dir "sglang[all]==0.4.8" && pip install torch-memory-saver --no-cache-dir + +# Fix packages +RUN pip install --no-cache-dir "tensordict==0.6.2" "transformers[hf_xet]>=4.51.0" accelerate datasets peft hf-transfer \ + "numpy<2.0.0" "pyarrow>=19.0.1" pandas \ + ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \ + pytest py-spy pyext pre-commit ruff + +RUN pip uninstall -y pynvml nvidia-ml-py && \ + pip install --resume-retries 999 --no-cache-dir --upgrade "nvidia-ml-py>=12.560.30" "fastapi[standard]>=0.115.0" "optree>=0.13.0" "pydantic>=2.9" "grpcio>=1.62.1" + +RUN pip install --resume-retries 999 --no-cache-dir nvidia-cudnn-cu12==9.8.0.87 + +# Install TransformerEngine +RUN export NVTE_FRAMEWORK=pytorch && pip3 install --resume-retries 999 --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.5 + +# Install Megatron-LM +RUN pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.12.2 + +# Install mbridge +RUN pip3 install --no-cache-dir mbridge \ No newline at end of file diff --git a/verl/docker/verl0.5-cu126-torch2.7.1-fa2.8.0/Dockerfile.base b/verl/docker/verl0.5-cu126-torch2.7.1-fa2.8.0/Dockerfile.base new file mode 100644 index 0000000000000000000000000000000000000000..53089fba400d56a3a37e1e943439c367ee40df55 --- /dev/null +++ b/verl/docker/verl0.5-cu126-torch2.7.1-fa2.8.0/Dockerfile.base @@ -0,0 +1,132 @@ +# Base Docker Image of verl, with CUDA/Torch/FlashAttn/Apex/TransformerEngine, without other frameworks +# Target: verlai/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.1-fa2.8.0-fi0.2.6 +# Start from the NVIDIA official image (ubuntu-22.04 + cuda-12.6 + python-3.10) +# https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-08.html +FROM nvcr.io/nvidia/pytorch:24.08-py3 + +# Define environments +ENV MAX_JOBS=16 +ENV VLLM_WORKER_MULTIPROC_METHOD=spawn +ENV DEBIAN_FRONTEND=noninteractive +ENV NODE_OPTIONS="" +ENV PIP_ROOT_USER_ACTION=ignore +ENV HF_HUB_ENABLE_HF_TRANSFER="1" + +# Define installation arguments +ARG APT_SOURCE=https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ +ARG PIP_INDEX=https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple + +# Set apt source +RUN cp /etc/apt/sources.list /etc/apt/sources.list.bak && \ + { \ + echo "deb ${APT_SOURCE} jammy main restricted universe multiverse"; \ + echo "deb ${APT_SOURCE} jammy-updates main restricted universe multiverse"; \ + echo "deb ${APT_SOURCE} jammy-backports main restricted universe multiverse"; \ + echo "deb ${APT_SOURCE} jammy-security main restricted universe multiverse"; \ + } > /etc/apt/sources.list + +# Install systemctl +RUN apt-get update && \ + apt-get install -y -o Dpkg::Options::="--force-confdef" systemd && \ + apt-get clean + +# Install tini +RUN apt-get update && \ + apt-get install -y tini aria2 libfreeimage3 libfreeimage-dev zlib1g htop && \ + apt-get clean + +# Change pip source +RUN pip config set global.index-url "${PIP_INDEX}" && \ + pip config set global.extra-index-url "${PIP_INDEX}" && \ + python -m pip install --upgrade pip + +# Uninstall nv-pytorch fork +RUN pip uninstall -y torch torchvision torchaudio \ + pytorch-quantization pytorch-triton torch-tensorrt \ + xgboost transformer_engine flash_attn apex megatron-core grpcio + +RUN pip install --resume-retries 999 --no-cache-dir torch==2.7.1 torchvision==0.22.1 torchaudio==2.7.1 + +# Install flash-attn-2.8.0.post2 (cxx11abi=True) +RUN ABI_FLAG=$(python -c "import torch; print('TRUE' if torch._C._GLIBCXX_USE_CXX11_ABI else 'FALSE')") && \ + URL="https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.0.post2/flash_attn-2.8.0.post2+cu12torch2.7cxx11abi${ABI_FLAG}-cp310-cp310-linux_x86_64.whl" && \ + FILE="flash_attn-2.8.0.post2+cu12torch2.7cxx11abi${ABI_FLAG}-cp310-cp310-linux_x86_64.whl" && \ + wget -nv "${URL}" && \ + pip install --no-cache-dir "${FILE}" + +# Fix packages +RUN pip uninstall -y pynvml nvidia-ml-py && \ + pip install --no-cache-dir --upgrade "nvidia-ml-py>=12.560.30" "fastapi[standard]>=0.115.0" "optree>=0.13.0" "pydantic>=2.9" "grpcio>=1.62.1" + +# Install cudnn +RUN aria2c --max-tries=9999 https://developer.download.nvidia.com/compute/cudnn/9.8.0/local_installers/cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb && \ + dpkg -i cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb && \ + cp /var/cudnn-local-repo-ubuntu2204-9.8.0/cudnn-*-keyring.gpg /usr/share/keyrings/ && \ + apt-get update && \ + apt-get -y install cudnn-cuda-12 && \ + rm cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb + +# Install Apex +RUN pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" --resume-retries 999 git+https://github.com/NVIDIA/apex.git + +# Profiling tools +RUN aria2c --always-resume=true --max-tries=99999 https://developer.nvidia.com/downloads/assets/tools/secure/nsight-systems/2025_3/nsight-systems-2025.3.1_2025.3.1.90-1_amd64.deb && \ + apt-get update && apt-get install -y libxcb-cursor0 + +RUN apt-get install -y ./nsight-systems-2025.3.1_2025.3.1.90-1_amd64.deb && \ + rm -rf /usr/local/cuda/bin/nsys && \ + ln -s /opt/nvidia/nsight-systems/2025.3.1/target-linux-x64/nsys /usr/local/cuda/bin/nsys && \ + rm -rf /usr/local/cuda/bin/nsys-ui && \ + ln -s /opt/nvidia/nsight-systems/2025.3.1/target-linux-x64/nsys-ui /usr/local/cuda/bin/nsys-ui && \ + rm nsight-systems-2025.3.1_2025.3.1.90-1_amd64.deb + +RUN pip install --resume-retries 999 --no-cache-dir "tensordict==0.6.2" torchdata "transformers[hf_xet]>=4.53" accelerate datasets peft hf-transfer \ + "numpy<2.0.0" "pyarrow>=19.0.1" pandas cuda-bindings \ + ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \ + pytest py-spy pyext pre-commit ruff + +# Install DeepEP +## the dependency of IBGDA +RUN ln -s /usr/lib/x86_64-linux-gnu/libmlx5.so.1 /usr/lib/x86_64-linux-gnu/libmlx5.so + +## Clone and build deepep and deepep-nvshmem +RUN git clone -b v2.3.1 https://github.com/NVIDIA/gdrcopy.git && \ + git clone https://github.com/deepseek-ai/DeepEP.git && \ + cd DeepEP && git checkout a84a248 + +# Prepare nvshmem +RUN wget https://developer.nvidia.com/downloads/assets/secure/nvshmem/nvshmem_src_3.2.5-1.txz && \ + tar -xvf nvshmem_src_3.2.5-1.txz && mv nvshmem_src deepep-nvshmem && \ + cd deepep-nvshmem && git apply ../DeepEP/third-party/nvshmem.patch + +ENV CUDA_HOME=/usr/local/cuda +### Set MPI environment variables. Having errors when not set. +ENV CPATH=/usr/local/mpi/include:$CPATH +ENV LD_LIBRARY_PATH=/usr/local/mpi/lib:$LD_LIBRARY_PATH +ENV LD_LIBRARY_PATH=/usr/local/x86_64-linux-gnu:$LD_LIBRARY_PATH +ENV GDRCOPY_HOME=/workspace/gdrcopy + +## Build deepep-nvshmem +RUN cd deepep-nvshmem && \ + NVSHMEM_SHMEM_SUPPORT=0 \ + NVSHMEM_UCX_SUPPORT=0 \ + NVSHMEM_USE_NCCL=0 \ + NVSHMEM_MPI_SUPPORT=0 \ + NVSHMEM_IBGDA_SUPPORT=1 \ + NVSHMEM_PMIX_SUPPORT=0 \ + NVSHMEM_TIMEOUT_DEVICE_POLLING=0 \ + NVSHMEM_USE_GDRCOPY=1 \ + cmake -G Ninja -S . -B build/ -DCMAKE_INSTALL_PREFIX=/workspace/deepep-nvshmem/install && cmake --build build/ --target install + +ENV NVSHMEM_DIR=/workspace/deepep-nvshmem/install +ENV LD_LIBRARY_PATH=$NVSHMEM_DIR/lib:$LD_LIBRARY_PATH +ENV PATH=$NVSHMEM_DIR/bin:$PATH + +## Build deepep +RUN cd DeepEP && \ + python setup.py install + +# Reset pip config +RUN pip config unset global.index-url && \ + pip config unset global.extra-index-url + diff --git a/verl/docker/verl0.5-cu126-torch2.7.1-fa2.8.0/README.md b/verl/docker/verl0.5-cu126-torch2.7.1-fa2.8.0/README.md new file mode 100644 index 0000000000000000000000000000000000000000..4cbf4a974a5844376ddc36ade862140af55a6ba5 --- /dev/null +++ b/verl/docker/verl0.5-cu126-torch2.7.1-fa2.8.0/README.md @@ -0,0 +1,27 @@ +# verl image with verl v0.5 + +## Important packages version + +```txt +cuda==12.6 +cudnn==9.8.0 +torch==2.7.1 +flash_attn=2.8.0 ## +sglang==0.4.8 +vllm==0.8.5.post1 +nvidia-cudnn-cu12==9.8.0.87 +transformer_engine==2.3 +megatron.core==core_v0.12.2 +# Preview +transformer_engine==2.5 +megatron.core==core_r0.13.0 +``` + +## Target + +- Base image: + - `verlai/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.1-fa2.8.0`: We offer a base image with deep ep built in +- App image: + - `verlai/verl:app-verl0.5-sglang0.4.9-mcore0.12.2` + - `verlai/verl:app-verl0.5-sglang0.4.9-mcore0.13.0-preview` +- vllm temporarily not support latest version \ No newline at end of file diff --git a/verl/docker/verl0.5-preview-cu128-torch2.7.1-fa2.8.0/Dockerfile.app.sglang.megatron b/verl/docker/verl0.5-preview-cu128-torch2.7.1-fa2.8.0/Dockerfile.app.sglang.megatron new file mode 100644 index 0000000000000000000000000000000000000000..d41ea19d69bccb39436a79a0798a9cb0cbefec04 --- /dev/null +++ b/verl/docker/verl0.5-preview-cu128-torch2.7.1-fa2.8.0/Dockerfile.app.sglang.megatron @@ -0,0 +1,36 @@ +# Start from the verl base image +# Dockerfile.base +FROM verlai/verl:base-verl0.5-preview-cu128-cudnn9.8-torch2.7.1-fa2.8.0-fi0.2.6 + +# Define environments +ENV MAX_JOBS=8 +ENV VLLM_WORKER_MULTIPROC_METHOD=spawn +ENV DEBIAN_FRONTEND=noninteractive +ENV NODE_OPTIONS="" +ENV PIP_ROOT_USER_ACTION=ignore +ENV HF_HUB_ENABLE_HF_TRANSFER="1" + +# Install sglang-0.4.8 and torch-memory-saver +# Install FlashInfer Python package +RUN pip install --resume-retries 999 --no-cache-dir --no-build-isolation flashinfer-python==0.2.6.post1 +RUN pip install --resume-retries 999 --no-cache-dir "sglang[all]==0.4.8" && pip install torch-memory-saver --no-cache-dir + +# Fix packages +RUN pip install --no-cache-dir "tensordict==0.6.2" "transformers[hf_xet]>=4.51.0" accelerate datasets peft hf-transfer \ + "numpy<2.0.0" "pyarrow>=19.0.1" pandas \ + ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \ + pytest py-spy pre-commit ruff + +RUN pip uninstall -y pynvml nvidia-ml-py && \ + pip install --resume-retries 999 --no-cache-dir --upgrade "nvidia-ml-py>=12.560.30" "fastapi[standard]>=0.115.0" "optree>=0.13.0" "pydantic>=2.9" "grpcio>=1.62.1" + +RUN pip install --resume-retries 999 --no-cache-dir nvidia-cudnn-cu12==9.8.0.87 + +# Install TransformerEngine +RUN export NVTE_FRAMEWORK=pytorch && pip3 install --resume-retries 999 --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.5 + +# Install Megatron-LM +RUN pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/Megatron-LM.git@core_r0.13.0 + +# Install mbridge +RUN pip3 install --no-cache-dir mbridge \ No newline at end of file diff --git a/verl/docker/verl0.5-preview-cu128-torch2.7.1-fa2.8.0/Dockerfile.base b/verl/docker/verl0.5-preview-cu128-torch2.7.1-fa2.8.0/Dockerfile.base new file mode 100644 index 0000000000000000000000000000000000000000..29c49faa848c0e15350abaf48cad7f8a899f4eb6 --- /dev/null +++ b/verl/docker/verl0.5-preview-cu128-torch2.7.1-fa2.8.0/Dockerfile.base @@ -0,0 +1,91 @@ +# Base Docker Image of verl, with CUDA/Torch/FlashAttn/Apex/TransformerEngine, without other frameworks +# Target: verlai/verl:base-verl0.5-preview-cu128-cudnn9.8-torch2.7.1-fa2.8.0-fi0.2.6 +# Start from the NVIDIA official image (ubuntu-22.04 + cuda-12.6 + python-3.10) +# https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-08.html +FROM nvcr.io/nvidia/pytorch:25.02-py3 + +# Define environments +ENV MAX_JOBS=16 +ENV VLLM_WORKER_MULTIPROC_METHOD=spawn +ENV DEBIAN_FRONTEND=noninteractive +ENV NODE_OPTIONS="" +ENV PIP_ROOT_USER_ACTION=ignore +ENV HF_HUB_ENABLE_HF_TRANSFER="1" + +# Define installation arguments +ARG APT_SOURCE=https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ +ARG PIP_INDEX=https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple + +# Set apt source +RUN cp /etc/apt/sources.list /etc/apt/sources.list.bak && \ + { \ + echo "deb ${APT_SOURCE} jammy main restricted universe multiverse"; \ + echo "deb ${APT_SOURCE} jammy-updates main restricted universe multiverse"; \ + echo "deb ${APT_SOURCE} jammy-backports main restricted universe multiverse"; \ + echo "deb ${APT_SOURCE} jammy-security main restricted universe multiverse"; \ + } > /etc/apt/sources.list + +# Install systemctl +RUN apt-get update && \ + apt-get install -y -o Dpkg::Options::="--force-confdef" systemd && \ + apt-get clean + +# Install tini +RUN apt-get update && \ + apt-get install -y tini aria2 libfreeimage3 libfreeimage-dev zlib1g htop && \ + apt-get clean + +# Change pip source +RUN pip config set global.index-url "${PIP_INDEX}" && \ + pip config set global.extra-index-url "${PIP_INDEX}" && \ + python -m pip install --upgrade pip + +# Uninstall nv-pytorch fork +RUN pip uninstall -y torch torchvision torchaudio \ + pytorch-quantization pytorch-triton torch-tensorrt \ + xgboost transformer_engine flash_attn apex megatron-core grpcio + +RUN pip install --resume-retries 999 --no-cache-dir torch==2.7.1 torchvision==0.22.1 torchaudio==2.7.1 --index-url https://download.pytorch.org/whl/cu128 + +# Install flash-attn-2.8.0.post2 (cxx11abi=True) +RUN ABI_FLAG=$(python -c "import torch; print('TRUE' if torch._C._GLIBCXX_USE_CXX11_ABI else 'FALSE')") && \ + URL="https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.0.post2/flash_attn-2.8.0.post2+cu12torch2.7cxx11abi${ABI_FLAG}-cp312-cp312-linux_x86_64.whl" && \ + FILE="flash_attn-2.8.0.post2+cu12torch2.7cxx11abi${ABI_FLAG}-cp312-cp312-linux_x86_64.whl" && \ + wget -nv "${URL}" && \ + pip install --no-cache-dir "${FILE}" + +# Fix packages +RUN pip uninstall -y pynvml nvidia-ml-py && \ + pip install --no-cache-dir --upgrade "nvidia-ml-py>=12.560.30" "fastapi[standard]>=0.115.0" "optree>=0.13.0" "pydantic>=2.9" "grpcio>=1.62.1" + +# Install cudnn +RUN aria2c --max-tries=9999 https://developer.download.nvidia.com/compute/cudnn/9.8.0/local_installers/cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb && \ + dpkg -i cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb && \ + cp /var/cudnn-local-repo-ubuntu2204-9.8.0/cudnn-*-keyring.gpg /usr/share/keyrings/ && \ + apt-get update && \ + apt-get -y install cudnn-cuda-12 && \ + rm cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb + +# Install Apex +RUN pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" --resume-retries 999 git+https://github.com/NVIDIA/apex.git + +# Profiling tools +RUN aria2c --always-resume=true --max-tries=99999 https://developer.nvidia.com/downloads/assets/tools/secure/nsight-systems/2025_3/nsight-systems-2025.3.1_2025.3.1.90-1_amd64.deb && \ + apt-get update && apt-get install -y libxcb-cursor0 + +RUN apt-get install -y ./nsight-systems-2025.3.1_2025.3.1.90-1_amd64.deb && \ + rm -rf /usr/local/cuda/bin/nsys && \ + ln -s /opt/nvidia/nsight-systems/2025.3.1/target-linux-x64/nsys /usr/local/cuda/bin/nsys && \ + rm -rf /usr/local/cuda/bin/nsys-ui && \ + ln -s /opt/nvidia/nsight-systems/2025.3.1/target-linux-x64/nsys-ui /usr/local/cuda/bin/nsys-ui && \ + rm nsight-systems-2025.3.1_2025.3.1.90-1_amd64.deb + +RUN pip install --resume-retries 999 --no-cache-dir "tensordict==0.6.2" torchdata "transformers[hf_xet]>=4.51.0" accelerate datasets peft hf-transfer \ + "numpy<2.0.0" "pyarrow>=19.0.1" pandas cuda-bindings \ + ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \ + pytest py-spy pre-commit ruff + +# Reset pip config +RUN pip config unset global.index-url && \ + pip config unset global.extra-index-url + diff --git a/verl/docker/verl0.5-preview-cu128-torch2.7.1-fa2.8.0/README.md b/verl/docker/verl0.5-preview-cu128-torch2.7.1-fa2.8.0/README.md new file mode 100644 index 0000000000000000000000000000000000000000..bb16d2abece29efc5b839279d78dc21191240b63 --- /dev/null +++ b/verl/docker/verl0.5-preview-cu128-torch2.7.1-fa2.8.0/README.md @@ -0,0 +1,26 @@ +# verl image with verl v0.5 + +## Important packages version + +```txt +cuda==12.8 +cudnn==9.8.0 +torch==2.7.1 +flash_attn=2.8.0 ## +sglang==0.4.8 +transformer_engine==2.5 +megatron.core==core_r0.13.0 +nvidia-cudnn-cu12==9.8.0.87 +``` + +## Target + +- Base image: + - `verlai/verl:base-verl0.5-preview-cu128-cudnn9.8-torch2.7.1-fa2.8.0`: We offer a base image with flash infer 0.2.6.post1 built in +- App image: + - `verlai/verl:app-verl0.5-preview-sglang0.4.8-mcore0.13.0-preview` +- vllm temporarily not support latest version + +## !!!Notice!!! + +- pyext is lack of maintainace and cannot work with python 3.12, consider using replacement and deprecating this package. \ No newline at end of file diff --git a/verl/docker/verl0.6-cu128-torch2.8.0-fa2.7.4/Dockerfile.app.sglang b/verl/docker/verl0.6-cu128-torch2.8.0-fa2.7.4/Dockerfile.app.sglang new file mode 100644 index 0000000000000000000000000000000000000000..23dbea72ccce3030beb76691ba1832fc04dcdc76 --- /dev/null +++ b/verl/docker/verl0.6-cu128-torch2.8.0-fa2.7.4/Dockerfile.app.sglang @@ -0,0 +1,4 @@ +FROM verlai/verl:base-verl0.6-cu128-cudnn9.8-torch2.8.0-fa2.7.4 + +RUN pip install --no-cache-dir "sglang[all]==0.5.2" +RUN pip install --no-cache-dir "torch-memory-saver==0.0.9rc1" diff --git a/verl/docker/verl0.6-cu128-torch2.8.0-fa2.7.4/Dockerfile.base b/verl/docker/verl0.6-cu128-torch2.8.0-fa2.7.4/Dockerfile.base new file mode 100644 index 0000000000000000000000000000000000000000..7bfd48169cc0ed666507d6560669e4c6f3511a11 --- /dev/null +++ b/verl/docker/verl0.6-cu128-torch2.8.0-fa2.7.4/Dockerfile.base @@ -0,0 +1,108 @@ +# Start from the NVIDIA official image (ubuntu-24.04 + cuda-12.8 + python-3.12) +# https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-25-03.html +FROM nvcr.io/nvidia/pytorch:25.03-py3 + +# Define environments +ENV MAX_JOBS=32 +ENV VLLM_WORKER_MULTIPROC_METHOD=spawn +ENV DEBIAN_FRONTEND=noninteractive +ENV NODE_OPTIONS="" +ENV PIP_ROOT_USER_ACTION=ignore +ENV HF_HUB_ENABLE_HF_TRANSFER="1" +ENV PIP_CONSTRAINT="" + +ARG PIP_INDEX=https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple + +# Change pip source +RUN pip config set global.index-url "${PIP_INDEX}" && \ + pip config set global.extra-index-url "${PIP_INDEX}" && \ + pip config set global.no-cache-dir "true" && \ + python -m pip install --upgrade pip + +# Install systemctl +RUN apt-get update && \ + apt-get install -y -o Dpkg::Options::="--force-confdef" systemd && \ + apt-get clean + +# Install libxml2 +RUN apt-get update && \ + apt-get install -y libxml2 aria2 && \ + apt-get clean + +# Uninstall nv-pytorch fork +RUN pip uninstall -y torch torchvision torchaudio \ + pytorch-quantization pytorch-triton torch-tensorrt \ + transformer_engine flash_attn apex megatron-core \ + xgboost opencv grpcio + +# Fix packages +RUN pip install --no-cache-dir tensordict torchdata "transformers[hf_xet]==4.55.4" accelerate datasets peft hf-transfer \ + "numpy<2.0.0" "pyarrow>=19.0.1" pandas \ + ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \ + pytest py-spy pre-commit ruff + +# Fix cv2 +RUN rm -rf /usr/local/lib/python3.11/dist-packages/cv2 + +# Install torch +RUN pip install --no-cache-dir torch==2.8.0 --index-url https://download.pytorch.org/whl/cu128 + +# Install flash-attn +RUN pip install --no-cache-dir --no-build-isolation flash_attn==2.7.4.post1 + +# Install DeepEP +# the dependency of IBGDA +RUN ln -s /usr/lib/x86_64-linux-gnu/libmlx5.so.1 /usr/lib/x86_64-linux-gnu/libmlx5.so + +# Clone and build deepep and deepep-nvshmem +RUN git clone -b v2.3.1 https://github.com/NVIDIA/gdrcopy.git && \ + git clone https://github.com/deepseek-ai/DeepEP.git && \ + cd DeepEP && git checkout a84a248 + +# Prepare nvshmem +RUN wget https://developer.nvidia.com/downloads/assets/secure/nvshmem/nvshmem_src_3.2.5-1.txz && \ + tar -xvf nvshmem_src_3.2.5-1.txz && mv nvshmem_src deepep-nvshmem && \ + cd deepep-nvshmem && git apply ../DeepEP/third-party/nvshmem.patch + +## Build deepep-nvshmem +RUN apt-get install -y ninja-build cmake + +ENV CUDA_HOME=/usr/local/cuda +### Set MPI environment variables. Having errors when not set. +ENV CPATH=/usr/local/mpi/include:$CPATH +ENV LD_LIBRARY_PATH=/usr/local/mpi/lib:$LD_LIBRARY_PATH +ENV LD_LIBRARY_PATH=/usr/local/x86_64-linux-gnu:$LD_LIBRARY_PATH +ENV GDRCOPY_HOME=/workspace/gdrcopy +ENV GDRCOPY_INCLUDE=/workspace/gdrcopy/include + +RUN cd deepep-nvshmem && \ + NVSHMEM_SHMEM_SUPPORT=0 \ + NVSHMEM_UCX_SUPPORT=0 \ + NVSHMEM_USE_NCCL=0 \ + NVSHMEM_MPI_SUPPORT=0 \ + NVSHMEM_IBGDA_SUPPORT=1 \ + NVSHMEM_PMIX_SUPPORT=0 \ + NVSHMEM_TIMEOUT_DEVICE_POLLING=0 \ + NVSHMEM_USE_GDRCOPY=1 \ + cmake -G Ninja -S . -B build/ -DCMAKE_INSTALL_PREFIX=/workspace/deepep-nvshmem/install && cmake --build build/ --target install + +ENV NVSHMEM_DIR=/workspace/deepep-nvshmem/install +ENV LD_LIBRARY_PATH=$NVSHMEM_DIR/lib:$LD_LIBRARY_PATH +ENV PATH=$NVSHMEM_DIR/bin:$PATH + +## Build deepep +RUN cd DeepEP && \ + python setup.py install + +# Install Apex +RUN pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" git+https://github.com/NVIDIA/apex.git + +# Install TransformerEngine +RUN export NVTE_FRAMEWORK=pytorch && pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@v2.2.1 + +# Install Megatron-LM +RUN git clone -b core_v0.13.0 https://github.com/NVIDIA/Megatron-LM.git && \ + cd Megatron-LM && pip3 install --no-deps -e . + +# Install mbridge +RUN pip3 install --no-cache-dir git+https://github.com/ISEEKYAN/mbridge.git diff --git a/verl/docker/verl0.6-cu128-torch2.8.0-fa2.7.4/Dockerfile.vllm011.mcore_gpt-oss b/verl/docker/verl0.6-cu128-torch2.8.0-fa2.7.4/Dockerfile.vllm011.mcore_gpt-oss new file mode 100644 index 0000000000000000000000000000000000000000..4d57448686a1ae1270429fa0c02168682d36731f --- /dev/null +++ b/verl/docker/verl0.6-cu128-torch2.8.0-fa2.7.4/Dockerfile.vllm011.mcore_gpt-oss @@ -0,0 +1,15 @@ +FROM nvcr.io/nvidia/nemo:25.07.gpt_oss + +RUN git clone -b v0.11.0 --depth 1 https://github.com/vllm-project/vllm.git /opt/vllm + +RUN pip install setuptools_scm + +RUN cd /opt/vllm && pip install --no-deps --no-build-isolation --no-cache-dir -e . + +RUN pip install cbor2 setproctitle blake3 openai_harmony pybase64 msgspec partial_json_parser py-cpuinfo diskcache gguf + +RUN pip install --upgrade transformers tokenizers + +RUN pip install codetiming tensordict mathruler pylatexenc + +RUN pip3 install --no-cache-dir mbridge \ No newline at end of file diff --git a/verl/docker/verl0.6.1-experimental/Dockerfile.sglang056exp b/verl/docker/verl0.6.1-experimental/Dockerfile.sglang056exp new file mode 100644 index 0000000000000000000000000000000000000000..9d422183544d4659c7ffbfb1631bfce20757c5b7 --- /dev/null +++ b/verl/docker/verl0.6.1-experimental/Dockerfile.sglang056exp @@ -0,0 +1,63 @@ +# Dockerfile for verlai/verl:sgl056.exp +FROM lmsysorg/sglang:v0.5.6.post1 + +RUN pip install pybind11 + +RUN pip install nvidia-mathdx + +RUN pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" git+https://github.com/NVIDIA/apex.git + +RUN export NVTE_FRAMEWORK=pytorch && MAX_JOBS=128 NVTE_BUILD_THREADS_PER_JOB=4 pip3 install --resume-retries 999 --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.11 + +RUN pip install --upgrade --no-cache-dir transformers tokenizers + +RUN pip install codetiming tensordict mathruler pylatexenc qwen_vl_utils + +RUN pip install --no-cache-dir --no-build-isolation flash_attn==2.8.1 + +RUN NSIGHT_VERSION=2025.6.1_2025.6.1.190-1_$(if [ "$(uname -m)" = "aarch64" ]; then echo "arm64"; else echo "amd64"; fi) && \ + wget https://developer.nvidia.com/downloads/assets/tools/secure/nsight-systems/2025_6/nsight-systems-${NSIGHT_VERSION}.deb && \ + apt-get update && apt-get install -y libxcb-cursor0 && \ + apt-get install -y ./nsight-systems-${NSIGHT_VERSION}.deb && \ + rm -rf /usr/local/cuda/bin/nsys && \ + ln -s /opt/nvidia/nsight-systems/2025.6.1/nsys /usr/local/cuda/bin/nsys && \ + rm -rf /usr/local/cuda/bin/nsys-ui && \ + ln -s /opt/nvidia/nsight-systems/2025.6.1/nsys-ui /usr/local/cuda/bin/nsys-ui && \ + rm nsight-systems-${NSIGHT_VERSION}.deb + + +# ========================= +# Install HybridEP +# ========================= +WORKDIR /home/ +RUN git clone --branch hybrid-ep https://github.com/deepseek-ai/DeepEP.git && \ + cd DeepEP && git checkout 3f601f7ac1c062c46502646ff04c535013bfca00 && \ + TORCH_CUDA_ARCH_LIST="9.0;10.0" pip install --no-build-isolation . + +# ========================= +# Install Qwen3-Next dependencies +# ========================= +WORKDIR /home/ +# Install causal-conv1d and flash-linear-attention +RUN cd /tmp && \ + git clone https://github.com/Dao-AILab/causal-conv1d.git && \ + cd causal-conv1d && \ + unset PIP_CONSTRAINT && \ + CAUSAL_CONV1D_FORCE_BUILD=TRUE pip install --no-build-isolation . && \ + cd .. && \ + rm -rf causal-conv1d && \ + pip install flash-linear-attention + +RUN pip install --no-cache-dir torch-memory-saver + +RUN pip3 install --no-cache-dir --no-deps trl + +RUN pip3 install nvtx matplotlib liger_kernel + +RUN pip install -U git+https://github.com/ISEEKYAN/mbridge.git + +RUN pip install --no-deps --no-cache-dir git+https://github.com/NVIDIA/Megatron-LM.git@1d462bd37dac21cfa14177405d4921eedb987052 # latest dev branch on 20251209 + +RUN pip install git+https://github.com/verl-project/verl.git@v0.6.1 + +RUN pip uninstall -y verl \ No newline at end of file diff --git a/verl/docker/verl0.6.1-experimental/Dockerfile.vllm012exp b/verl/docker/verl0.6.1-experimental/Dockerfile.vllm012exp new file mode 100644 index 0000000000000000000000000000000000000000..832a645627365ff9ea4cfc8b3a8f6566a4db377e --- /dev/null +++ b/verl/docker/verl0.6.1-experimental/Dockerfile.vllm012exp @@ -0,0 +1,69 @@ +# dockerfile for verlai/verl:vll012.exp +FROM nvcr.io/nvidia/pytorch:25.11-py3 + +RUN git clone -b v0.12.0 --depth 1 https://github.com/vllm-project/vllm.git /opt/vllm + +RUN pip install setuptools_scm + +RUN cd /opt/vllm && pip install --no-deps --no-build-isolation --no-cache-dir -e . + +RUN pip install -r /opt/vllm/requirements/common.txt + + +RUN pip install pybind11 + +RUN export NVTE_FRAMEWORK=pytorch && MAX_JOBS=128 NVTE_BUILD_THREADS_PER_JOB=4 pip3 install --resume-retries 999 --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.11 + +RUN pip install --upgrade --no-cache-dir transformers tokenizers + +RUN pip install codetiming tensordict mathruler pylatexenc qwen_vl_utils + +RUN pip install flash_attn +#==2.8.1 + +RUN apt update && apt install numactl + +RUN NSIGHT_VERSION=2025.6.1_2025.6.1.190-1_$(if [ "$(uname -m)" = "aarch64" ]; then echo "arm64"; else echo "amd64"; fi) && \ + wget https://developer.nvidia.com/downloads/assets/tools/secure/nsight-systems/2025_6/nsight-systems-${NSIGHT_VERSION}.deb && \ + apt-get update && apt-get install -y libxcb-cursor0 && \ + apt-get install -y ./nsight-systems-${NSIGHT_VERSION}.deb && \ + rm -rf /usr/local/cuda/bin/nsys && \ + ln -s /opt/nvidia/nsight-systems/2025.6.1/nsys /usr/local/cuda/bin/nsys && \ + rm -rf /usr/local/cuda/bin/nsys-ui && \ + ln -s /opt/nvidia/nsight-systems/2025.6.1/nsys-ui /usr/local/cuda/bin/nsys-ui && \ + rm nsight-systems-${NSIGHT_VERSION}.deb + + +# ========================= +# Install HybridEP +# ========================= +WORKDIR /home/ +RUN git clone --branch hybrid-ep https://github.com/deepseek-ai/DeepEP.git && \ + cd DeepEP && git checkout 3f601f7ac1c062c46502646ff04c535013bfca00 && \ + TORCH_CUDA_ARCH_LIST="9.0;10.0" pip install --no-build-isolation . + +# ========================= +# Install Qwen3-Next dependencies +# ========================= +WORKDIR /home/ +# Install causal-conv1d and flash-linear-attention +RUN cd /tmp && \ + git clone https://github.com/Dao-AILab/causal-conv1d.git && \ + cd causal-conv1d && \ + unset PIP_CONSTRAINT && \ + CAUSAL_CONV1D_FORCE_BUILD=TRUE pip install --no-build-isolation . && \ + cd .. && \ + rm -rf causal-conv1d && \ + pip install flash-linear-attention + +RUN pip3 install --no-cache-dir --no-deps trl + +RUN pip3 install nvtx matplotlib liger_kernel + +RUN pip install -U git+https://github.com/ISEEKYAN/mbridge.git + +RUN pip install --no-deps --no-cache-dir git+https://github.com/NVIDIA/Megatron-LM.git@1d462bd37dac21cfa14177405d4921eedb987052 # latest dev branch on 20251209 + +RUN pip install git+https://github.com/verl-project/verl.git@v0.6.1 + +RUN pip uninstall -y verl \ No newline at end of file diff --git a/verl/docs/Makefile b/verl/docs/Makefile new file mode 100644 index 0000000000000000000000000000000000000000..8bda904a9b0b29dfcf538cb52b806dd910710a4a --- /dev/null +++ b/verl/docs/Makefile @@ -0,0 +1,20 @@ +# Minimal makefile for Sphinx documentation +# + +# You can set these variables from the command line. +SPHINXOPTS = +SPHINXBUILD = sphinx-build +SPHINXPROJ = verl +SOURCEDIR = . +BUILDDIR = _build + +# Put it first so that "make" without argument is like "make help". +help: + @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + +.PHONY: help Makefile + +# Catch-all target: route all unknown targets to Sphinx using the new +# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). +%: Makefile + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) diff --git a/verl/docs/README.md b/verl/docs/README.md new file mode 100644 index 0000000000000000000000000000000000000000..8c5db04874138435ef986342a7b8be668b81d0b0 --- /dev/null +++ b/verl/docs/README.md @@ -0,0 +1,22 @@ +# verl documentations + +## Build the docs + +```bash +# If you want to view auto-generated API docstring, please make sure verl is available in python path. For instance, install verl via: +# pip install .. -e[test] + +# Install dependencies needed for building docs. +pip install -r requirements-docs.txt + +# Build the docs. +make clean +make html +``` + +## Open the docs with your browser + +```bash +python -m http.server -d _build/html/ +``` +Launch your browser and navigate to http://localhost:8000 to view the documentation. Alternatively you could drag the file `_build/html/index.html` to your local browser and view directly. diff --git a/verl/docs/README_vllm0.7.md b/verl/docs/README_vllm0.7.md new file mode 100644 index 0000000000000000000000000000000000000000..ffef3200767c32d0a714b96af4722065fb016fd5 --- /dev/null +++ b/verl/docs/README_vllm0.7.md @@ -0,0 +1,73 @@ +# Upgrading to vllm >= 0.7 + +Note: verl+vllm 0.8.3 is now stable. Please see ``docs/README_vllm0.8.md`` for upgrade guide. + +## Installation + +Note: At time of writing, verl+vllm 0.7.x supports **FSDP** for training and **vLLM** for rollout. + +``` +# Create the conda environment +conda create -n verl python==3.10 +conda activate verl + +# Install verl +git clone https://github.com/verl-project/verl.git +cd verl +pip3 install -e . + +# Install the latest stable version of vLLM +pip3 install vllm==0.7.3 + +# Install flash-attn +pip3 install flash-attn --no-build-isolation + +``` + +Note that if you are installing lower versions of vLLM (0.7.0, 0.7.1, 0.7.2), you need to make some tiny patches manually on vllm (/path/to/site-packages/vllm after installation) after the above steps: + +- vllm/distributed/parallel_state.py: Remove the assertion below: + +``` +if (world_size + != tensor_model_parallel_size * pipeline_model_parallel_size): + raise RuntimeError( + f"world_size ({world_size}) is not equal to " + f"tensor_model_parallel_size ({tensor_model_parallel_size}) x " + f"pipeline_model_parallel_size ({pipeline_model_parallel_size})") + +``` + +- vllm/executor/uniproc_executor.py: change `local_rank = rank` to `local_rank = int(os.environ["LOCAL_RANK"])` +- vllm/model_executor/model_loader/weight_utils.py: remove the `torch.cuda.empty_cache()` in `pt_weights_iterator` + +## Features + +### Use cuda graph + +After installation, examples using FSDP as training backends can be used. By default, the `enforce_eager` is set to True, which disables the cuda graph. To enjoy cuda graphs and the sleep mode of vLLM>=0.7, add the following lines to the bash script: + +``` +actor_rollout_ref.rollout.enforce_eager=False \ +actor_rollout_ref.rollout.free_cache_engine=True \ + +``` + +For a typical job like examples/ppo_trainer/run_qwen3_8b_fsdp.sh, the rollout generation time is 85 seconds with vLLM0.7.0. By enabling the cudagraph, the generation duration is further reduced to 62 seconds. + +**Note:** Currently, if the `n` is greater than 1 in `SamplingParams` in vLLM>=0.7, there is a potential performance issue on the stability of rollout generation time (Some iterations would see generation time bursts) using vLLM's V0 Engine. + +### Use vLLM V1 Engine + +Using the vLLM V1 engine can avoid instability issues and achieve additional performance improvements. To use the V1 engine, you can first uninstall the previously installed vLLM and then follow the steps below to install the newer version. + +``` +git clone https://github.com/vllm-project/vllm.git +cd vllm +git checkout 2275784 +sed -i "903a\ data_parallel_size = world_size // pipeline_model_parallel_size // tensor_model_parallel_size" ./vllm/distributed/parallel_state.py +VLLM_USE_PRECOMPILED=1 pip install --editable . +``` + +Then you can enable the V1 engine by setting `export VLLM_USE_V1=1`. In some benchmark tests, the V1 engine demonstrates a 1.5x speed improvement over the vLLM V0 engine. +The stable support of the vLLM V1 engine is available on verl main. diff --git a/verl/docs/README_vllm0.8.md b/verl/docs/README_vllm0.8.md new file mode 100644 index 0000000000000000000000000000000000000000..3641b30155853e3cfa1c51e83d0dc94b224dd50b --- /dev/null +++ b/verl/docs/README_vllm0.8.md @@ -0,0 +1,52 @@ +# Upgrading to vLLM >= 0.8 + +Last updated: 05/04/2025. + +## Installation + +Note: This version of verl+vLLM 0.8+ supports **FSDP** for training and **vLLM** for rollout. + +```bash +# Create the conda environment +conda create -n verl python==3.10 +conda activate verl + +# Install verl +git clone https://github.com/verl-project/verl.git +cd verl +pip3 install -e . + +# Install the latest stable version of vLLM +pip3 install vllm==0.8.3 + +# Install flash-attn +pip3 install flash-attn --no-build-isolation + +``` + +We have a pre-built docker image for verl+vLLM 0.8.3. You can direct import it with the following command: + +```bash +docker pull hiyouga/verl:ngc-th2.6.0-cu126-vllm0.8.3-flashinfer0.2.2-cxx11abi0 +``` + +## Features + +vLLM 0.8+ supports cuda graph and V1 engine by default in verl. To enable these features, remember to add the following lines to the bash script: + +```bash +actor_rollout_ref.rollout.enforce_eager=False \ +actor_rollout_ref.rollout.free_cache_engine=True \ +``` + +and also **remove** the environment variable if it exists: + +## Notes + +When you just directly upgrade vllm>=0.8, some dependency packages may undergo version changes. If you encounter the following problems: + +```bash +in from torch.multiprocessing.reductions import ForkingPickler ImportError: cannot import name 'ForkingPickler' from 'torch.multiprocessing.reductions' (/opt/conda/lib/python3.11/site-packages/torch/multiprocessing/reductions.py) +``` + +You need to upgrade `tensordict` to version 0.6.2 using the command `pip install tensordict==0.6.2`. diff --git a/verl/docs/_static/custom.css b/verl/docs/_static/custom.css new file mode 100644 index 0000000000000000000000000000000000000000..32f08475754bc280bca407d1643ec3aa68eeacf3 --- /dev/null +++ b/verl/docs/_static/custom.css @@ -0,0 +1,217 @@ +/* Make the documentation use full screen width */ +.wy-nav-content { + max-width: none !important; + width: 100% !important; + padding: 1.618em 3.236em !important; +} + +/* Adjust the content wrapper - will be set by JavaScript */ +.wy-nav-content-wrap { + margin-left: 300px; + transition: margin-left 0.2s ease; + width: auto !important; + position: relative !important; + background: white !important; + min-height: 100vh !important; +} + +/* Make the main content area responsive */ +.rst-content { + max-width: none !important; + width: 100% !important; +} + +/* Optional: Adjust table widths to prevent overflow */ +.rst-content table.docutils { + width: 100% !important; + table-layout: auto !important; +} + +/* Optional: Better code block width handling */ +.rst-content .highlight { + width: 100% !important; +} + +/* Content area positioning already handled above */ + +/* Optional: Improve readability with some margin on very wide screens */ +@media (min-width: 1400px) { + .wy-nav-content { + max-width: none !important; + margin: 0 auto !important; + } +} + +/* Resizable sidebar styles */ +.wy-nav-side { + position: fixed !important; + top: 0 !important; + bottom: 0 !important; + left: 0 !important; + width: 300px; + min-width: 200px; + max-width: 600px; + display: flex; + flex-direction: column; + z-index: 200 !important; +} + +/* Ensure sidebar header (logo, search) adapts to width */ +.wy-side-nav-search { + width: 100% !important; + box-sizing: border-box !important; + padding: 0.809em 0.809em !important; +} + +.wy-side-nav-search input[type="text"] { + width: 100% !important; + box-sizing: border-box !important; +} + +/* Make logo/title area responsive */ +.wy-side-nav-search > div.version { + width: 100% !important; +} + +.wy-side-nav-search > a { + width: 100% !important; + display: block !important; + white-space: nowrap !important; + overflow: hidden !important; + text-overflow: ellipsis !important; +} + +/* Responsive adjustments for narrow sidebar */ +@media (max-width: 300px) { + .wy-side-nav-search > a { + font-size: 0.9em !important; + } + + .wy-side-nav-search input[type="text"] { + font-size: 0.8em !important; + } +} + +/* Ensure search input doesn't overflow */ +.wy-side-nav-search form { + width: 100% !important; + margin: 0 !important; +} + +/* Make search icon responsive */ +.wy-side-nav-search .wy-dropdown { + width: 100% !important; +} + +/* Adjust search results dropdown width */ +.wy-side-nav-search .wy-dropdown-menu { + width: 100% !important; + max-width: none !important; + left: 0 !important; + right: 0 !important; +} + +/* Resize handle is created by JavaScript */ + +/* Make sure the sidebar content doesn't overflow */ +.wy-side-scroll { + width: 100% !important; + flex: 1 !important; + overflow-y: auto !important; + overflow-x: hidden !important; + padding-right: 10px !important; + box-sizing: border-box !important; + scroll-behavior: auto !important; /* Prevent smooth scrolling on sidebar itself */ +} + +/* Ensure proper scroll behavior for main content area */ +html { + scroll-behavior: smooth !important; +} + +/* Ensure anchor links work properly in main content */ +.wy-nav-content-wrap { + scroll-behavior: smooth !important; +} + +/* Fix scroll to target for anchor links */ +.rst-content { + scroll-behavior: smooth !important; +} + +/* Fix anchor scroll offset to account for fixed header */ +.rst-content .section { + scroll-margin-top: 60px; +} + +/* Fix anchor scroll offset for headers */ +.rst-content h1, .rst-content h2, .rst-content h3, .rst-content h4, .rst-content h5, .rst-content h6 { + scroll-margin-top: 60px; +} + +/* Fix anchor scroll offset for specific scroll targets */ +.rst-content .headerlink { + scroll-margin-top: 60px; +} + +/* Fix sidebar navigation styling */ +.wy-menu-vertical { + width: 100% !important; +} + +.wy-menu-vertical li { + width: 100% !important; +} + +.wy-menu-vertical a { + width: 100% !important; + word-wrap: break-word !important; + white-space: normal !important; +} + +/* Content area margin is handled by JavaScript */ + +/* Custom drag handle (more visible) */ +.resize-handle { + position: absolute; + top: 0; + right: 0; + width: 8px; + height: 100%; + background: #ccc; + cursor: col-resize; + z-index: 1001; + opacity: 0.3; + transition: opacity 0.2s ease; +} + +.resize-handle:hover { + opacity: 0.8; + background: #999; +} + +.resize-handle::before { + content: ''; + position: absolute; + top: 50%; + left: 50%; + width: 2px; + height: 20px; + background: #666; + transform: translate(-50%, -50%); + border-radius: 1px; +} + +.resize-handle:hover::before { + background: #333; +} + +/* Ensure smooth resizing */ +.wy-nav-side.resizing { + user-select: none; + pointer-events: none; +} + +.wy-nav-side.resizing .wy-side-scroll { + overflow: hidden; +} \ No newline at end of file diff --git a/verl/docs/_static/js/resizable-sidebar.js b/verl/docs/_static/js/resizable-sidebar.js new file mode 100644 index 0000000000000000000000000000000000000000..2a51fa90043bb0ecf78149b092fd3447740fdaee --- /dev/null +++ b/verl/docs/_static/js/resizable-sidebar.js @@ -0,0 +1,251 @@ +// Resizable sidebar functionality +document.addEventListener('DOMContentLoaded', function() { + const sidebar = document.querySelector('.wy-nav-side'); + const content = document.querySelector('.wy-nav-content-wrap'); + + if (!sidebar || !content) return; + + // Create resize handle + const resizeHandle = document.createElement('div'); + resizeHandle.className = 'resize-handle'; + sidebar.appendChild(resizeHandle); + + let isResizing = false; + let startX = 0; + let startWidth = 0; + + // Get initial width + const getInitialWidth = () => { + return 300; // Default width + }; + + // Save width to localStorage + const saveWidth = (width) => { + localStorage.setItem('sidebar-width', width); + }; + + // Load width from localStorage + const loadWidth = () => { + const savedWidth = localStorage.getItem('sidebar-width'); + if (savedWidth) { + const width = parseInt(savedWidth, 10); + if (width >= 200 && width <= 600) { + return width; + } + } + return getInitialWidth(); + }; + + // Apply width to sidebar and content + const applyWidth = (width) => { + // Update sidebar width + sidebar.style.width = width + 'px'; + + // Update content margin with !important to override any CSS + content.style.setProperty('margin-left', width + 'px', 'important'); + + // Also update any other content wrapper that might exist + const contentInner = document.querySelector('.wy-nav-content'); + if (contentInner) { + contentInner.style.setProperty('margin-left', '0px', 'important'); + } + + // Force reflow and repaint + sidebar.offsetHeight; + content.offsetHeight; + + // Trigger window resize event to notify other components + window.dispatchEvent(new Event('resize')); + }; + + // Initialize with saved width + const initialWidth = loadWidth(); + applyWidth(initialWidth); + + // Mouse down on resize handle + resizeHandle.addEventListener('mousedown', (e) => { + isResizing = true; + startX = e.clientX; + startWidth = parseInt(window.getComputedStyle(sidebar).width, 10); + + sidebar.classList.add('resizing'); + document.body.style.cursor = 'col-resize'; + document.body.style.userSelect = 'none'; + + // Add overlay to prevent iframe issues + const overlay = document.createElement('div'); + overlay.style.cssText = ` + position: fixed; + top: 0; + left: 0; + width: 100%; + height: 100%; + z-index: 9999; + cursor: col-resize; + `; + overlay.id = 'resize-overlay'; + document.body.appendChild(overlay); + + e.preventDefault(); + }); + + // Mouse move + document.addEventListener('mousemove', (e) => { + if (!isResizing) return; + + const width = startWidth + e.clientX - startX; + const clampedWidth = Math.max(200, Math.min(600, width)); + applyWidth(clampedWidth); + }); + + // Mouse up + document.addEventListener('mouseup', () => { + if (!isResizing) return; + + isResizing = false; + sidebar.classList.remove('resizing'); + document.body.style.cursor = ''; + document.body.style.userSelect = ''; + + // Remove overlay + const overlay = document.getElementById('resize-overlay'); + if (overlay) { + overlay.remove(); + } + + // Save the current width + const currentWidth = parseInt(window.getComputedStyle(sidebar).width, 10); + saveWidth(currentWidth); + }); + + // Handle window resize - removed to prevent infinite loop + // The sidebar width is fixed and managed by drag functionality, no need to recalculate on window resize + + // Double-click to reset to default width + resizeHandle.addEventListener('dblclick', () => { + const defaultWidth = 300; + applyWidth(defaultWidth); + saveWidth(defaultWidth); + }); +}); + +// Fix navigation issues - Using MutationObserver for reliable initialization +document.addEventListener('DOMContentLoaded', function() { + let navigationFixed = false; + + function setupNavigationFix() { + if (navigationFixed) return; + + // Find all links in the sidebar + const sidebarLinks = document.querySelectorAll('.wy-menu-vertical a'); + + // Only proceed if we have sidebar links + if (sidebarLinks.length === 0) return; + + console.log('Setting up navigation fix...'); + + sidebarLinks.forEach(function(link) { + const href = link.getAttribute('href'); + + // Clone the link to remove all existing event listeners + const newLink = link.cloneNode(true); + + // Add our own click handler + newLink.addEventListener('click', function(e) { + console.log('Link clicked:', href); + + // If it's an anchor link within the same page + if (href && href.startsWith('#') && href !== '#') { + e.preventDefault(); + e.stopPropagation(); + + const targetId = href.substring(1); + const targetElement = document.getElementById(targetId); + + if (targetElement) { + // Calculate offset for fixed header + const headerHeight = 60; + const elementPosition = targetElement.getBoundingClientRect().top; + const offsetPosition = elementPosition + window.pageYOffset - headerHeight; + + window.scrollTo({ + top: offsetPosition, + behavior: 'smooth' + }); + + // Update URL hash + if (history.pushState) { + history.pushState(null, null, '#' + targetId); + } else { + location.hash = '#' + targetId; + } + } + } + // For external links, navigate normally + else if (href && !href.startsWith('#') && !href.startsWith('javascript:')) { + console.log('Navigating to external link:', href); + window.location.href = href; + } + }); + + // Replace the old link with the new one + link.parentNode.replaceChild(newLink, link); + }); + + navigationFixed = true; + + // Handle initial page load with hash + if (window.location.hash) { + // Use requestAnimationFrame for better timing + requestAnimationFrame(() => { + const targetId = window.location.hash.substring(1); + const targetElement = document.getElementById(targetId); + if (targetElement) { + const headerHeight = 60; + const elementPosition = targetElement.getBoundingClientRect().top; + const offsetPosition = elementPosition + window.pageYOffset - headerHeight; + + window.scrollTo({ + top: offsetPosition, + behavior: 'smooth' + }); + } + }); + } + } + + // Try to set up navigation fix immediately + setupNavigationFix(); + + // If it didn't work, use MutationObserver to watch for when sidebar links are added + if (!navigationFixed) { + const observer = new MutationObserver(function(mutations) { + mutations.forEach(function(mutation) { + if (mutation.type === 'childList' && mutation.addedNodes.length > 0) { + // Check if sidebar links were added + const sidebarLinks = document.querySelectorAll('.wy-menu-vertical a'); + if (sidebarLinks.length > 0) { + setupNavigationFix(); + if (navigationFixed) { + observer.disconnect(); + } + } + } + }); + }); + + // Start observing the document for changes + observer.observe(document.body, { + childList: true, + subtree: true + }); + + // Fallback timeout in case MutationObserver doesn't work + setTimeout(function() { + if (!navigationFixed) { + setupNavigationFix(); + } + observer.disconnect(); + }, 5000); + } +}); \ No newline at end of file diff --git a/verl/docs/_static/js/runllm-widget.js b/verl/docs/_static/js/runllm-widget.js new file mode 100644 index 0000000000000000000000000000000000000000..bec345cacc5b943693e1bf1973a7a6d863b0d85e --- /dev/null +++ b/verl/docs/_static/js/runllm-widget.js @@ -0,0 +1,14 @@ +document.addEventListener("DOMContentLoaded", function () { + var script = document.createElement("script"); + script.type = "module"; + script.id = "runllm-widget-script"; + script.src = "https://widget.runllm.com"; + script.setAttribute("version", "stable"); + script.setAttribute("crossorigin", "true"); + script.setAttribute("runllm-keyboard-shortcut", "Mod+j"); + script.setAttribute("runllm-name", "verl Chatbot"); + script.setAttribute("runllm-position", "TOP_RIGHT"); + script.setAttribute("runllm-assistant-id", "679"); + script.async = true; + document.head.appendChild(script); + }); \ No newline at end of file diff --git a/verl/docs/_static/logo.png b/verl/docs/_static/logo.png new file mode 100644 index 0000000000000000000000000000000000000000..6a3b61308d73eb72071eaec8c95544ab36cd3970 Binary files /dev/null and b/verl/docs/_static/logo.png differ diff --git a/verl/docs/advance/agent_loop.rst b/verl/docs/advance/agent_loop.rst new file mode 100644 index 0000000000000000000000000000000000000000..d5b41ff7cfdbe6864e02fbac2b1753f0fbef4271 --- /dev/null +++ b/verl/docs/advance/agent_loop.rst @@ -0,0 +1,238 @@ +Agent Loop +========== + +Last updated: 07/17/2025. + +.. versionadded:: 0.4.2 + [status: alpha] + +.. warning:: + Agent Loop is ready for use, but the API may change in future releaes. + +Agent Loop is designed as general interface for multi-turn rollout and agentic reinforcement learning. + +**Design goal**: + +- Plugable user defined agent loop +- Provide standard request generate api with different inference frameworks +- Provide request level load balance between multiple inference servers + +**Non-goal**: + +- How tool is defined and how to call tool + +In high level overview, agent loop is given a prompt, run user defined loop: call LLM generate api, call tools, ... +and return the final output. The final output is then calculated reward and used as trajectory for RL training. + +.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/agent_loop_overview.svg?raw=true + + +API Design +---------- + +``AgentLoopBase`` class is the abstraction of agent loop, and ``run`` method is the only interface that user need to implement. +The run method, given prompt messages in format: [{"role": "user"}, {"content": "..."}], and additional sampling params, +could do whatever user wants, such as + +- call LLM generate api +- call tools: web search, database query, code sandbox, ... +- environment interaction +- reflection +- ... + +.. code:: python + + class AgentLoopBase(ABC): + @abstractmethod + async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentLoopOutput: + """Run agent loop to interact with LLM server and environment. + + Args: + sampling_params (Dict[str, Any]): LLM sampling params. + **kwargs: dataset fields from `verl.utils.dataset.RLHFDataset`. + + Returns: + AgentLoopOutput: Agent loop output. + """ + raise NotImplementedError + +After running user defined loop, run method should return ``AgentLoopOutput``, including prompt token ids, +response token ids, and response mask. + +.. code:: python + + class AgentLoopOutput(BaseModel): + """Agent loop output.""" + + prompt_ids: list[int] + """Prompt token ids.""" + response_ids: list[int] + """Response token ids including LLM generated token, tool response token.""" + response_mask: list[int] + """Response mask, 1 for LLM generated token, 0 for tool response token.""" + +.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/agent_loop_output.svg?raw=true + +.. note:: AgentLoopOutput only output one trajectory for a given prompt, multiple trajectories output is still under discussion. + +Architecture Design +------------------- + +.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/agent_loop_architecture.png?raw=true + +A single PPO step contain two phase: rollout and train. In rollout phase: + +1. PPOTrainer sample a batch from dataset and call ``AgentLoopManager.generate_sequences``. +2. AgentLoopManager ``wake_up`` all async LLM server instances, which will sync weights between inference engine(vLLM/SGLang) and training engine(FSDP/Megatron-LM). +3. AgentLoopManager split batch into chunks and send each chunk to ``AgentLoopWorker``. +4. AgentLoopWorker receive chunk and for each prompt, spawn a user defined ``AgentLoopBase`` instance, run ``run`` coroutine until end and get ``AgentLoopOutput``. + +.. tip:: + AgentLoopWorker schedules multiple coroutines concurrently. If number of AgentLoopWorker equals batch_size, then each worker is response for one prompt. + +In agent loop, when user need LLM generate response: + +5. Call ``LLMServerClient.generate`` with prompt_ids. +6. LLMServerClient select a server instance with least request in first turn and send request to it. (In following turns, the request will be sent to the same server instance). +7. AsyncLLMServer receive a request, issue ipc/rpc with model_runner, and generate response. (There's slight differences between vLLM and SGLang, see below). + +When all prompts in all AgentLoopWorker finish, AgentLoopManager gather results and return to PPOTrainer. + +8. AgentLoopManager ``sleep`` all server instances, which will free kv cache and offload weights to CPU memory. + +AsyncLLMServer +~~~~~~~~~~~~~~ + +AsyncLLMServer is the abstraction of LLM server with two types of generation api: + +- `OpenAI chat completion `_: generate response for the given chat conversation. +- Token in token out: generate response ids for the given token ids. + +We have officially supported vLLM and SGLang AsyncLLMServer, both of them implement the two api and are well tested. +Other inference engine should be easy to plug-in by implement the ``AsyncServerBase`` class. + +.. code:: python + + class AsyncServerBase(ABC): + @abstractmethod + async def chat_completion(self, raw_request: Request) -> JSONResponse: + """OpenAI chat completion API. + + Args: + raw_request (Request): raw json request + + Returns: + JSONResponse: json response + + API reference: https://platform.openai.com/docs/api-reference/chat/create + """ + raise NotImplementedError + + @abstractmethod + async def generate(self, prompt_ids: list[int], sampling_params: dict[str, Any], request_id: str) -> list[int]: + """Generate response ids given prompt ids. + + Args: + prompt_ids (List[int]): prompt ids + sampling_params (Dict[str, Any]): sampling params + request_id (str): request id + + Returns: + List[int]: response ids + """ + raise NotImplementedError + + +Chat completion vs Token in token out +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. warning:: + The following conclusion is based on our recent experience and is still open to investigation and discussion. + +Almost all agent frameworks (LangGraph, CrewAI, LlamaIndex, etc) call LLM with OpenAI chat completion api, and +keep chat history as messages. So user may expect that we should use the chat completion api in multi-turn rollout. + +But based on our recent experience on single-turn training on DAPO and multi-turn training on `retool `_, +we found the token_ids from apply the final messages may not equal to the token_ids by concat prompt_ids and response_ids in each turn. + +.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/multi_turn.png?raw=true + +**Where does this inconsistency happened?** + +First, the tool parser may alter the content. For example + +.. code:: json + + {"role": "assistant", "content": "Let me call a ... and get the result"} + +After tool_calls extraction, the messages is like this: + +.. code:: json + + {"role": "assistant", "content": "Let me call a and get the result", "tool_calls": [{"name": "foo", "arguments": "{}"}]} + +Encode the extracted message back is not equal to the original LLM generated response_ids. + +Second, the `decode-encode` may also lead to inconsistency: `Agent-R1 issue#30 `_. + +**What is the impact of this inconsistency?** + +This inconsistency is not a big problem for serving/agent system, but is critical to RL training. +It causes the trajectory deviate from the policy model distribution. We have observed that apply_chat_template +to the final chat history messages make PPO training not even converged in single-turn. + +vLLM +^^^^ + +.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/async_vllm.png?raw=true + +For vLLM, the Async LLM Engine is running in same process as the server, and ModelRunner is running in same process as FSDP/Megatron-LM workers. +Async LLM Engine communicate with ModelRunner through ZeroMQ. When server receive a request, it directly call engine to generate response_ids. + +SGLang +^^^^^^ + +.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/async_sglang.png?raw=true + +For SGLang, the Async LLM Engine is running in same process as FSDP/Megatron-LM worker-0, and it spawn multiple subprocesses as ModelRunner. +Also, Async LLM Engine communicate with ModelRunner through ZeroMQ. When server receive a request, it remote call the worker-0 and get response_ids. + +LLMServerClient +~~~~~~~~~~~~~~~~~~~~~ + +LLMServerClient serve as proxy to multiple AsyncLLMServer instances, provides: + +- load balance: select a server instance with least request in first turn and send request to it. +- sticky session: bind request_id to server instance, so that the same request_id will be sent to the same server instance in following turns. + +LLMServerClient is passed to ``AgentLoopBase.__init__``, whenever user want to interact with LLM in agent loop, +they can call ``LLMServerClient.generate`` to generate response_ids. + +.. code:: python + + class LLMServerClient: + async def generate( + self, + request_id, + *, + prompt_ids: list[int], + sampling_params: dict[str, Any], + ) -> list[int]: + """Generate tokens from prompt ids. + + Args: + request_id (str): request id for sticky session. + prompt_ids (List[int]): List of prompt token ids. + sampling_params (Dict[str, Any]): Sampling parameters for the chat completion. + + Returns: + List[int]: List of generated token ids. + """ + ... + +Next +---- + +- :doc:`Agentic RL Training<../start/agentic_rl>`: Quick start agentic RL training with gsm8k dataset. +- `LangGraph MathExpression `_: Demonstrate how to use LangGraph to build agent loop. +- `Retool `_: End-to-end retool paper reproduction using tool agent. diff --git a/verl/docs/advance/async-on-policy-distill.md b/verl/docs/advance/async-on-policy-distill.md new file mode 100644 index 0000000000000000000000000000000000000000..55b8d392206c94968d6ade5a29ce82eb8d267c8f --- /dev/null +++ b/verl/docs/advance/async-on-policy-distill.md @@ -0,0 +1,242 @@ +# Recipe: Async On-Policy Knowledge Distillation Trainer + +**Authors:** Brilliant Hanabi, furunding + +**Last updated:** 2025-11-08 + +## 1. Background + +On-policy knowledge distillation (KD) trains a student policy to imitate a stronger teacher using samples drawn from the student's current policy. For each on-policy rollout the teacher returns soft, top-k token distributions and the student is optimized with a token-wise sparse KL objective that focuses learning on the teacher's high-probability modes. Because training examples come from the student's own state distribution, KD reduces distributional mismatch relative to off-policy distillation or supervised fine-tuning (SFT), improving stability and sample efficiency. Compared with reinforcement learning, KD avoids high-variance reward-based optimization and complex reward design by providing dense, informative per-token targets, which typically yields faster convergence and simpler scaling. Recent empirical and implementation-focused writeups (e.g., [ThinkingMachines' blog on on-policy distillation](https://thinkingmachines.ai/blog/on-policy-distillation/)) also demonstrate that on-policy distillation can deliver high-quality behavior with substantially lower compute and data requirements than many alternative approaches. + +Built on verl’s Ray-based single-controller components, we initially assembled a strictly on-policy KD pipeline where rollout generation, teacher knowledge acquisition, and policy optimization ran in lockstep. In practice, this synchronous design proved highly inefficient: the three stages had to wait for one another, creating pipeline bubbles and underutilized GPUs. To address this, we extend the asynchronous schedulers introduced by the One-Step-Off Policy pipeline to overlap these phases. This overlap preserves the same distillation objective while trading some strict on-policy guarantees for substantial gains in end-to-end throughput and hardware utilization. + +## 2. Distillation Overview and Objective + +This recipe centers on on-policy knowledge distillation: the student policy learns from a stronger teacher on samples generated by the current policy (on-policy). For each input prompt, the student (actor) generates responses; the teacher provides top-k token distributions, and the student is trained to match them token-wise. + +Core components: + +1. Teacher signal: top-k log-probabilities and token indices per valid token position. +2. Student objective: sparse, token-level KL divergence between student logits and teacher top-k distribution. + +Objective: encourage student probabilities $Q$ to cover teacher modes $P$ using token-wise $\mathrm{KL}(P\,\|\,Q)$ computed on the teacher's top-k support. + +## 3. Efficient System Design + +### 3.1 Schedulers (One-Step / Two-Step Off-Policy) + +The native (serial) on-policy distillation process is shown in the figure below. + +![Zero-Step-Off Scheduler](https://raw.githubusercontent.com/eric-haibin-lin/verl-community/refs/heads/main/docs/zero-step-off-distill.png) + +This recipe supports optional schedulers that overlap generation, teacher querying, and updates to improve throughput without changing the distillation objective. + +#### 3.1.1 One-Step-Off-Policy + +![One-Step-Off Scheduler](https://raw.githubusercontent.com/eric-haibin-lin/verl-community/refs/heads/main/docs/one-step-off-distill.png) + +- Warm-up: 2 steps. +- Overlap pattern: rollout while actor update; weight sync while teacher retrieving. +- Timing keys: `sync_rollout_weights`, `wait_prev_gen`, `wait_prev_teacher`. + +#### 3.1.2 Two-Step-Off-Policy + +![Two-Step-Off Scheduler](https://raw.githubusercontent.com/eric-haibin-lin/verl-community/refs/heads/main/docs/two-step-off-distill.png) + +- Warm-up: 3 steps. +- Overlap pattern: rollout, actor update while teacher retrieving; interleave weight sync. +- Timing keys: `sync_rollout_weights`, `max(wait_prev_gen, wait_prev_prev_teacher)`. + +Tip: Use `two_step_off` when teacher takes much more time than sync; `one_step_off` for simpler overlapping. + +Practical details: + +- Inputs per batch: `teacher_topk_logps`, `teacher_topk_indices`, `attention_mask` (to select valid token positions). +- Loss injection: last pipeline stage computes KL via a logits processor; earlier stages remain unchanged. +- Optional dynamic micro-batching groups sequences by density to reduce padding overhead. + +The pipeline: + +1. Actor parameters are synchronized to a rollout worker group (nccl broadcast) with a little bit latency. +2. Rollout workers (vLLM-backed) generate sequences asynchronously (`async_generate_sequences`). +3. Teacher client service (ZeroMQ based) returns top-k log-probabilities + token indices for each sequence (batched micro-requests), enabling KL-based guidance. +4. Megatron actor performs a KL divergence computation between student logits and teacher top-k distributions (custom TP-aware kernel in `megatron_kl_loss.py`). +5. Scheduling strategies (`one_step_off_scheduler`, `two_step_off_scheduler`) can overlap phases (optional for throughput): + +### 3.2 Weights sync between actor and rollout + +We initially followed the weight synchronization path from the One-Step-Off-Policy recipe (Ray collective broadcast across all actor and rollout ranks, plus Megatron-side allgather of parameter shards). In practice this became the dominant bottleneck, so we made three changes: + +1. Batch-and-bulk load on the rollout side: instead of streaming tensors one-by-one (in one-step-off-policy recipe), we stage a bundle of parameter tensors and issue a single batched load into the rollout engine. In our setup this reduced the weight-loading time by roughly 3×. +2. Batch-and-bulk broadcast between the actor and rollout: instead of streaming tensors one-by-one (in one-step-off-policy recipe), we stage a bundle of parameter tensors and issue a single batched broadcast between the actor and rollout workers. +3. Replace allgather with gather-to-root in Megatron: parameter shards are gathered to actor rank 0 (rather than allgathered to everyone), and that root then serves as the single source for broadcasting to rollout ranks. On top of the previous change, 2 and 3 changes delivered an additional ~4× speedup in the synchronization phase. + +## 4. High-Level Data & Control Flow + +``` +Driver (TaskRunner) + ├─ Initialize Ray, tokenizer, datasets, worker groups + ├─ Build ResourcePoolManager (actor vs rollout GPU layouts) + ├─ Trainer.fit() + ├─ init_workers(): build actor + rollout groups, broadcast weight metadata, create nccl collective group + ├─ continuous_iterator(): epochs → batches + ├─ scheduler (see Section 6) + • _async_gen_next_batch(): optional weight sync + non-blocking rollout + • _async_get_teacher_knowledge(): submit teacher requests, store future + ├─ For each step: + • Sync rollout weights + • Retrieve (batch, gen_output, teacher_output) from futures + • Merge gen + teacher outputs → DataProto + • Compute metrics (response length stats, timing, throughput) + • Update actor (forward_backward_batch + KL loss + optimizer step) + • (Optional) save checkpoint +``` + +> Note: Schedulers are optional and explained later; the distillation objective is independent of how phases are overlapped. + +## 5. Key Components + +### 5.1 `OnPolicyDistillTrainer` (`ray_trainer.py`) +- Creates `GenerationBatchFuture` objects holding rollout and (later) teacher futures. +- Adds scheduling + teacher integration + modified metric emission (KL, timing, MFU). + +### 5.2 Actor Worker (Megatron) +- `OnPolicyDistillActor.update_policy()` orchestrates micro-batch forward/backward. +- KL Loss injection via `logits_processor` during forward on pipeline last stage. + +### 5.3 Rollout Worker (vLLM / SGLang) +- Pure inference mode (`init_model` builds model; no optimizer). +- `async_generate_sequences` returns a Ray future for overlapping. + +### 5.4 Teacher Service (`teacher/`) +- Proxy + worker architecture (ZMQ REQ/REP) for batched top-k retrieval. +- `TeacherClient.submit()` returns a `Future`; aggregator composes micro-batches. +- Configurable temperature, max tokens, only-response mode. + +### 5.5 KL Loss (`megatron_kl_loss.py`) +- Performs normalization & stable per-token probability construction across TP shards. +- Gradient is (student_probs - teacher_sparse_probs) scaled by upstream grad. + +## 6. Configuration Highlights (`on_policy_distill_trainer.yaml`) + +| Section | Purpose | Notable Keys | +|---------|---------|-------------| +| actor_rollout_ref.teacher | Teacher server | server_ip, server_port, n_server_workers | +| trainer | Global training control | total_epochs, save_freq, scheduler (one_step_off | two_step_off), n_gpus_per_node, nnodes | +| rollout | Resource split for rollout | n_gpus_per_node, nnodes | + +**Remember to set `trainer.n_gpus_per_node`, `trainer.nnodes`, `rollout.n_gpus_per_node` and `rollout.nnodes` to allocate GPU resources.** + +### Dynamic Batch Size + +Enable by: + +``` +actor_rollout_ref.actor.use_dynamic_bsz=True +actor_rollout_ref.actor.max_token_len=6000 # cap post-group token length +``` + +Improves utilization under variable sequence lengths. + +### Resource Guidelines + +- Actor pool: `trainer.nnodes * trainer.n_gpus_per_node` GPUs. +- Rollout pool: `rollout.nnodes * rollout.n_gpus_per_node` GPUs. +- Ensure teacher server capacity ≈ `n_server_workers` to avoid stalls (monitor `wait_prev_teacher`). + +## 7. Usage Examples + +### 7.1 Launch Teacher Server + +Before training process, you should have a teacher server to provide logp information. + +We provide a toy teacher server example with vLLM. It needs `telnet` to check proxy status, and `python` command to run. So if you have not installed `telnet`, you can just delete these code in `start_server.sh`. And some OS use `python3` rather than `python`, so you also need to modify it. Also you can change the port of teacher if you meet port conflict. + +There are 3 arguments can be set for vllm backend `--tp-size`, `--n-logprobs` and `--ckpt-path` in `start_server.sh` / `worker.py`. You should set before you start server. + +We also provide a toy multi-node teacher server. You can start the main node using `start_server.sh` and start the slave nodes using `join_server.sh`. Still remember to set args in `join_server.sh`, especially the `$PROXY_IP` and `$PROXY_BACKEND_PORT` of main node. + +When training, student will automatically use the teacher's topk (n-logprobs) to set its own topk argument at line 83 of `recipe/gkd/megatron_kl_loss.py`, so you don't need to set student's topk argument. + +```bash +cd recipe/gkd/teacher +bash start_server.sh +# Exports ports and launches proxy + worker (default vLLM backend) +``` + +Verify with: + +```bash +telnet localhost 15555 +``` + +### 7.2 Minimal Local (Megatron + vLLM) Run + +```bash +python3 -m recipe.gkd.main_gkd \ + --config-path=recipe/gkd/config \ + --config-name=on_policy_distill_trainer \ + actor_rollout_ref.model.path=/path/to/MODEL \ + data.train_files=/path/to/train.parquet \ + trainer.total_epochs=2 \ + trainer.n_gpus_per_node=4 rollout.n_gpus_per_node=2 \ + actor_rollout_ref.teacher.server_ip=127.0.0.1 \ + actor_rollout_ref.teacher.server_port=15555 \ + trainer.scheduler=one_step_off +``` + +(Requires a running teacher server). + +### 7.3 Ray Job Submission (Distilled 16B Example) + +See `run_moonlight_dsv3_training.sh` for a full script including: + +- Dist ckpt path setup (`dist_checkpointing_path`) +- Expert parallel sizing (EP / ETP) +- Dynamic batch sizing +- Two-step-off scheduling for deeper overlap. + +Submit (after adjusting paths): + +```bash +bash recipe/gkd/run_moonlight_dsv3_training.sh +``` + +## 8. Metrics & Monitoring + +Emitted metrics include (prefixes may vary): + +- Timing: `timing/wait_prev_gen`, `timing/sync_rollout_weights`, `timing/get_teacher_knowledge`, `timing/update_actor`. +- Sequence stats: `response_seq_len/*` (avg, max, min, counts). +- Performance: `perf/mfu/actor`, `perf/max_memory_allocated_gb`, `perf/cpu_memory_used_gb`. +- Distillation: `actor/kl_loss`, `actor/grad_norm`, `actor/lr`. + +Interpretation Tips: + +- High `wait_prev_teacher` → scale `n_server_workers` and allocate more teacher GPUs or reduce per-request batch size, or just use `two_step_off`. +- High `wait_prev_gen` with uniform lengths → allocate more rollout GPUs. +- High `sync_rollout_weights` → check NCCL env / network congestion and try to modify `actor_rollout_ref.rollout.update_weights_bucket_megabytes`. + +## 9. Extensibility Notes + +- Add new schedulers by following interface returning `(epoch, batch, gen_output, teacher_output, timing_dict)`. +- Integrate different distillation signals (e.g., hidden states, intermediate reasoning tokens) by extending `teacher_utils.get_teacher_knowledge` and modifying `logits_processor`. + +## 10. Functional Support Summary + +| Category | Supported | +|----------|-----------| +| Train engine | Megatron | +| Rollout engine | vLLM | +| Distillation signal | Teacher top-k logprobs & indices | +| Scheduling | one_step_off, two_step_off | + +## 11. Quick Checklist Before Running + +- Teacher server reachable (`telnet `). +- `actor_rollout_ref.model.path` contains the correct Megatron/HF config artifacts. +- `train_files` points to a parquet dataset compatible with this recipe's dataset loader. +- NCCL environment vars set (see `config/runtime_env.yaml`). + +--- +Feel free to open issues or PRs to extend scheduler variants, add new distillation objectives, or broaden engine support, and more improvement. diff --git a/verl/docs/advance/attention_implementation.rst b/verl/docs/advance/attention_implementation.rst new file mode 100644 index 0000000000000000000000000000000000000000..c068bd92115d38a86b4ba9414ae4c5e5a18a2218 --- /dev/null +++ b/verl/docs/advance/attention_implementation.rst @@ -0,0 +1,119 @@ +.. _attention-implementation-override: + +Attention Implementation Override +================================== + +Last updated: 10/31/2025. + +By default, VERL's FSDP workers use ``flash_attention_2`` as the attention implementation for improved performance. +However, you can now override this setting to use different attention implementations based on your needs. + +Supported Attention Implementations +----------------------------------- + +The following attention implementations are supported (subject to model and hardware compatibility): + +- ``flash_attention_2``: High-performance attention implementation (default) +- ``eager``: Standard PyTorch attention implementation +- ``sdpa``: Scaled Dot-Product Attention (PyTorch native) + +When to Override +---------------- + +You might want to override the attention implementation in the following scenarios: + +- **Debugging**: Use ``eager`` for easier debugging and better error messages +- **Compatibility**: Some models or hardware configurations may not support ``flash_attention_2`` +- **Memory constraints**: Different implementations have different memory characteristics +- **Performance tuning**: Testing different implementations for optimal performance + +Configuration Examples +----------------------- + +PPO Training with Eager Attention +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +To override the attention implementation for the actor, rollout, and reference models: + +.. code:: bash + + python3 ppo_trainer.py \ + +actor_rollout_ref.model.override_config.attn_implementation=eager \ + [other parameters...] + +PPO Training with SDPA Attention +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. code:: bash + + python3 ppo_trainer.py \ + +actor_rollout_ref.model.override_config.attn_implementation=sdpa \ + [other parameters...] + +Critic Model Override +~~~~~~~~~~~~~~~~~~~~~ + +For training configurations that include a critic model, you can also override its attention implementation: + +.. code:: bash + + python3 ppo_trainer.py \ + +actor_rollout_ref.model.override_config.attn_implementation=eager \ + +critic.model.override_config.attn_implementation=eager \ + [other parameters...] + +YAML Configuration +~~~~~~~~~~~~~~~~~~ + +You can also specify the attention implementation in your YAML configuration file: + +.. code:: yaml + + actor_rollout_ref: + model: + override_config: + attn_implementation: eager + # other overrides... + + critic: # if using a critic model + model: + override_config: + attn_implementation: eager + # other overrides... + +Important Notes +--------------- + +**Backward Compatibility**: If you don't specify ``attn_implementation`` in the override config, +VERL will continue to use ``flash_attention_2`` by default, ensuring backward compatibility with existing configurations. + +**Model Support**: Not all models support all attention implementations. Ensure your model is compatible +with the chosen attention implementation before training. + +**Performance Impact**: Different attention implementations have varying performance characteristics. +``flash_attention_2`` typically offers the best performance, while ``eager`` provides better debugging capabilities. + +**Hardware Dependencies**: Some attention implementations (like ``flash_attention_2``) may require +specific hardware or CUDA versions. If you encounter compatibility issues, try using ``eager`` or ``sdpa``. + +Troubleshooting +--------------- + +If you encounter errors when using a specific attention implementation: + +1. **Check model compatibility**: Verify that your model supports the chosen attention implementation +2. **Try eager attention**: Use ``attn_implementation=eager`` as a fallback for debugging +3. **Check hardware requirements**: Ensure your hardware supports the attention implementation +4. **Review error messages**: Attention implementation errors often provide clear guidance on supported options + +Example Error Resolution +~~~~~~~~~~~~~~~~~~~~~~~~ + +If you see an error like "flash_attention_2 is not supported", you can resolve it by switching to eager attention: + +.. code:: bash + + # Instead of the default flash_attention_2 + python3 ppo_trainer.py +actor_rollout_ref.model.override_config.attn_implementation=eager + +This override ensures your training can proceed while you investigate the flash attention compatibility issue. diff --git a/verl/docs/advance/checkpoint.rst b/verl/docs/advance/checkpoint.rst new file mode 100644 index 0000000000000000000000000000000000000000..4e224f6efc3759aa898b9f755cbc1253fd10abef --- /dev/null +++ b/verl/docs/advance/checkpoint.rst @@ -0,0 +1,188 @@ +.. _checkpoint-page: + +Using Checkpoints to Support Fault Tolerance Training +===================================================== + +Last updated: 03/22/2026. + +There could be training errors or machine failure during the whole RLHF training process, +so it is recommended to enable checkpoints to minimize your loss. + +The API Interface has already been listed in :ref:`config-explain-page`, +and we will not repeat them. But there are still some technique details +we hope to clarify. + +.. note:: + + Notice that the ``checkpoint.contents`` field has no effect to FSDP checkpoint except ``hf_model``, + the other 3 fields are binded together to save and load. We recommend to include ``model``, ``optimizer`` and ``extra`` all. + +Checkpoint Saving Directory Structure +------------------------------------- + +Commonly, we use the ``default_local_dir`` declared in ``ppo_trainer.yaml`` or ``ppo_megatron_trainer.yml`` +to work as preffix when saving checkpoints, which is ``checkpoints/${trainer.project_name}/${trainer.experiment_name}``. + +So the inner checkpoint structure of **FSDP** is like: + +.. code:: + + checkpoints/${trainer.project_name}/${trainer.experiment_name} + ├── global_steps_${i} + │ ├── actor + │ │ ├── huggingface # default save config and tokenizer, save huggingface model if include ``hf_model`` in checkpoint.contents + │ │ └── fsdp_config.json # FSDP config file, including world_size and fsdp version + │ │ ├── model_world_size_{self.world_size}_rank_{self.rank}.pt + │ │ ├── optim_world_size_{self.world_size}_rank_{self.rank}.pt + │ │ └── extra_state_world_size_{self.world_size}_rank_{self.rank}.pt + │ ├── critic + │ │ ├── huggingface + │ │ └── fsdp_config.json + │ │ ├── model_world_size_{self.world_size}_rank_{self.rank}.pt + │ │ ├── optim_world_size_{self.world_size}_rank_{self.rank}.pt + │ │ └── extra_state_world_size_{self.world_size}_rank_{self.rank}.pt + └── latest_checkpointed_iteration.txt + +All model shards, optimizers and extra states are stored together, in a sharded and distributed way. + +While **Megatron** current checkpoint structure is: + +.. code:: + + checkpoints/${trainer.project_name}/${trainer.experiment_name} + ├── global_steps_${i} + │ ├── actor + │ │ ├── huggingface # default save config and tokenizer, save huggingface model if include ``hf_mode`` in checkpoint.contents + │ │ └── dist_ckpt # save sharded model/optimizer/rng_states, naming the same as Megatron + │ └── critic + │ │ ├── huggingface + │ │ └── dist_ckpt + └── latest_checkpointed_iteration.txt + +Convert FSDP and Megatron Checkpoints to HuggingFace Format Model +----------------------------------------------------------------- + +We provide a tool to convert the FSDP and Megatron checkpoints to HuggingFace format model. +The tool is located in ``verl/model_merger``. For older versions of verl that don't include fsdp_config.json in checkpoints, you can use the legacy model merger located at ``verl/scripts/legacy_model_merger.py``. + +The script supports two main sub-commands: `merge` (to convert and save checkpoints) and `test` (to validate merged checkpoints against a reference model). +The arguments for the `merge` sub-command are as follows: + +.. code:: bash + + usage: python -m verl.model_merger merge [-h] --backend {fsdp,megatron} [--local_dir LOCAL_DIR] [--tie-word-embedding] [--is-value-model] [--use_cpu_initialization] [--target_dir TARGET_DIR] + [--hf_upload_path HF_UPLOAD_PATH] [--private] + + options: + -h, --help show this help message and exit + --backend {fsdp,megatron} + The backend of the model + --local_dir LOCAL_DIR + Path to the saved model checkpoints + --tie-word-embedding Whether to tie word embedding weights (currently only Megatron supported) + --is-value-model Whether the model is a value model (currently only Megatron supported) + --use_cpu_initialization + Whether to use CPU initialization for the model. This is useful for large models that cannot fit into GPU memory during initialization. + --target_dir TARGET_DIR + Directory to save the merged huggingface model + --hf_upload_path HF_UPLOAD_PATH + Hugging Face repository ID to upload the model + --private Whether to upload the model to a private Hugging Face repository + +Example usage for merging Megatron checkpoints: + +.. code:: bash + + python -m verl.model_merger merge \ + --backend megatron \ + --tie-word-embedding \ + --local_dir checkpoints/verl_megatron_gsm8k_examples/qwen2_5_0b5_megatron_saveload/global_step_1/actor \ + --target_dir /path/to/merged_hf_model + +Example usage for distributed merging Megatron checkpoints: + +.. code:: bash + + torchrun --nproc_per_node 1 --nnodes 8 --node_rank ${RANK} -m verl.model_merger merge \ + --backend megatron \ + --tie-word-embedding \ + --local_dir checkpoints/verl_megatron_gsm8k_examples/qwen2_5_0b5_megatron_saveload/global_step_1/actor \ + --target_dir /path/to/merged_hf_model + +Example usage for merging FSDP checkpoints: + +.. code:: bash + + python -m verl.model_merger merge \ + --backend fsdp \ + --local_dir checkpoints/verl_fsdp_gsm8k_examples/qwen2_5_0b5_fsdp_saveload/global_step_1/actor \ + --target_dir /path/to/merged_hf_model + + +Megatron Merger details +----------------------- + +Current implement of decoder layers uses ``nn.ModuleList`` to store the layers, +and thus the model layers on every PP rank and VPP rank starts their index from 0. + +There are 3 ways to correct this behavior: + +1. Modify the decoder layer's state_dict, add ``offset`` to each layer's index, thus rewrite ``nn.ModuleList`` implementation. +2. Modify the layer index when saving checkpoint and recover them when loading checkpoint. +3. The Checkpoint merger do this work, calculate the actual ``offset`` from ``state_dict`` only, a little complex. + +Current implementation use solution 2. + + +HuggingFace to Megatron DistCheckpoint details +---------------------------------------------- + +Through ``mbridge``, we can directly save the mcore model to huggingface format during training. +No need to convert the model to Megatron dist-checkpoint format. + +.. note:: + + Megatron provides multiple optimizer checkpoint formats controlled by: + + - ``dist_ckpt_optim_fully_reshardable``: + + - ``False`` (default, dp-reshardable): + The optimizer checkpoint supports resuming with different data parallel sizes. + This format is faster and has lower memory overhead during checkpoint saving. + + - ``True`` (fully-reshardable): + The optimizer checkpoint supports resuming with arbitrary parallelism configurations. + However, this format is slower and introduces additional memory overhead. + + - ``distrib_optim_fully_reshardable_mem_efficient``: + + When using fully-reshardable format, enabling this option switches communication + from NCCL to Gloo to reduce CUDA memory usage, at the cost of performance. + +.. warning:: + + When ``dist_ckpt_optim_fully_reshardable=True``, saving optimizer checkpoints requires + gathering optimizer states on data parallel rank 0. Although the final checkpoint is + sharded, this introduces a temporary aggregation step during saving. + + This may increase CPU memory usage and lead to OOM issues for large models. + We recommend using the default dp-reshardable format in most cases. + + +Original Checkpoint Utils +------------------------- + +Original Checkpoint Utils refer to original checkpoint implementation in ``verl/models/[model]/megatron/checkpoint_utils``. + +We only need ``[model]_loader.py`` in original checkpoint utils now, since we get rid of storing ``hf_model`` every time (which is not recommended for large model training, try only saving sharded models if you can). + +.. note:: + + Note that ``[model]_loader`` only support environments where **storage clusters are able to connect with every calculation nodes**. + Because it utilizes **sharded load way to minimize the loading checkpoint overhead**. + Every rank loads its own data from ``state_dict`` which can be accessed by all of them. + While there is also no need to broadcast among DP ranks, since the saved state_dict is only produced by DP rank 0. + + For users who can **only place the huggingface model on one device**, we keep the original costly implementation in ``[model]_loader_deprecated``. In this implementation, rank 0 broadcast all weights to each tp and pp rank, and then dp rank 0 broadcast to all dp ranks. There may be at risks of OOM. + + To use deprecated loader, change the import package of ``load_state_dict_to_megatron_llama``. diff --git a/verl/docs/advance/dpo_extension.rst b/verl/docs/advance/dpo_extension.rst new file mode 100644 index 0000000000000000000000000000000000000000..25d159cea6aa774367021358498d6f1509bf0b92 --- /dev/null +++ b/verl/docs/advance/dpo_extension.rst @@ -0,0 +1,273 @@ +Extend to other RL(HF) algorithms +================================= + +Last updated: 02/25/2025. + +We already implemented the complete training pipeline of the PPO +algorithms. To extend to other algorithms, we analyze the high-level +principle to use verl and provide a tutorial to implement the DPO +algorithm. Users can follow the similar paradigm to extend to other RL algorithms. + +.. note:: **Key ideas**: Single process drives multi-process computation and data communication. + +Overall Approach +---------------- + +Step 1: Consider what multi-machine multi-GPU computations are needed +for each model, such as ``generate_sequence`` , ``compute_log_prob`` and +``update_policy`` in the actor_rollout model. Implement distributed +single-process-multiple-data (SPMD) computation and encapsulate them +into APIs + +Step 2: Based on different distributed scenarios, including FSDP and 3D +parallelism in Megatron-LM, implement single-process control of data +interaction among multi-process computations. + +Step 3: Utilize the encapsulated APIs to implement the control flow + +Example: Online DPO +------------------- + +We use verl to implement a simple online DPO algorithm. The algorithm +flow of Online DPO is as follows: + +1. There is a prompt (rollout) generator which has the same weight as + the actor model. After a batch of prompts are fed into the generator, + it generates N responses for each prompt. +2. Send all the prompts + responses to a verifier for scoring, which can + be reward model or a rule-based function. Then sort them in pairs to + form a training batch. +3. Use this training batch to train the actor model using DPO. During + the process, a reference policy is needed. + +Step 1: What are the multi-machine multi-GPU computations +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +**Sample Generator** + +Implementation details: + +.. code:: python + + from verl.single_controller.base import Worker + from verl.single_controller.ray import RayWorkerGroup, RayClassWithInitArgs, RayResourcePool + import ray + + @ray.remote + class SampleGenerator(Worker): + def __init__(self, config): + super().__init__() + self.config = config + + def generate_sequences(self, data): + pass + +Here, ``SampleGenerator`` can be viewed as a multi-process pulled up by +``torchrun``, with each process running the same code (SPMD). +``SampleGenerator`` needs to implement a ``generate_sequences`` API for +the control flow to call. The implementation details inside can use any +inference engine including vllm, sglang and huggingface. Users can +largely reuse the code in +verl/verl/workers/rollout/vllm_rollout/vllm_rollout.py and we won't +go into details here. + +**ReferencePolicy inference** + +API: compute reference log probability + +.. code:: python + + from verl.single_controller.base import Worker + import ray + + @ray.remote + class ReferencePolicy(Worker): + def __init__(self): + super().__init__() + self.model = Model() + + def infer(self, data): + return self.model(data) + +**Actor update** + +API: Update actor model parameters + +.. code:: python + + from verl.single_controller.base import Worker + import ray + + @ray.remote + class DPOActor(Worker): + def __init__(self): + super().__init__() + self.model = Model() + self.model = FSDP(self.model) # or other distributed strategy + self.optimizer = optim.Adam(self.model.parameters(), lr=1e-3) + self.loss_fn = xxx + + def update(self, data): + self.optimizer.zero_grad() + logits = self.model(data) + loss = self.loss_fn(logits) + loss.backward() + self.optimizer.step() + +**Notes: How to distinguish between control processes and distributed computation processes** +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +- Control processes are generally functions directly decorated with + ``@ray.remote`` +- Computation processes are all wrapped into a ``RayWorkerGroup``. + +Users can reuse most of the distribtued computation logics implemented +in PPO algorithm, including FSDP and Megatron-LM backend in +verl/verl/trainer/ppo. + +Step 2: Based on different distributed scenarios, implement single-process control of multi-process data interaction +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +**The core problem to solve here is how a single process sends data to +multiple processes, drives multi-process computation, and how the +control process obtains the results of multi-process computation.** +First, we initialize the multi-process ``WorkerGroup`` in the control +process. + +.. code:: python + + @ray.remote(num_cpus=1) + def main_task(config): + # construct SampleGenerator + resource_pool = RayResourcePool(process_on_nodes=[8] * 2) # 16 GPUs + ray_cls = RayClassWithInitArgs(SampleGenerator, config=config) + # put SampleGenerator onto resource pool + worker_group = RayWorkerGroup(resource_pool, ray_cls) + + # construct reference policy + +As we can see, in the control process, multiple processes are wrapped +into a ``RayWorkerGroup``. Inside this ``WorkerGroup``, there is a +``self._workers`` member, where each worker is a RayActor +(https://docs.ray.io/en/latest/ray-core/actors.html) of SampleGenerator. +ray_trainer.md also provide an implementation of +``MegatronRayWorkerGroup``. + +Assuming the model is distributed using FSDP, and there is a batch of +data on the control process, for data parallelism, the underlying +calling process is: + +.. code:: python + + data = xxx + data_list = data.chunk(dp_size) + + output = [] + for d in data_list: + # worker_group._workers[i] is a SampleGenerator + output.append(worker_group._workers[i].generate_sequences.remote(d)) + + output = ray.get(output) + output = torch.cat(output) + +Single process calling multiple processes involves the following 3 +steps: + +1. Split the data into DP parts on the control process. +2. Send the data to remote, call the remote computation through RPC, and + utilize multi-process computation. +3. Obtain the computation results of each worker on the control process + and merge them. + +Frequently calling these 3 steps on the controller process greatly hurts +code readability. **In verl, we have abstracted and encapsulated these 3 +steps, so that the worker's method + dispatch + collect can be +registered into the worker_group** + +.. code:: python + + from verl.single_controller.base.decorator import register + + def dispatch_data(worker_group, data): + return data.chunk(worker_group.world_size) + + def collect_data(worker_group, data): + return torch.cat(data) + + dispatch_mode = { + 'dispatch_fn': dispatch_data, + 'collect_fn': collect_data + } + + @register(dispatch_mode=dispatch_mode) + def generate_sequences(self, data): + pass + +In this way, we can directly call the method inside the worker through +the ``worker_group`` on the control (driver) process (which is a single +process): + +.. code:: python + + output = worker_group.generate_sequences(data) + +This single line includes data splitting, data distribution and +computation, and data collection. + +Furthermore, the model parallelism size of each model is usually fixed, +including dp, tp, pp. So for these common distributed scenarios, we have +pre-implemented specific dispatch and collect methods,in `decorator.py `_, which can be directly used to wrap the computations. + +.. code:: python + + from verl.single_controller.base.decorator import register, Dispatch + + @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO) + def generate_sequences(self, data: DataProto) -> DataProto: + pass + +Here it requires the data interface to be ``DataProto``. Definition of +``DataProto`` is in `protocol.py `_. + +Step 3: Main training loop +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +With the above training flows, we can implement the algorithm's control +flow. It is recommended that ``main_task`` is also a ray remote process. + +.. code:: python + + @ray.remote(num_cpus=1) + def main_task(config): + # construct SampleGenerator + resource_pool = RayResourcePool(process_on_nodes=[8] * 2) # 16 GPUs + ray_cls = RayClassWithInitArgs(SampleGenerator, config=config) + # put SampleGenerator onto resource pool + sample_gen = RayWorkerGroup(resource_pool, ray_cls) + + # construct reference policy + ray_cls = RayClassWithInitArgs(ReferencePolicy) + ref_policy = RayWorkerGroup(resource_pool, ray_cls) + + # construct actor + ray_cls = RayClassWithInitArgs(DPOActor) + dpo_policy = RayWorkerGroup(resource_pool, ray_cls) + + dataloader = DataLoader() + + for data in dataloader: + # generate data + data = sample_gen.generate_sequences(data) + # generate scores for each data + data = generate_scores(data) + # generate pairwise data using scores + data = generate_pairwise_data(data) + # generate ref_log_prob + data.batch['ref_log_prob'] = ref_policy.infer(data) + # update using dpo + dpo_policy.update(data) + # logging + +Here, different ``WorkerGroups`` can be placed in the same resource pool or +in different resource pools using ``create_colocated_worker_cls`` +similar as in `ray_trainer.py `_. diff --git a/verl/docs/advance/fsdp_extension.rst b/verl/docs/advance/fsdp_extension.rst new file mode 100644 index 0000000000000000000000000000000000000000..181e109082262f26334034337c5915d522049759 --- /dev/null +++ b/verl/docs/advance/fsdp_extension.rst @@ -0,0 +1,97 @@ + +Add models with the FSDP backend +================================== + +Last updated: 02/09/2025. + +Model +-------------------------- + +In principle, our FSDP backend can support any HF model and we can +sychronoize the actor model weight with vLLM using `hf_weight_loader.py` under `third_party/vllm`. +However, ``hf_weight_loader`` is will gather the full state_dict of a +model during synchronization, which may cause OOM. We suggest using +``dtensor_weight_loader`` which gather the full model parameter layer by +layer to reduce the peak memory usage. We already support dtensor weight +loader for the models below in `dtensor_weight_loader.py` under `third_party/vllm`: + +- ``GPT2LMHeadModel`` +- ``LlamaForCausalLM`` +- ``LLaMAForCausalLM`` +- ``MistralForCausalLM`` +- ``InternLMForCausalLM`` +- ``AquilaModel`` +- ``AquilaForCausalLM`` +- ``Phi3ForCausalLM`` +- ``GemmaForCausalLM`` +- ``Gemma2ForCausalLM`` +- ``GPTBigCodeForCausalLM`` +- ``Starcoder2ForCausalLM`` +- ``Qwen2ForCausalLM`` +- ``DeepseekV2ForCausalLM`` + +To implement ``dtensor_weight_loader`` of a model that's supported in +vLLM, follow the guide of gemma model below: + +1. Copy the + ``load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]])`` from the vllm model class + to ``dtensor_weight_loaders.py`` +2. Modify the arguments to + ``(actor_weights: Dict, vllm_model: nn.Module)`` +3. Replace the ``self`` to ``vllm_model`` +4. Add the + ``local_loaded_weight = redistribute_dtensor(param_name=name, loaded_weights=loaded_weight)`` + before each ``param = params_dict[name]`` and modify the following + weight loading using ``local_loaded_weight``. +5. Register the implemented dtensor weight loader to ``__MODEL_DTENSOR_WEIGHT_LOADER_REGISTRY__``. + +.. code-block:: diff + + - def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + + def gemma_dtensor_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + - params_dict = dict(self.named_parameters()) + + params_dict = dict(vllm_model.named_parameters()) + loaded_params = set() + - for name, loaded_weight in weights: + + for name, loaded_weight in actor_weights.items(): + for (param_name, shard_name, shard_id) in stacked_params_mapping: + if shard_name not in name: + continue + name = name.replace(shard_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + local_loaded_weight = redistribute_dtensor(param_name=name, loaded_weights=loaded_weight) + param = params_dict[name] + weight_loader = param.weight_loader + - weight_loader(param, loaded_weight, shard_id) + + weight_loader(param, local_loaded_weight.to(dtype=param.dtype), shard_id) + break + else: + # lm_head is not used in vllm as it is tied with embed_token. + # To prevent errors, skip loading lm_head.weight. + if "lm_head.weight" in name: + continue + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + local_loaded_weight = redistribute_dtensor(param_name=name, loaded_weights=loaded_weight) + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", + default_weight_loader) + - weight_loader(param, loaded_weight) + + weight_loader(param, local_loaded_weight.to(dtype=param.dtype)) + loaded_params.add(name) + unloaded_params = params_dict.keys() - loaded_params + if unloaded_params: + raise RuntimeError( + "Some weights are not initialized from checkpoints: " + f"{unloaded_params}") \ No newline at end of file diff --git a/verl/docs/advance/fully_async.md b/verl/docs/advance/fully_async.md new file mode 100644 index 0000000000000000000000000000000000000000..9a021131c130ea18d11cf7368e1c4b5240533e2f --- /dev/null +++ b/verl/docs/advance/fully_async.md @@ -0,0 +1,579 @@ +# Recipe: Fully Async Policy Trainer + +**Author:** `https://github.com/meituan-search` + +Last updated: 05/09/2026. + +This document introduces a fully asynchronous PPO training system that completely decouples the Trainer and Rollouter, +supporting asynchronous sample generation and training. +Under this system, we achieved a 2.35x-2.67x performance improvement when training the Qwen2.5-7B model with 128 GPUs, +without significantly affecting the results. + +## Introduction + +### Background + +The separated rollout and train architecture, compared to the colocate architecture, can allocate resources more +flexibly and design more flexible training logic, thereby addressing issues such as low GPU utilization and training +efficiency caused by long-tail problems. +The one_step_off_policy alleviates the problem of long rollout times and achieves some gains in training efficiency by +designing a separated architecture and performing asynchronous training between rollout and train for one round. +However, it forcibly uses data from one round of asynchronous training, which is not flexible enough and cannot +completely eliminate the impact of long-tail on training efficiency. +In other frameworks such as AReaL, Magistral, StreamRL, and AsyncFlow, asynchronous training and streaming training have +been implemented based on the separated architecture and have achieved gains. +We borrow from their methods and implemented them in VERL. The fully_async_policy supports asynchronous, streaming, and +partial +rollout training. +By reasonably setting parameters such as resource allocation and parameter synchronization frequency, fully_async_policy +can significantly improve training efficiency. + +> Magistral https://arxiv.org/abs/2506.10910 +> +> AReaL: A Large-Scale Asynchronous Reinforcement Learning System for Language +> Reasoning https://arxiv.org/abs/2505.24298 +> +> StreamRL: Scalable, Heterogeneous, and Elastic RL for LLMs with Disaggregated Stream +> Generation https://arxiv.org/abs/2504.15930 +> +> AsyncFlow: An Asynchronous Streaming RL Framework for Efficient LLM Post-Training https://arxiv.org/abs/2507.01663 +> + +### Core Contributions + +* **Resource Isolation**: Unlike using hybrid_engine, Rollouter and Trainer use separate computing resources and need to + specify the resources they occupy separately. +* **Parallel Generation and Training**: While the Trainer is training, the Rollouter is generating new samples. +* **Multi-step Asynchronous**: Compared to one step off policy, it supports asynchronous settings from 0.x steps to + multiple steps, making the asynchronous solution more flexible. +* **NCCL Parameter Synchronization**: Based on the nccl communication primitive, refer + to [checkpoint-engine](https://github.com/MoonshotAI/checkpoint-engine) to + achieve efficient parameter synchronization between Rollouter and Trainer. +* **Stream Inference and Training**: Rollouter generates data sample by sample, and data transmission uses a single + sample as the minimum transmission unit. +* **Asynchronous Training and Freshness Control**: By setting the parameter async_training.staleness_threshold, it + supports training with samples generated by old parameters. +* **PartialRollout**: The Rollouter's inference process supports partial rollout logic. During parameter + synchronization, by adding `sleep() and resume()` logic, it + saves samples from ongoing rollouts and continues using them in the next rollout, reducing the time spent waiting for + ongoing tasks to finish during parameter synchronization. + +Currently, the supported usage mode is megatron/fsdp+vllm. vllm must use the server mode based on AgentLoop. + +## Design + +The overall architecture of fully_async_policy is shown in the figure below. fully_async_policy mainly consists of four +parts: Rollouter, MessageQueue, Trainer, and ParameterSynchronizer. + +![fully_async_policy_structure]( +https://github.com/ArronHZG/verl-community/blob/main/docs/fully_async_policy_structure.svg?raw=true) + +1. Rollouter generates sequences sample by sample and puts the generated samples into the MessageQueue, with the + production speed controlled by freshness. +2. MessageQueue is used to temporarily store samples generated by Rollouter. +3. Trainer fetches samples from MessageQueue sample by sample. After fetching `require_batches*ppo_mini_batch_size` + samples, it will perform training. After training for async_training.trigger_parameter_sync_step rounds, it triggers + a parameter synchronization with Rollouter. +4. ParameterSynchronizer implements the NCCL synchronous parameter synchronization capability. + +The source of benefits compared to the base scheme lies in the fact that in the colocate case, using more resources for +rollout cannot solve the idleness caused by long-tail samples. +After we perform resource isolation, the time for rollout and train may be longer than before (because fewer resources +are used), +but the overlap in their time consumption reduces the end-to-end time consumption. + +![fully_async_policy_revenue]( +https://github.com/ArronHZG/verl-community/blob/main/docs/fully_async_policy_revenue.svg?raw=true) + +## Usage + +### Parameter Description + +| super params | implication | +|------------------------------------------------------------------|------------------------------------------------------------------------------------------------| +| `trainer.nnodes` | Number of nodes for Trainer | +| `trainer.n_gpus_per_node` | Number of GPUs per node for Trainer | +| `rollout.nnodes` | Number of nodes for Rollouter | +| `rollout.n_gpus_per_node` | Number of GPUs per node for Rollouter | +| `data.train_batch_size` | In the fully async strategy, this value is not effective (default is 0) | +| `data.gen_batch_size` | In the fully async strategy, uses streaming sample production logic (default is 1) | +| `rollout.total_rollout_steps` | Total number of rollout samples | +| `rollout.test_freq` | How many times Rollouter updates parameters before performing a validation | +| `actor_rollout_ref.actor.ppo_mini_batch_size` | The ppo_mini_batch_size is a global num across all workers/gpus | +| `actor_rollout_ref.actor.use_rollout_log_probs=True` | Use log_probs generated by rollout | +| `algorithm.rollout_correction.bypass_mode` | Whether to compute log_prob using the training model's parameters during the training phase. | +| `async_training.require_batches` | Number of ppo_mini_batch_size that FullyAsyncTrainer fetches at once | +| `async_training.trigger_parameter_sync_step` | Indicates how many local updates FullyAsyncTrainer performs before a parameter synchronization | +| `async_training.staleness_threshold` | Freshness control | +| `async_training.partial_rollout` | Whether to perform partial_rollout | +| `async_training.use_trainer_do_validate` | Whether use trainer node to do validate process, default `False` | + +**Further Explanation:** + +* `rollout.total_rollout_steps` + + Compared to colocate, the quantity can be aligned by multiplying train_batch_size and step: + `rollout.total_rollout_steps = data.train_batch_size * step`. + +* `async_training.trigger_parameter_sync_step` + + In the fully async strategy, it indicates how many local updates the Trainer performs (i.e., how many times it fetches + `require_batches * ppo_mini_batch_size` samples) before a parameter synchronization with Rollouter. + Between every two parameter synchronizations between Rollouter and Trainer, the Trainer will process + `trigger_parameter_sync_step* require_batches*ppo_mini_batch_size` samples. + To fairly compare speed with colocate, `trigger_parameter_sync_step` should be set to + `data.train_batch_size / (require_batches * ppo_mini_batch_size)`. + +* `async_training.staleness_threshold` + + In the fully async strategy, it indicates the maximum proportion of stale samples allowed to be used. + + * `staleness_threshold`=0, indicates synchronous training. + Rollouter will generate a fixed number of samples between two parameter updates, the sample count is: + + `rollout_num = (trigger_parameter_sync_step*require_batches*ppo_mini_batch_size)` + * `staleness_threshold`>0, indicates asynchronous training, can be set to a decimal for more flexible asynchronous + calls. + Rollouter will generate at most the following number of samples between two parameter updates: + + `rollout_num = (1+staleness_threshold)*(trigger_parameter_sync_step*require_batches*ppo_mini_batch_size) - num_staleness_sample` + + `num_staleness_sample` represents the number of stale samples generated in excess during the last rollout. + + Since it's a streaming system, rollout continues to generate and trainer continues to consume. If rollouter is slower, + trainer will trigger parameter synchronization earlier, and rollouter will not actually produce rollout_num samples. + When rollout is fast enough, setting `staleness_threshold` to 1 is basically equivalent to one_step_off policy. + To avoid too many expired samples affecting training accuracy, it is recommended to set this value to less than 1. + +* `async_training.partial_rollout` + + partial_rollout only actually takes effect when staleness_threshold>0. + +* `async_training.require_batches` + + In streaming training, require_batches should be set to 1, indicating that training is performed after producing + enough ppo_mini_batch_size samples. + In actual testing, we found that if fewer samples are issued at once, due to the order of data distribution, it can + cause training instability and longer response lengths. + Here, we additionally provide require_batches for streaming distribution and control the number of samples + participating in training at once. + +* `actor_rollout_ref.actor.use_rollout_log_probs=True` + + In reinforcement learning algorithms, log_probs have implicit correlations with parameter versions and tokens. Due to + the settings of algorithms like PPO/GRPO/DAPO, when calculating importance sampling, + old_log_prob must use the log_probs corresponding to the rollout parameters and tokens to ensure algorithm + correctness. In the fully + async strategy, we default to old_log_prob being calculated by rollout rather than by trainer. + +* `algorithm.rollout_correction.bypass_mode` + + > algorithm.rollout_correction.bypass_mode default is True, using rollout log prob. + + During the training process, we observed that metrics and response lengths may become unstable in the later + stages of training. To mitigate this issue, we can use + the [Rollout Importance Sampling](https://verl.readthedocs.io/en/latest/advance/rollout_is.html) + technique for importance sampling. To utilize Rollout Importance Sampling, we need to compute log_prob using + the training engine, which requires enabling this switch. + Additionally, when `algorithm.rollout_correction.bypass_mode=False` and Rollout Importance Sampling are enabled under + mode d + (async stream pipeline with partial rollout), our implementation approximates `Areal's Decoupled PPO`. + +* `async_training.use_trainer_do_validate` + + It controls whether to use the trainer's `do_validate` method for validation. + If set to True, the trainer will perform validation after each parameter update. It can reduce the validation time + overhead and trainer node idle time. + If set to False, the trainer will not perform validation. + +### Supported Modes + +1. on policy pipeline: + 1. **trigger_parameter_sync_step=1, staleness_threshold=0** + 2. Rollouter produces `require_batches*ppo_mini_batch_size` samples at once, Trainer fetches these samples for + training, and after training completes, Trainer and Rollouter perform a parameter synchronization; + 3. During the rollout phase, if there are long-tail samples but few rollout samples, shorter samples cannot fill + idle resources, causing some resource waste. + 4. As shown in figure a; + +2. stream off policy pipeline: + 1. **trigger_parameter_sync_step>1, staleness_threshold=0** + 2. Synchronous streaming training will be performed. Rollouter produces + `require_batches*ppo_mini_batch_size*trigger_parameter_sync_step` samples at once, Trainer performs a local + training every time it fetches `require_batches*ppo_mini_batch_size` samples, and after training + trigger_parameter_sync_step times, Trainer and Rollouter perform a parameter synchronization; + 3. Compared to a, since more samples are generated at once, resource idleness will be lower. + 4. In one step training, there will be two periods of resource idleness: when fetching the first batch of samples, + train waits for `require_batches*ppo_mini_batch_size` samples to be produced, and during the last parameter + update, rollout waits for training to complete. + 5. As shown in figure b; + +3. async stream pipeline with stale samples: + 1. **trigger_parameter_sync_step>=1, staleness_threshold>0, partial_rollout=False** + 2. After each parameter update, Rollouter will plan to produce at most rollout_num samples (in practice, the number + of samples generated may be less than this value depending on rollout speed). + 3. If the rollout process is relatively fast, Rollouter will generate some additional samples num_stale_samples + before parameter synchronization for immediate use by Trainer after synchronization. + When triggering parameter synchronization, if Rollouter has ongoing tasks, it will wait for the tasks to complete + and not add new tasks; + 4. Compared to b, except for the first step training, subsequent training will not have the time to wait for the + first batch rollout to finish, but will have the time to wait for active tasks to finish. + 5. As shown in figure c; + +4. async stream pipeline with partial rollout: + 1. **trigger_parameter_sync_step>=1, staleness_threshold>0, partial_rollout=True** + 2. Compared to c, when triggering parameter synchronization, if Rollouter has samples being produced, it will + interrupt the rollout process and perform parameter synchronization. The interrupted samples will continue to be + generated after synchronization. This reduces the time to wait for active tasks to finish. + 3. As shown in figure d; + +![fully_async_policy_mode]( +https://github.com/ArronHZG/verl-community/blob/main/docs/fully_async_policy_mode.svg?raw=true) + +### Key Metrics + +| metrics | implication | +|------------------------------------------------|--------------------------------------------------------------------------------------------------------| +| `trainer/idle_ratio` | Trainer idle rate | +| `rollouter/idle_ratio` | Rollouter idle rate | +| `fully_async/count/stale_samples_processed` | Total number of old samples used in training | +| `fully_async/count/stale_trajectory_processed` | Total number of old trajectories used in training (one sample produces rollout.n trajectories) | +| `fully_async/partial/total_partial_num` | Number of partial samples processed by Trainer between two trigger_parameter_sync_step | +| `fully_async/partial/partial_ratio` | Ratio of partial samples processed by Trainer between two trigger_parameter_sync_step | +| `fully_async/partial/max_partial_span` | Maximum parameter span of partial samples processed by Trainer between two trigger_parameter_sync_step | + +### Parameter Tuning Recommendations + +* Resource Allocation and Adjustment: + * Reasonable resource allocation is the prerequisite for achieving good training efficiency. The ideal resource + allocation should make the rollout time and train time close, thereby minimizing pipeline bubbles in the entire + training process, + avoiding resource idleness, and ensuring Trainer does not use old samples. In real training scenarios, resource + allocation can be adjusted based on the idle time of rollout and train during actual training, + which can be obtained from rollouter/idle_ratio and trainer/idle_ratio. If rollouter/idle_ratio is high and + trainer/idle_ratio is low, + Trainer resources should be increased and Rollouter resources should be reduced, and vice versa. + +* Key Parameters: + * staleness_threshold: Setting it too high will cause more old samples to be used, affecting model performance. It + is recommended to set it to less than 1. + * require_batches: The closer to 1, the closer to a pure streaming process, the smaller the training bubbles, and + the faster the acceleration effect that can be achieved in terms of speed, but it will affect the order of sample + processing; + * trigger_parameter_sync_step: The smaller the setting, the closer to on policy, but it will cause frequent + parameter synchronization. Long-tail samples waste resources that cannot be filled by short samples, resulting in + low resource utilization. + The larger the setting, the higher the computational efficiency, but the accuracy will be affected by off policy. + * rollout.test_freq: It will occupy Rollouter resources and is not recommended to be set too small. + +* Mode Selection: By adjusting different parameters, the Fully Async architecture supports optimization acceleration at + different levels, suitable for tasks in different scenarios. + * For small-scale tasks that need to ensure training stability and on-policy nature, and have low speed + requirements, the on policy pipeline mode (Mode 1) can be tried. + * For scenarios that need to improve training throughput but are sensitive to staleness, the stream off policy + pipeline mode can be tried. That is, by + setting trigger_parameter_sync_step>1 to improve training efficiency, but still maintaining the synchronization + mechanism (staleness_threshold=0) (Mode 2). + * For large-scale tasks with high training speed requirements and can tolerate a certain degree of off-policy and + staleness, setting staleness_threshold> + 0 and partial_rollout=True can improve training efficiency, using the async stream pipeline mode (Mode 3 or 4). + +### Quick Start + +```shell +rollout_mode="async" +rollout_name="vllm" # sglang or vllm +if [ "$rollout_mode" = "async" ]; then + export VLLM_USE_V1=1 + return_raw_chat="True" +fi + +train_prompt_bsz=0 +gen_prompt_bsz=1 +n_resp_per_prompt=16 +train_prompt_mini_bsz=32 +total_rollout_steps=$(((512*400))) +test_freq=10 +staleness_threshold=0 +trigger_parameter_sync_step=16 +partial_rollout=False + + +python -m verl.experimental.fully_async_policy.fully_async_main \ + train_batch_size=${train_prompt_bsz} \ + data.gen_batch_size=${gen_prompt_bsz} \ + data.return_raw_chat=${return_raw_chat} \ + actor_rollout_ref.rollout.n=${n_resp_per_prompt} \ + actor_rollout_ref.actor.strategy=fsdp2 \ + critic.strategy=fsdp2 \ + actor_rollout_ref.hybrid_engine=False \ + actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \ + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \ + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \ + actor_rollout_ref.rollout.name=${rollout_name} \ + actor_rollout_ref.rollout.mode=${rollout_mode} \ + trainer.nnodes="${NNODES_TRAIN}" \ + trainer.n_gpus_per_node="${NGPUS_PER_NODE}" \ + rollout.nnodes="${NNODES_ROLLOUT}" \ + rollout.n_gpus_per_node="${NGPUS_PER_NODE}" \ + rollout.total_rollout_steps="${total_rollout_steps}" \ + rollout.test_freq="${test_freq}" \ + async_training.staleness_threshold="${staleness_threshold}" \ + async_training.trigger_parameter_sync_step="${trigger_parameter_sync_step}" \ + async_training.partial_rollout="${partial_rollout}" +``` + +## Experiments + +### Asynchronous Training on 7B Model + +We used Qwen2.5-Math-7B to verify the benefits of the fully async strategy under long candidates and multiple resources. +Using the `async stream pipeline with stale samples` strategy, we achieved about 2x performance improvement on 32 cards, +64 cards, and 128 cards without significantly affecting experimental results. + +* Machine: H20 +* Model: Qwen2.5-Math-7B +* Rollout length: max_response_length FSDP2: 28K tokens; +* Algorithm: DAPO +* Dataset: TRAIN_FILE: dapo-math-17k.parquet TEST_FILE: aime-2024.parquet +* Engine: vllm+FSDP2 +* rollout.n: 16 +* ppo_mini_batch_size: 32 +* test_freq: 20 + +* colocate sync: + * step: 400 + * train_batch_size: 512 + +* fully_async_policy + * total_rollout_steps: 512*400 + * require_batches: 4 + * trigger_parameter_sync_step: 4 + * staleness_threshold: 0.5 + * partial_rollout: True + +| training mode | resource allocation | step | gen | old_log_prob | update_actor | total time
100 step | total time
200 step | total time
300 step | total time
400 step | acc/mean@1 | +|:--------------------:|:---------------------:|:--------:|:--------:|:--------------:|:---------------:|:------------------------:|:------------------------:|:------------------------:|:------------------------:|:-------------------------------:| +| colocate sync | 32 | 790.10 | 357.41 | 107.71 | 269.80 | 13h 44m | 1d 3h 43m | 2d 9h 22m | 3d 17h 5m | max: 0.3313
last: 0.2448 | +| fully_async_policy | 16:16 | 294.77 | 21.26 | \ | 313.81 | 7h 58m
(1.72x) | 16h 21m
(1.70x) | 1d 0h 53m
(2.31x) | 1d 9h 26m
(2.66x) | max: 0.3302
last: 0.2333 | +| colocate sync | 64 | 365.28 | 150.72 | 70.26 | 133.41 | 10h 22m | 20h 45m | 1d 7h 6m | 1d 17h 32m | max: 0.3365
last: 0.2333 | +| fully_async_policy | 32:32 | 189.26 | 28.46 | \ | 156.98 | 4h 57m
(2.09x) | 10h 14m
(2.03x) | 16h 58m
(1.83x) | 21h 40m
(1.92x) | max: 0.3677
last: 0.3406 | +| colocate sync | 128 | 356.30 | 177.85 | 53.92 | 113.81 | 8h 36m | 17h 56m | 1d 5h 6m | 1d 16h 48m | max: 0.3573
last: 0.2958 | +| fully_async_policy | 64:64 | 150.63 | 33.14 | \ | 113.16 | 3h 13m
(2.67x) | 6h 46m
(2.65x) | 10h 53m
(2.67x) | 17h 22m
(2.35x) | max: 0.3521
last: 0.3094 | + +> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-colocate_async?nw=nwuserhouzg + +### 128-card 7B Asynchronous Mode Experiment + +We used Qwen2.5-Math-7B to verify the effects of various modes supported by fully async. +We can see that the benefit brought by streaming is approximately 1.6x, and after combining staleness and +partial_rollout, the benefit reaches 2.35x. + +| mode | step | gen | old_log_prob | update_actor | total time
100 step | total time
200 step | total time
300 step | total time
400 step | acc/mean@1 | +|:-------------------------------------------------------------------------------------------------------:|:--------:|:--------:|:--------------:|:--------------:|:------------------------:|:------------------------:|:------------------------:|:------------------------:|:------------------------------:| +| colocate sync | 356.30 | 177.85 | 53.92 | 113.81 | 8h 36m | 17h 56m | 1d 5h 6m | 1d 16h 48m | max: 0.3573
last: 0.2958 | +| `stream off policy pipeline`
(+fully async: trigger_parameter_sync_step= 4,
require_batches= 4) | 231.34 | 128.47 | \ | 98.77 | 4h 25m | 9h 41m | 15h 2m | 1d 1h 53m | max: 0.2844
last: 0.2604 | +| `async stream pipeline with stale samples`
(+staleness_threshold=0.5) | | | | | | | | | | +| `async stream pipeline with partial rollout`
(+partial_rollout=True) | 150.63 | 33.14 | \ | 113.16 | 3h 13m | 6h 46m | 10h 53m | 17h 22m | max: 0.3521
last: 0.3094 | + +> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-stream_stale_partial?nw=nwuserhouzg + +### 128-card Stale Ablation Experiment + +Under the `async stream pipeline with partial rollout` mode, we verified the impact of staleness settings on training +efficiency. +We found that the larger the staleness, the more obvious the final gains. +We also noticed that the times for staleness values of 0.3 and 0.5 are quite close, because as the training steps +increase, the response length changes significantly, causing training instability. +Further analysis and optimization are needed for this issue. + +| staleness_threshold | step | gen | old_log_prob | update_actor | total time
100 step | total time
200 step | total time
300 step | total time
400 step | acc/mean@1 | +|:---------------------:|:--------:|:--------:|:--------------:|:--------------:|:------------------------:|:------------------------:|:------------------------:|:------------------------:|:-----------------------------:| +| 0 | 231.34 | 128.47 | \ | 98.77 | 4h 25m | 9h 41m | 15h 2m | 1d 1h 53m | max: 0.2844
last: 0.2604 | +| 0.1 | 171.30 | 58.17 | \ | 109.12 | 3h 53m | 8h 37m | 14h 25m | 19h 59m | max: 0.3542
last: 0.2979 | +| 0.3 | 146.11 | 38.88 | \ | 103.22 | 3h 18m | 6h 49m | 11h 40m | 17h 20m | max: 0.3469
last: 0.2865 | +| 0.5 | 150.63 | 33.14 | \ | 113.16 | 3h 13m | 6h 46m | 10h 53m | 17h 22m | max: 0.3521
last: 0.3094 | + +> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-stream_stale_partial?nw=nwuserhouzg + +### 128-card 7B require_batches Ablation Experiment + +In multiple tests, we found that the number of samples issued each time in streaming affects the response length during +training, which in turn affects training time. We verified the impact on results by modifying +`async_training.require_batches`. + +| require_batches | step | gen | old_log_prob | update_actor | total time
100 step | total time
200 step | total time
300 step | acc/mean@1 | +|:-----------------:|:--------:|:-------:|:--------------:|:--------------:|:------------------------:|:------------------------:|:------------------------:|:-----------------------------:| +| 1 | 203.47 | 30.88 | \ | 181.08 | 3h 31m | 8h 29m | 17h 36m | max: 0.349
last: 0.326 | +| 2 | 158.72 | 26.32 | \ | 128.08 | 3h 35m | 7h 38m | 13h 57m | max: 0.351
last: 0.3406 | +| 4 | 124.64 | 25.62 | \ | 95.06 | 3h 13m | 6h 46m | 10h 53m | max: 0.3521
last: 0.3521 | + +> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-ablation_require_batches?nw=nwuserhouzg + +### 30B Model Mode Experiment + +We achieved a 1.7x performance improvement with `async stream pipeline with staleness samples` strategy on the +Qwen3-30B-A3B-Base model compared to the colocate setup. It is worth noting that this is far from the upper limit of +performance gains achievable through asynchrony. Firstly, the comparative experiments used a maximum response length of +only 8k, which is much shorter than the 20k sequence length in previous experiments, resulting in a less pronounced +rollout tail effect. Secondly, we adopted a highly skewed resource allocation, with rollout using 96 GPUs and trainer +using 32 GPUs, which is not an optimal configuration. During the experiments, we observed that the current verl +implementation imposes certain constraints, such as requiring data to be evenly divisible by the number of GPUs, making +resource adjustment less flexible. Additionally, as asynchronous training and deployment accelerate, the performance gap +is gradually narrowing. Therefore, enabling more flexible resource allocation and dynamic resource adjustment in the +future will be our next focus. + +* Machine: H20 +* Model: Qwen3-30B-A3B-Base +* Rollout length: max_response_length : 8K tokens; +* Algorithm: GRPO +* Dataset: TRAIN_FILE: dapo-math-17k.parquet TEST_FILE: aime-2024.parquet +* Engine: vllm+Megatron +* rollout.n: 16 +* ppo_mini_batch_size: 128 +* test_freq: 20 + +* colocate sync: + * step:400 + * train_batch_size: 512 + +* fully_async_policy + * total_rollout_steps: 512*400 + * trigger_parameter_sync_step: 512/128 = 4 + * staleness_threshold: 0.5 + * partial_rollout: True + +| Training Mode | Resource Allocation | Step | Gen | Old Log Prob | Ref | Update Actor | Total Time 100 Step | Total Time 200 Step | Total Time 300 Step | Total Time 400 Step | Acc/Mean@1 | +|--------------------|---------------------|--------|--------|--------------|-------|--------------|---------------------|---------------------|---------------------|---------------------|-----------------------------| +| Colocate Sync | 128 | 497.89 | 348.05 | 28.73 | 20.86 | 86.27 | 13h 36m | 1d 3h 48m | 1d 19h 4m | 2d 11h 39m | max: 0.3500
last: 0.3208 | +| Fully Async Policy | 96:32 | 282.75 | 22.06 | \ | 50.05 | 206.63 | 6h 45m (2.01x) | 14h 48m (1.88x) | 1d 0h 9m (1.78x) | 1d 10h 41m (1.72x) | max: 0.3813
last: 0.3448 | + +> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-30B?nw=nwuserhouzg | | | + +### checkpoint-engine Ablation Experiment +We tested the single-step parameter synchronization time of the checkpoint-engine on three models: Qwen2.5-Math-7B, Qwen3-30B-A3B, and Qwen3-235B-A22B, using default checkpoint-engine configurations. All experiments were performed on H20 machines, and the Megatron engine was used for training. + +| model | trainer rank | rollout rank | checkpoint-engine | total sync time | +|:---------------:|:--------------:|:-------------:|:-------------------:|:-----------------:| +| Qwen2.5-Math-7B | 4 | 4 | False | 0.12s | +| Qwen2.5-Math-7B | 4 | 4 | True | 0.02s | +| Qwen3-30B-A3B | 16 | 16 | False | 15.76s | +| Qwen3-30B-A3B | 16 | 16 | True | 4.38s | +| Qwen3-235B-A22B | 64 | 64 | False | 58.57s | +| Qwen3-235B-A22B | 64 | 64 | True | 23.70s | + + +### use_trainer_do_validate Experiment +We tested the effect of setting `use_trainer_do_validate=True` on the training process. The results show that setting +this parameter to True can reduce the validation time overhead and trainer node idle time. +We used Qwen2.5-Math-7B to verify the benefits of `use_trainer_do_validate=True` on the training process, we achieved about 2x performance improvement on validation time, and the trainer node idle time is reduced by about 40%. + +* Machine: H20 +* Model: Qwen2.5-Math-7B +* Rollout length: max_response_length FSDP2: 10K tokens; +* Algorithm: DAPO +* Dataset: TRAIN_FILE: dapo-math-17k.parquet TEST_FILE: aime-2024.parquet +* Engine: vllm+FSDP2 +* rollout.n: 16 +* ppo_mini_batch_size: 32 +* test_freq: 10 + +* fully_async_policy + * total_rollout_steps: 512*400 + * require_batches: 4 + * trigger_parameter_sync_step: 4 + * staleness_threshold: 0.5 + * partial_rollout: True + +| training mode | resource allocation | step | gen | old_log_prob | update_actor | validate time | total time
50 step | acc/mean@2 | +|:------------------:|:-------------------:|:-------:|:-------:|:------------:|:------------:|:-------------:|:---------------------:|:----------:| +| colocate sync | 16 | 484.623 | 52.939 | 0 | 430.263 | 205.080 | 7h9m | 22.6 | +| fully_async_policy | 8:8 | 489.953 | 52.622 | 0 | 435.874 | 95.699 | 7h2m | 21.0 | + + +## Multi-Turn Tool Calling + +Referencing **recipe/retool** and **ToolAgentLoop**, we implemented **AsyncPartialToolAgentLoop**, a multi-turn +tool-calling loop that supports partial_rollout for **fully_async_policy**. + +### Core Design + +`AsyncPartialToolAgentLoop` inherits from `ToolAgentLoop` and is adapted for the asynchronous training mode of +`fully_async_policy`. When `partial_rollout=True`, the Rollouter interrupts ongoing generation tasks before +synchronizing parameters with the Trainer. `AsyncPartialToolAgentLoop` is capable of: + +1. **Interrupting Tasks**: Responding to an interrupt signal to save the current state. Currently, interruptions occur + during the `GENERATING` process or after other states have completed. +2. **Resuming Tasks**: Resuming execution from the saved state after parameter synchronization is complete, rather than + starting over. + +### How to Use + +RL training with multi-turn tool calling in `fully_async_policy` is similar to `recipe/retool`. It is enabled by +specifying `multi_turn` configurations in the config file. + +1. **SFT Stage**: First, the model should undergo SFT to learn how to follow tool-calling format instructions. +2. **Multi-turn Configuration**: In the `fully_async_policy` training configuration, set the following parameters: + ```yaml + actor_rollout_ref: + rollout: + multi_turn: + enable: True # AsyncPartialToolAgentLoop will be used by default in fully_async_policy mode + # Other multi_turn related configurations + ``` +3. **Async Parameters**: To improve efficiency, enable `partial_rollout` and `staleness_threshold` when using multi-turn + tool calling: + ```yaml + async_training: + partial_rollout: True + staleness_threshold: 0.5 + # Other async parameters + ``` +4. **Example**: See `recipe/fully_async_policy/shell/dapo_7b_async_retool.sh`. + +### Experimental Results + +To validate the performance of `fully_async_policy` on multi-turn tool-calling tasks, we compared it with the standard +`colocate` synchronous mode. Key parameter settings are as follows. + +* **SFT Model**: Based on `Qwen2.5-7B-Instruct`, trained for 6 epochs on the `ReTool-SFT` dataset +* **RL Algorithm**: DAPO +* **Dataset**: + * Train: `DAPO-Math-17k` + * Test: `aime_2025` +* **Resource and Mode Comparison**: + * `colocate sync`: 32 H20 gpus + * `fully_async_policy`: 16 gpus for Trainer + 16 gpus for Rollouter +* **Key Configurations**: + 1. **Tool Calling Configuration**: + * `multi_turn.enable: True` + * `multi_turn.max_user_turns: 16` + * `multi_turn.max_assistant_turns: 16` + * `multi_turn.tool_config_path: recipe/retool/sandbox_fusion_tool_config.yaml` + 2. **`colocate sync` Configuration**: + * `ppo_mini_batch_size: 16` + * `train_batch_size: 64` + 3. **`fully_async_policy` Configuration**: + * `ppo_mini_batch_size: 16` + * `trigger_parameter_sync_step: 4` + * `require_batches: 1` + * `staleness_threshold: 1` + * `partial_rollout: True` + +| training mode | Resource allocation | step | gen | old_log_prob | update_actor | total time
100 step | total time
200 step | aime_2025
acc/mean@30 | +|:--------------------:|:---------------------:|:---------:|:---------:|:--------------:|:--------------:|:------------------------:|:------------------------:|:-------------------------------:| +| colocate | 32 | 375.47 | 228.03 | 35.19 | 111.84 | 9h 46m | 22h 28m | start:0.1078
last:0.2056 | +| fully_async_policy | 16: 16 | 221.36 | 40.59 | \ | 179.58 | 6h 19m
(1.55x) | 14h 4m
(1.60x) | start:0.11
last:0.2044 | + +> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-multiturn-tool?nw=nwuserhouzg + +## Future Plans + +* GRPO experiments +* Megatron adaptation +* SGLang integration +* Transfer queue integration +* Asynchronous parameter synchronization +* AReaL asynchronous algorithm implementation +* TPPO algorithm implementation +* Multi-turn and Tool support diff --git a/verl/docs/advance/grafana_prometheus.md b/verl/docs/advance/grafana_prometheus.md new file mode 100644 index 0000000000000000000000000000000000000000..3b59f936728e2142df8765b6f886804069566cd9 --- /dev/null +++ b/verl/docs/advance/grafana_prometheus.md @@ -0,0 +1,193 @@ +# Use Prometheus and Grafana to Monitor Rollout + +**Author:** `https://github.com/meituan-search` + +Last updated: 12/05/2025. + +Monitor the rollout computation process using Prometheus and Grafana when using verl to enhance system observability and facilitate further performance optimization. + +We provide an additional training monitoring capability, leveraging Prometheus and Grafana to display rollout information during training and enhance system observability to facilitate further performance optimization. + +The system automatically configures Prometheus to scrape metrics from rollout servers, eliminating manual configuration steps. + +## Overview + +The figures below show the performance of Qwen235B on the AIME2024 dataset with a response length of 20k, where the emergence of a long-tail problem is clearly observable. + +![fully_async_policy_structure](https://github.com/ArronHZG/verl-community/blob/main/docs/grafana_validate.png?raw=true) + +The following figure presents the fully asynchronous training of the Qwen235B model. Here, resource idleness is distinctly noticeable, indicating that rollout resources can be reduced. + +![fully_async_policy_structure](https://github.com/ArronHZG/verl-community/blob/main/docs/grafana_fully_async_train.png?raw=true) + +Through the above two examples, we also illustrate the necessity of system observability. + +## Architecture Overview + +The overall workflow consists of the following steps: + +1. **Multi-node Ray Cluster Setup**: Start Ray cluster across multiple nodes with Grafana and Prometheus information configured in environment variables on the master node +2. **Start Grafana Service**: Launch Grafana on the master node for visualization of monitoring dashboards +3. **Start Prometheus Service**: Launch Prometheus on the master node for metrics collection and storage +4. **verl Async Rollout Mode**: verl uses async rollout mode to obtain rollout server ports and IP addresses +5. **Automatic Prometheus Configuration**: verl automatically rewrites the Prometheus configuration to add monitoring for rollout servers and notifies Prometheus to reload the configuration +6. **Metrics Collection**: After program execution, metrics can be viewed in Prometheus +7. **Dashboard Visualization**: Upload and view monitoring metrics in Grafana dashboards + +## Detailed Setup Steps + +### Step 1: Environment Variables and Start Ray Cluster + +First, set the necessary environment variables and start the Ray service. + +> Reference: [configure-manage-dashboard](https://docs.ray.io/en/latest/cluster/configure-manage-dashboard.html) + +```bash +# Master node environment variables +export GF_SERVER_HTTP_PORT=3000 # Grafana service default port (customizable) +export PROMETHEUS_PORT=9090 # Prometheus service default port (customizable) +export RAY_HEAD_PORT=6379 # Ray master node port (customizable) +export RAY_DASHBOARD_PORT=8265 # Ray dashboard default port (customizable) +export GRAFANA_PATHS_DATA=/tmp/grafana # Grafana data storage directory (customizable) +export RAY_GRAFANA_HOST="http://${master_ip}:${GF_SERVER_HTTP_PORT}" # Ray-associated Grafana address +export RAY_PROMETHEUS_HOST="http://${master_ip}:${PROMETHEUS_PORT}" # Ray-associated Prometheus address + +# Start Ray on master node +ray start --head --port=${RAY_HEAD_PORT} --dashboard-port=${RAY_DASHBOARD_PORT} + +# Start Ray on worker nodes +ray start --address={master_addr}:${RAY_HEAD_PORT} +``` + +**Verification:** Visit `http://master_ip:8265` to confirm Ray has started successfully. + +### Step 2: Start Grafana (Visualization Dashboard) + +Grafana is used to display metrics collected by Prometheus (such as cache hit rate, throughput, etc.): + +```bash +# Master node +nohup grafana-server \ + --config /tmp/ray/session_latest/metrics/grafana/grafana.ini \ + --homepath /usr/share/grafana \ + web > grafana.log 2>&1 & +``` + +**Verification:** Visit `http://master_ip:3000` to confirm Grafana has started successfully (default credentials: `admin/admin`). + +If you need to change the port, modify the `GF_SERVER_HTTP_PORT` environment variable, and grafana-server will automatically recognize it. + +### Step 3: Start Prometheus (Metrics Collection) + +Prometheus is responsible for scraping metrics from vLLM services and storing them as time-series data: + +```bash +# Master node +nohup prometheus \ + --config.file /tmp/ray/session_latest/metrics/prometheus/prometheus.yml \ + --web.enable-lifecycle \ + --web.listen-address=:${PROMETHEUS_PORT} \ + > prometheus.log 2>&1 & +``` + +**Verification:** Visit `http://master_ip:9090` to confirm Prometheus service has started successfully. + +### Step 4 & 5: Start verl Training + +Start verl training with the following parameters configured: + +**Required Configuration:** + +- `actor_rollout_ref.rollout.mode="async"` +- `actor_rollout_ref.rollout.disable_log_stats=False` +- `actor_rollout_ref.rollout.prometheus.enable=True` + +If use default port, this parameter can be omitted. + +- `actor_rollout_ref.rollout.prometheus.port=9090` + +If use default path, this parameter can be omitted. + +- `actor_rollout_ref.rollout.prometheus.file="/tmp/ray/session_latest/metrics/prometheus/prometheus.yml"` + +served_model_name uses `model_path.split("/")[-1]` for data statistics by default. +Users can also customize other aliases: + +- `actor_rollout_ref.rollout.prometheus.served_model_name="Qwen3-235B"` + +**Shell Script Example:** + +```bash +WORKING_DIR=${WORKING_DIR:-"${PWD}"} +RUNTIME_ENV=${RUNTIME_ENV:-"${WORKING_DIR}/verl/trainer/runtime_env.yaml"} + +rollout_mode="async" +rollout_name="vllm" # Options: sglang or vllm +if [ "$rollout_mode" = "async" ]; then + export VLLM_USE_V1=1 + return_raw_chat="True" +fi + +# Synchronous training +ray job submit --no-wait --runtime-env="${RUNTIME_ENV}" \ + --working-dir "${WORKING_DIR}" \ + -- python3 -m verl.trainer.main_ppo \ + data.return_raw_chat=${return_raw_chat} \ + actor_rollout_ref.rollout.name=${rollout_name} \ + actor_rollout_ref.rollout.mode=${rollout_mode} \ + actor_rollout_ref.rollout.disable_log_stats=False \ + actor_rollout_ref.rollout.prometheus.enable=True + ... + +# Asynchronous training +ray job submit --no-wait --runtime-env="${RUNTIME_ENV}" \ + --working-dir "${WORKING_DIR}" \ + -- python3 verl.experimental.fully_async_policy.fully_async_main \ + data.return_raw_chat=${return_raw_chat} \ + actor_rollout_ref.rollout.name=${rollout_name} \ + actor_rollout_ref.rollout.mode=${rollout_mode} \ + actor_rollout_ref.rollout.disable_log_stats=False \ + actor_rollout_ref.rollout.prometheus.enable=True + ... +``` + +### Step 6: View Metrics in Prometheus + +After task execution, verify that Prometheus is correctly collecting metrics. + +**Verification:** Visit the Prometheus interface at `http://master_ip:9090` and search for `vllm:` or `sglang:` to +confirm metrics are being reported correctly. + +**Troubleshooting:** + +If no metrics appear: + +1. Check logs for `AgentLoopManager` to find the server port +2. Visit `http://master_ip:server_port/metrics` to verify server metrics are available +3. Confirm that `actor_rollout_ref.rollout.disable_log_stats=False` is set + +### Step 7: View Metrics in Grafana + +After task execution, log in to Grafana to view and customize monitoring dashboards. + +**Login:** Visit `http://master_ip:3000` (default credentials: `admin/admin`) + +**Import Dashboard:** + +1. Select `Dashboards` → `New` → `Import` → `Upload dashboard JSON file` +2. Upload a pre-built dashboard JSON file + +**Available Dashboards:** + +- [vLLM Grafana Dashboard style 1](https://github.com/ArronHZG/verl-community/blob/main/docs/grafana/vllm_grafana.json) +- [vLLM Grafana Dashboard style 2](https://github.com/vllm-project/vllm/blob/main/examples/online_serving/dashboards/grafana/performance_statistics.json) +- [vLLM Grafana Dashboard style 2](https://github.com/vllm-project/vllm/blob/main/examples/online_serving/dashboards/grafana/query_statistics.json) +- [SGLang Grafana Dashboard](https://github.com/sgl-project/sglang/blob/main/examples/monitoring/grafana/dashboards/json/sglang-dashboard.json) + +## Additional Resources + +- [Ray Monitoring Documentation](https://docs.ray.io/en/latest/cluster/configure-manage-dashboard.html) +- [Prometheus Documentation](https://prometheus.io/docs/) +- [Grafana Documentation](https://grafana.com/docs/) +- [vLLM GitHub Repository](https://github.com/vllm-project/vllm) +- [SGLang GitHub Repository](https://github.com/sgl-project/sglang) diff --git a/verl/docs/advance/megatron_extension.rst b/verl/docs/advance/megatron_extension.rst new file mode 100644 index 0000000000000000000000000000000000000000..a96b15b408061d7c05368e8e3db0ba81f81fe3e2 --- /dev/null +++ b/verl/docs/advance/megatron_extension.rst @@ -0,0 +1,20 @@ +Add models with the Megatron-LM backend +========================================= + +Last updated: 04/25/2025. + +Model +----------- + + +If use latest verl, we have direct support of ``GPTModel`` for Megatron backend. +You can use the similar way of using Megatron to pretrain custom models. +We list the steps here: + +1. Find `model_initializer.py `_ +2. If your model is configurable by ``TransformerLayerSpec`` , you can + directly use ``GPTModel``. Otherwise, Please implement a new + ``ModelLayerSpec`` and ``ModelLayer`` here. +3. Use the right ``LayerSpec`` , ``TransformerConfig`` and ``HuggingfaceConfig`` + as arguments to initialize the GPTModel. +4. Return the model at last. diff --git a/verl/docs/advance/mtp.md b/verl/docs/advance/mtp.md new file mode 100644 index 0000000000000000000000000000000000000000..b1c4a77f117f631a6225ef997a3894213dd5b4b7 --- /dev/null +++ b/verl/docs/advance/mtp.md @@ -0,0 +1,112 @@ +# Guide to Using MTP in SFT/RL Training and Inference + +**Author**: `https://github.com/meituan-search` + +Last updated: 02/15/2026 + +## 1. Scope of Support + +Currently, RL training can be performed on mimo-7B-RL, Qwen-next, and Deepseek series models based on the MTP architecture. The support rules for training and inference engines are as follows: + +- **Training Engine**: Only supports the `mbridge/Megatron-Bridge + megatron` combination; other training engines are not compatible at this time; + +- **Inference Engine**: Compatible with all engines, but the model must be in the corresponding engine's compatibility list; + +- **Dependency Versions**: + + - mbridge: Apply the patches and review suggestions from PR: [#62](https://github.com/ISEEKYAN/mbridge/pull/62) (Already merged into the main branch); + + - Megatron-Bridge: Apply the patches and review suggestions from PR if you want to try out mimo-7B-RL: [#2387](https://github.com/NVIDIA-NeMo/Megatron-Bridge/pull/2387) (will be merged into the main branch in the future); + + - megatron: Use the latest dev version (commit: [23e092f41ec8bc659020e401ddac9576c1cfed7e](https://github.com/NVIDIA/Megatron-LM/tree/23e092f41ec8bc659020e401ddac9576c1cfed7e)), which supports MTP + CP training methods. + + - sglang: Use the specified branch: [https://github.com/ArronHZG/sglang/tree/fix_mtp_update_weights_from_tensor](https://github.com/ArronHZG/sglang/tree/fix_mtp_update_weights_from_tensor), [PR](https://github.com/sgl-project/sglang/pull/17870) , which fix the MTP update weights from tensor OOM issue. + +## 2. MTP Training Configuration (Core Parameters) + +The MTP training process can be flexibly controlled through the following configurations. All configurations are based on the `actor_rollout_ref.model.mtp` prefix: + +| Configuration Scenario | Core Parameters | Description | +|------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|---------------------------------------------------------| +| Load MTP Parameters Only | `enable=True` | VRAM usage will increase, but the exported parameters include the MTP module and can be directly used for online deployment | +| Full-Parameter MTP Training | `enable=True`
`enable_train=True`
`mtp_loss_scaling_factor=0.1` | MTP Loss will apply to all model parameters | +| MTP Parameter-Only Training | `enable=True`
`enable_train=True`
`detach_encoder=True` | Freeze the Encoder layer, update only MTP module parameters, MTP Loss applies only to MTP parameters | +| MTP Accelerated Rollout | 1. vLLM configuration:
`enable=True`
`enable_rollout=True`
`method="mtp"`
`num_speculative_tokens=1`
2. SGLang configuration:
`enable=True`
`enable_rollout=True`
`speculative_algorithm="EAGLE"`
`speculative_num_steps=2`
`speculative_eagle_topk=2`
`speculative_num_draft_tokens=4` | Achieve inference acceleration during the Rollout phase based on MTP | + +## 3. Experimental Results + +The experiment was conducted as follows: + +* model = mimo-7B-math +* max_response_length = 8k + +Experiment chart: + +![fully_async_policy_revenue]( +https://github.com/ArronHZG/verl-community/blob/main/docs/mimo-7b-mtp.png?raw=true) + +The wandb link for the graph: [wandb](https://wandb.ai/hou-zg-meituan/mimo-7b-sft-mtp?nw=nwuserhouzg) + +**Scenarios with No Significant Effect** + +The following configurations will not have a noticeable impact on training results: + +1. The base model does not carry MTP parameters; + +2. The base model carries MTP parameters, but the MTP module is not trained; + +3. The base model carries MTP parameters and trains MTP, with `mtp_loss_scaling_factor=0`; + +4. The base model carries MTP parameters, trains MTP and detaches the encoder, with `mtp_loss_scaling_factor=0.1`. + +**Scenarios with Significant Effect** + +Only the following configuration will have a noticeable impact on training results: + +- The base model carries MTP parameters, MTP Loss applies to all model parameters, and `mtp_loss_scaling_factor=0.1`. + +**Recommended Training Method** + +It is recommended to adopt the `detach_encoder=True` approach for MTP training. + +## 4. Performance Notes for MTP in Rollout Inference + +Enabling MTP improves the rollout acceptance rate by around 14%. However, on H20 GPUs, overall throughput does not increase and even decreases slightly. + +![spec_log]( +https://github.com/ArronHZG/verl-community/blob/main/docs/spec_log.png?raw=true) + +The effectiveness of MTP-accelerated Rollout is significantly affected by **model size** and **inference hardware**. Key reference information is as follows: + +**Hardware Tensor Core Performance** + +| Hardware Model | FP16 Performance (TFLOPS) | +|----------------|---------------------------| +| H20 | 148 | +| H800 | 1,671 | +| H200 | 1,979 | + +**Measured Performance and Recommendations** + +Taking the mimo-7B model deployed separately on H20 hardware using SGLang as an example: After enabling MTP speculative decoding, the Rollout throughput decreases by approximately 50%. + +- Current priority recommendation: Do not enable MTP acceleration during the inference phase for now; + +- Future planning: Further optimization of the speculative logic in the Rollout phase will be conducted to improve throughput performance. + +## 5. SFT training + +The SFT training with MTP is supported, using the same MTP training configuration as RL training. + +An example configuration for running SFT can be found in `examples/sft/gsm8k/run_mimo_7b_mtp_megatron.sh` + +**SFT result** + +The experiment was conducted using following data: +- model = mimo-7B-math +- dataset = gsm8k + +The result: [wandb link](https://wandb.ai/hou-zg-meituan/mimo-7b-sft-mtp?nw=nwuserhouzg) + +The presence of mtp layer has limited effect on main loss. However, when MTP layer is detached, the mtp_loss converges to a higher value. + diff --git a/verl/docs/advance/one_step_off.md b/verl/docs/advance/one_step_off.md new file mode 100644 index 0000000000000000000000000000000000000000..99170d75edc3112b5eba00ab562d8c2316acb9c0 --- /dev/null +++ b/verl/docs/advance/one_step_off.md @@ -0,0 +1,319 @@ +# Recipe: One Step Off Policy Async Trainer + +**Author:** `https://github.com/meituan-search` + +Last updated: 07/17/2025. + +## Introduction + +### Background + +The current reinforcement learning training process implemented by verl is synchronous, adhering to the algorithmic +workflows of established methods like PPO, GRPO, and DAPO. In each step, training samples are generated by the latest +model, and the model is updated after training completes. While this approach aligns with off-policy reinforcement +learning and stabilizes RL training, but it suffers from severe efficiency issues. +Model updates must wait for the longest output in the generation phase to complete. +During the generation of long-tail samples, GPUs remain idle, resulting in significant underutilization. +The more severe the long-tail problem in sample generation, the lower the overall training efficiency. +For example, in DAPO 32B training, the Rollout phase accounts for approximately 70% of the total time, +and increasing resources does not reduce the Rollout duration. + +![DAPO 32B Math Performance](https://raw.githubusercontent.com/eric-haibin-lin/verl-community/refs/heads/main/docs/dapo_32b_math.png) + +> source data: https://wandb.ai/verl-org/DAPO%20Reproduction%20on%20verl/workspace?nw=nwusertongyuxuan361 + +### Solution + +We have implemented the **One Step Off Async Trainer** to help alleviate this issue. This approach parallelizes the +generation and training processes, utilizing samples generated in the previous step for current training. +It also involves appropriately partitioning resources, allocating dedicated resources for generation while automatically +assigning the remainder to training. By reducing resources allocated to the generation phase, we mitigate GPU idle time +during long-tail sample generation. Throughout this process, generation and training parameters maintain a one-step off +policy. + +![One Step Off Policy Diagram](https://raw.githubusercontent.com/eric-haibin-lin/verl-community/refs/heads/main/docs/one_step_off_policy.png) + +> reference: [AReaL: A Large-Scale Asynchronous Reinforcement Learning System for Language Reasoning](https://arxiv.org/abs/2505.24298) + +Our core contributions include: + +1. **Parallel Generation and Training**: + Samples for the next batch are asynchronously generated while the current batch is being trained. + +2. **Resource Isolation**: + Unlike `hybrid_engine`, this method requires explicit resource allocation for rollout, with remaining resources + automatically assigned to training. + +3. **NCCL Parameter Synchronization**: + Employs NCCL communication primitives for seamless parameter transfer between generation and training modules. + +### Experimental Results + +- **Machine Configuration**: 2 nodes with 16 H20 GPUs each + - Generation: 4 GPUs + - Training: 12 GPUs +- **Model**: Qwen2.5-Math-7B +- **Rollout Configuration**: +- **Max Response Length**: FSDP2: 20,480 tokens; Megatron: 8,192 tokens +- **Algorithm**: DAPO +- **Rollout Engine**: vLLM + +| training mode | engine | step | gen | wait_prev_gen | generate_sequences | old_log_prob | update_actor | total time | acc/best@32/mean | acc/maj@32/mean | +| ---------------------- | ------------- | ---- | --- | ------------- | ------------------ | ------------ | ------------ | -------------- | ---------------- | --------------- | +| colocate sync | VLLM+FSDP2 | 749 | 321 | - | 247 | 88 | 286 | 19h18m | 0.5948 | 0.417 | +| one-step-overlap async | VLLM+FSDP2 | 520 | - | 45 | 458 | 108 | 337 | 15h34m（+23%） | 0.6165 | 0.494 | +| colocate sync | VLLM+Megatron | 699 | 207 | - | 162 | 119 | 344 | 18h21m | 0.605 | 0.4217 | +| one-step-overlap async | VLLM+Megatron | 566 | - | 59 | 501 | 120 | 347 | 13h06m (+40%) | 0.6569 | 0.4038 | + +- colocate sync: step ≈ gen + old_log_prob + update_actor +- one-step-overlap async: step ≈ wait_prev_gen + old_log_prob + update_actor + +![One Step Off Megatron Performance](https://raw.githubusercontent.com/eric-haibin-lin/verl-community/refs/heads/main/docs/one_step_off_megatron.png) + +> source data: https://wandb.ai/hou-zg-meituan/one-step-off-policy?nw=nwuserhouzg + +## Implementation + +### One Step Off Policy Async Pipeline + +Our implemented **One Step Off Policy Async Pipeline** integrates seamlessly into existing training logic at minimal +cost, +eliminating the need for additional sample storage management. The core mechanism uses `async_gen_next_batch` +for asynchronous rollout generation while maintaining continuous operation during epoch transitions +via `create_continuous_iterator`. + +```python +# iterator generator, simplify one-step integration of the training process +def _create_continuous_iterator(self): + for epoch in range(self.config.trainer.total_epochs): + iterator = iter(self.train_dataloader) + for batch_dict in iterator: + yield epoch, batch_dict + + +# read next batch samples, parameters sync and launch asyn gen_seq +def _async_gen_next_batch(self, continuous_iterator): + # read train_data + try: + epoch, batch_dict = next(continuous_iterator) + except StopIteration: + return None + batch = DataProto.from_single_dict(batch_dict) + gen_batch = batch_pocess(batch) + # sync weights from actor to rollout + self.sync_rollout_weights() + # async generation + gen_batch_output = self.rollout_wg.async_generate_sequences(gen_batch) + # future encapsulated + return GenerationBatchFuture(epoch, batch, gen_batch_output) + + +continuous_iterator = self._create_continuous_iterator() +# run rollout first to achieve one-step-off +batch_data_future = self._async_gen_next_batch(continuous_iterator) + +while batch_data_future is not None: + # wait for the gen_seq result from the previous step + batch = batch_data_future.get() + # launch the next async call to generate sequences + batch_data_future = self._async_gen_next_batch(continuous_iterator) + + # compute advantages + batch = critic.compute_values(batch) + batch = reference.compute_log_prob(batch) + batch = reward.compute_reward(batch) + batch = compute_advantages(batch) + + # model update + critic_metrics = critic.update_critic(batch) + actor_metrics = actor.update_actor(batch) +``` + +### Parameter Synchronization + +The exciting point is that our nccl based weights updating for rollout model has great performance. +At most of time, the latency is under 300ms, which is negligible for RLHF. + +> **sync_rollout_weights**：The time for synchronizing parameters from actor to rollout is extremely fast and can almost +> be ignored because it is implemented with nccl. + +```python +class ActorRolloutRefWorker: + # actor acquires the meta-info of model parameters for parameter sync + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def get_actor_weights_info(self): + params = self._get_actor_params() + ret = [] + for key, tensor in params.items(): + ret.append((key, tensor.size(), tensor.dtype)) + self._weights_info = ret + return ret + + # rollout sets the meta-info of model parameters for parameter sync + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def set_actor_weights_info(self, weights_info): + self._weights_info = weights_info + + +class AsyncRayPPOTrainer(RayPPOTrainer): + def init_workers(self): + ... + # rollout obtains the meta-info of model parameters from the actor for parameter sync + weights_info = self.actor_wg.get_actor_weights_info()[0] + self.rollout_wg.set_actor_weights_info(weights_info) + + # Create an actor-rollout communication group for parameter sync + self.create_weight_sync_group +``` + +```python +# The driving process invokes the actor and rollout respectively to create a weight synchronization group based on nccl/hccl. +def create_weight_sync_group(self): + master_address = ray.get(self.actor_wg.workers[0]._get_node_ip.remote()) + master_port = ray.get(self.actor_wg.workers[0]._get_free_port.remote()) + world_size = len(self.actor_wg.workers + self.rollout_wg.workers) + self.actor_wg.create_weight_sync_group( + master_address, + master_port, + 0, + world_size, + ) + ray.get( + self.rollout_wg.create_weight_sync_group( + master_address, + master_port, + len(self.actor_wg.workers), + world_size, + ) + ) + +# drive process call the actor and rollout respectively to sync parameters by nccl +def sync_rollout_weights(self): + self.actor_wg.sync_rollout_weights() + ray.get(self.rollout_wg.sync_rollout_weights()) + + +# fsdp model parameter sync +@register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False) +def sync_rollout_weights(self): + params = self._get_actor_params() if self._is_actor else None + if self._is_rollout: + inference_model = ( + self.rollout.inference_engine.llm_engine.model_executor.driver_worker.worker.model_runner.model + ) + from verl.utils.vllm.patch import patch_vllm_moe_model_weight_loader + patch_vllm_moe_model_weight_loader(inference_model) + # Model parameters are broadcast tensor-by-tensor from actor to rollout + for key, shape, dtype in self._weights_info: + tensor = torch.empty(shape, dtype=dtype, device=get_torch_device().current_device()) + if self._is_actor: + assert key in params + origin_data = params[key] + if hasattr(origin_data, "full_tensor"): + origin_data = origin_data.full_tensor() + if torch.distributed.get_rank() == 0: + tensor.copy_(origin_data) + from ray.util.collective import collective + + collective.broadcast(tensor, src_rank=0, group_name="actor_rollout") + if self._is_rollout: + inference_model.load_weights([(key, tensor)]) +``` + +### PPO Correctness + +To ensure the correctness of the PPO algorithm, we use rollout log_probs for PPO importance sampling. +For the related algorithm details, please refer to: https://verl.readthedocs.io/en/latest/algo/rollout_corr_math.html +The default mode is `bypass_ppo_clip`, but other modification strategies can also be explored. + +### AgentLoop + +In the current implementation, we no longer provide SPMD model rollout mode. +Instead, we have switched to AgentLoop mode, which also supports multi-turn tool calling. + +## Usage + +### FSDP2 Configuration Example + +```shell +python3 -m verl.experimental.one_step_off_policy.async_main_ppo \ + --config-path=config \ + --config-name='one_step_off_ppo_trainer.yaml' \ + actor_rollout_ref.actor.strategy=fsdp2 \ + # actor and rollout are placed separately + actor_rollout_ref.hybrid_engine=False \ + # actor and rollout resource + trainer.nnodes=1 \ + trainer.n_gpus_per_node=6 \ + rollout.nnodes=1 \ + rollout.n_gpus_per_node=2 +``` + +### Megatron Configuration Example + +```shell +python3 -m verl.experimental.one_step_off_policy.async_main_ppo \ + --config-path=config \ + --config-name='one_step_off_ppo_megatron_trainer.yaml' \ + actor_rollout_ref.actor.strategy=megatron \ + # actor and rollout are placed separately + actor_rollout_ref.hybrid_engine=False \ + # actor and rollout resource + trainer.nnodes=1 \ + trainer.n_gpus_per_node=6 \ + rollout.nnodes=1 \ + rollout.n_gpus_per_node=2 +``` + +### Configuration Guidelines + +1. **Card Number Relationships** + Maintain either of these relationships for optimal batch distribution: + + - `actor_rollout_ref.rollout.n` should be an integer divisor of: + `trainer.n_gpus_per_node * trainer.nnodes` + - `actor_rollout_ref.rollout.n * data.train_batch_size` should be evenly divisible by: + `trainer.n_gpus_per_node * trainer.nnodes` + + > Rationale: Ensures training samples can be evenly distributed across training GPUs when using partial resources for + > generation. + +2. **Dynamic Resource Tuning** + Adjust `trainer.nnodes` `trainer.n_gpus_per_node` `rollout.nnodes` `rollout.n_gpus_per_node` based on phase + durations: + - **Ideal state**: Rollout and training phases have comparable durations + - **Diagnostic metrics**: + - Monitor `wait_prev_gen` duration + - Analyze `sequence_length` distribution + - **Adjustment strategy**: + - High `wait_prev_gen` + uniform sequence lengths → Increase rollout resources + - High `wait_prev_gen` + long-tail sequences → Optimize stopping criteria (resource increase won't help) + > **wait_prev_gen**：The time consumed waiting for the previous rollout to end (the part that is not fully + > overlapped). + > **Resource Configuration Strategies:** + - **Resource-constrained scenario**: Optimize resource utilization by adjusting GPU allocation ratios, + keeping the number of nodes equal to allow training and rollout to share nodes; + - Configure `trainer.nnodes = rollout.nnodes` with + `trainer.n_gpus_per_node + rollout.n_gpus_per_node = physical_gpus_per_node`. Control rollout resource + allocation by adjusting `n_gpus_per_node`. + - **Resource-abundant scenario**: Optimize performance by adjusting the number of nodes, + keeping the number of GPUs per node equal to enable independent scaling of training and rollout + parallelism. + - Configure `trainer.n_gpus_per_node = rollout.n_gpus_per_node` and control rollout resource allocation by + adjusting `trainer.nnodes` and `rollout.nnodes`to achieve optimal performance. + > **Note**: The total number of nodes required by the system is not simply `trainer.nnodes + rollout.nnodes`. The + > actual calculation depends on GPU capacity: + > + > - When `trainer.n_gpus_per_node + rollout.n_gpus_per_node <= physical_gpus_per_node`, + > the required node count is `max(trainer.nnodes, rollout.nnodes)` + > - When `trainer.n_gpus_per_node + rollout.n_gpus_per_node > physical_gpus_per_node`, + > the required node count is `trainer.nnodes + rollout.nnodes` + +## Functional Support + +| Category | Support Situation | +| ------------------ | --------------------------------------------------------------------------------------------------------------- | +| train engine | FSDP2
Megatron | +| rollout engine | vLLM | +| AdvantageEstimator | GRPO
GRPO_PASSK
REINFORCE_PLUS_PLUS
RLOO
OPO
REINFORCE_PLUS_PLUS_BASELINE
GPG | +| Reward | all | diff --git a/verl/docs/advance/placement.rst b/verl/docs/advance/placement.rst new file mode 100644 index 0000000000000000000000000000000000000000..ae944ec3c53384c3dffbb10635948cea9c262cad --- /dev/null +++ b/verl/docs/advance/placement.rst @@ -0,0 +1,13 @@ +Ray API Design Tutorial +======================================= + +Last updated: 10/30/2024. + +We provide a tutorial for our Ray API design, including: + +- Ray basic concepts +- Resource Pool and RayWorkerGroup +- Data Dispatch, Execution and Collection +- Initialize the RayWorkerGroup and execute the distributed computation in the given Resource Pool + +See details in `tutorial.ipynb `_. \ No newline at end of file diff --git a/verl/docs/advance/ppo_lora.rst b/verl/docs/advance/ppo_lora.rst new file mode 100644 index 0000000000000000000000000000000000000000..1e902785f5438778901e013fa1358612de01726e --- /dev/null +++ b/verl/docs/advance/ppo_lora.rst @@ -0,0 +1,215 @@ +RL(HF) algorithms with LoRA Support +=========================================== + +Last updated: 02/03/2026. + +We support LoRA (Low-Rank Adaptation) for reinforcement learning algorithms such as PPO, GRPO, and others. + +LoRA is a parameter-efficient fine-tuning technique that injects trainable low-rank matrices into pre-trained weights (typically linear layers). This reduces memory footprint and compute cost, making it possible to fine-tune large models with limited hardware. + +The benefits this brings include: + +- reinforcement learning with very large models (e.g. 70B+) with modest hardware (e.g. 8x80G GPUs), +- enable larger batch sizes due to reduced memory usage, +- simplify model transfer and deployment, as only LoRA adapters need to be saved, +- Combine with techniques like `SLoRA `_ or `CCoE `_ to serve multiple LoRA adapters efficiently + +This guide explains how to enable LoRA in RL training and configure related parameters. + +FSDP Backend Usage Guide +------------------------ + +.. note:: + + This section applies to **FSDP/FSDP2 backend only**. For Megatron backend, see the :ref:`megatron-lora` section below. + +1. Lora is available in the `verl.trainer.ppo.ray_trainer.RayPPOTrainer`. Examples are provided via the `verl.trainer.main_ppo` entry point. + +2. LoRA is supported via huggingface peft with fsdp/fsdp2 and both vllm and sglang rollout backends. + +- `strategy=fsdp` or `strategy=fsdp2` +- `rollout.name=vllm` or `rollout.name=sglang` + +3. Required configurations for LoRA: + +- `actor_rollout_ref.model.lora_rank`: int, set to a reasonable value greater than 0 (e.g., 8, 16, 32, 64) +- `actor_rollout_ref.model.lora_alpha`: float, the alpha term in LoRA +- `actor_rollout_ref.rollout.load_format="safetensors"`: required. This enables vLLM to load the base model. +- `actor_rollout_ref.model.target_modules`: the target modules for LoRA. Typically set to "all-linear". + +4. Optional configurations for LoRA: + +- `actor_rollout_ref.model.lora_adapter_path`: string, path to a pretrained LoRA adapter directory. + If provided, loads existing adapter instead of creating new one. Enables multi-stage training from previously saved adapters. + Directory need contain `adapter_model.safetensors` and `adapter_config.json`. +- `actor_rollout_ref.model.lora.merge`: bool, whether to merge LoRA adapters into the base model weights before transferring to the rollout engine (vLLM or SGLang). + If True, LoRA adapters are merged into base weights and full merged weights are synced. If False, only LoRA adapter deltas are transferred natively. + For SGLang, ``merge=True`` is currently required. Native adapter loading (``merge=False``) for SGLang is planned. + +5. Recommend options: + +- `actor_rollout_ref.model.use_shm=True`: preload the model into `/dev/shm` to improve model loading speed. +- `actor_rollout_ref.rollout.layered_summon=True`: this enables the actor-model to gather the FSDP shards per layers when synchronizing the LoRA Adapter to vLLM, thereby reducing GPU peak memory. Recommended if the model is very large (70B+) or the GPU memory is limited (< 48GB) + +.. _megatron-lora: + +Megatron Backend Usage Guide +---------------------------- + +.. warning:: + + The FSDP-specific config options are **NOT applicable** to Megatron backend, and they will be ignored if set. Only options listed under ``lora`` key are applicable: + + - ``actor_rollout_ref.model.lora.*`` + - ``critic.model.lora.*`` + +You need to install and enable Megatron-Bridge for Megatron LoRA support. + +Make sure you use Megatron-Bridge later than 0.2.0, and we recommended using `this commit `_ or later for proper support, and use the following settings to enable Megatron-Bridge: + +- ``actor_rollout_ref.actor.megatron.use_mbridge=True`` +- ``actor_rollout_ref.actor.megatron.vanilla_mbridge=False`` + +**Key Differences from FSDP LoRA:** + +1. **LoRA Implementation**: Verl Megatron backend uses Megatron-Bridge's native LoRA implementation, which differs from HuggingFace PEFT. + +2. **Weight Sync / Refit Mechanism**: Currently, Megatron-Bridge can support syncing weights by either merging LoRA adapters into the base model weights before transferring to vLLM (for better inference speed but more refit time and potential precision loss), as well as loading separate adapters. + +**Configuration for Megatron LoRA:** + +.. code-block:: yaml + + actor_rollout_ref: + model: + lora: + # LoRA type: "lora", "vlm_lora", "canonical_lora", or "dora" + type: lora + + # whether to sync weights / refit by either merging LoRA adapters into the base model weights before transferring to vLLM (for better inference speed but more refit time and potential precision loss). If this is False, it will load separate adapters. + merge: False + + # LoRA rank (Dimension of the low-rank projection space.). Set to 0 to disable LoRA + rank: 0 + + # Weighting factor for the low-rank projection. Defaults to 32 + alpha: 32 + + # Dropout rate for the low-rank projection. Defaults to 0.0 + dropout: 0.0 + + # A list of module names to apply LoRA to. + # For fused LoRA, Defaults to all linear layers ['linear_qkv', 'linear_proj', 'linear_fc1', 'linear_fc2']. + # For canonical LoRA: ["linear_q", "linear_k", "linear_v", "linear_proj", "linear_fc1_up", "linear_fc1_gate", "linear_fc2"] + # - 'linear_qkv': Apply LoRA to the fused linear layer used for query, key, and value projections in self-attention + # - 'linear_proj': Apply LoRA to the linear layer used for projecting the output of self-attention + # - 'linear_fc1': Apply LoRA to the first fully-connected layer in MLP + # - 'linear_fc2': Apply LoRA to the second fully-connected layer in MLP + # Target modules can also contain wildcards. For example, you can specify + # target_modules=['*.layers.0.*.linear_qkv', '*.layers.1.*.linear_qkv'] to add LoRA to only linear_qkv on the first two layers + # + # Note: + # For MLA (e.g., DeepSeek), you should use ["linear_kv_down_proj","linear_kv_up_proj","linear_q_down_proj","linear_q_up_proj","linear_q_proj"] + # Instead of "linear_qkv" or ["linear_q","linear_k","linear_v"] + # By default, MoE routers are excluded from LoRA adaptation, and you will need to specify "router" in target_modules to include them. + target_modules: + - linear_qkv + - linear_proj + - linear_fc1 + - linear_fc2 + + # A list of module names not to apply LoRa to. It will match all nn.Linear & nn.Linear-adjacent modules whose name + # does not match any string in exclude_modules. If used, will require target_modules to be empty list or None + exclude_modules: [] + + # Position for applying dropout, can be 'pre' (before the low-rank projection) or 'post' (after). Defaults to 'pre' + dropout_position: pre + + # Initialization method for the low-rank matrix A. Defaults to "xavier". + lora_A_init_method: xavier + + # Initialization method for the low-rank matrix B. Defaults to "zero". + lora_B_init_method: zero + + # Enables the experimental All-to-All (A2A) communication strategy. Defaults to False + a2a_experimental: False + + # Parameter data type for LoRA weights. Default to null, which will use model's dtype. + dtype: null + + # Path to pre-trained LoRA adapter weights (null to train from scratch) + adapter_path: null + + # Whether to fully shard LoRA adapters. Defaults to False + # https://docs.vllm.ai/en/latest/api/vllm/config/lora/#vllm.config.lora.LoRAConfig.fully_sharded_loras + fully_sharded_loras: bool + + # VLMLoRA additionally allows the user to specify whether the language or vision models should be frozen. + # For example, a common finetuning workload for multimodal models is to apply adapters to language model and fully + # finetune the vision model. + freeze_vision_model: True + freeze_vision_projection: True + freeze_language_model: True + +LoRA training experiment with Qwen3-8B on 8 * H200 single node comparing FSDP and Megatron backend (script adapted from examples/tuning/lora/run_qwen3_8b_fsdp.sh): + +.. image:: https://github.com/user-attachments/assets/0482f423-01a3-4e52-a7ee-8b9cd79b7b1a +.. image:: https://github.com/user-attachments/assets/6ce10400-8164-47d8-90a6-c1bf002fb9e8 +.. image:: https://github.com/user-attachments/assets/092d3a43-4eba-425e-a584-8d83c1f02de4 + + +Best Practices and Notes +------------------------- + +1. **Learning rate**: it is recommended to increase the value of learning rate by an order of magnitude. + +2. **LoRA Rank**: + +- Too small a rank can hurt convergence. +- LoRA rank recommendation from @thelongestusernameofall: + + - A very small lora_rank can lead to slower convergence or worse training performance. It is recommended to set lora_rank to be>=32. Tests have shown that for a 0.5B model, with lora_rank=32,the training convergence speed and final performance are almost identical to non-LoRA training + - For a 32B model,with lora_rank=128,the training convergence speed and final performance are also almost identical to non-LoRA training. + - More comprehensive reference results are coming soon. + +.. image:: https://github.com/eric-haibin-lin/verl-community/blob/f2b80b8b26829124dd393b7a795a0640eff11644/docs/lora.jpg?raw=true + +3. **FSDP-Specific:** Reference configuration for RL training with the Qwen2.5-72B model using 8 x 80GB GPUs (increase lora_rank if needed): + +.. code-block:: + + data.train_batch_size=64 \ + actor_rollout_ref.model.use_shm=True \ + actor_rollout_ref.model.lora_rank=32 \ + actor_rollout_ref.model.lora_alpha=32 \ + actor_rollout_ref.model.target_modules=all-linear \ + actor_rollout_ref.actor.optim.lr=3e-5 \ + actor_rollout_ref.actor.fsdp_config.fsdp_size=8 \ + actor_rollout_ref.actor.fsdp_config.param_offload=True \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \ + actor_rollout_ref.rollout.tensor_model_parallel_size=8 \ + actor_rollout_ref.rollout.name=vllm \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \ + actor_rollout_ref.rollout.n=5 \ + actor_rollout_ref.rollout.max_num_seqs=64 \ + actor_rollout_ref.rollout.max_model_len=1536 \ + actor_rollout_ref.rollout.max_num_batched_tokens=1536 \ + actor_rollout_ref.rollout.load_format=safetensors \ + actor_rollout_ref.rollout.layered_summon=True \ + actor_rollout_ref.ref.fsdp_config.param_offload=True \ + actor_rollout_ref.actor.ulysses_sequence_parallel_size=1 \ + +Example Scripts +------------------- + +For end-to-end examples, refer to the scripts below: + +**FSDP Examples:** + +- LoRA training from scratch: examples/tuning/lora/run_qwen3_8b_fsdp.sh +- LoRA training from adapter path: examples/tuning/lora/run_qwen3_8b_from_adapter_fsdp.sh +- LoRA training for VLMs: examples/tuning/lora/run_qwen2_5_vl_7b_fsdp.sh + +**Megatron Examples:** + +- LoRA training with MoE: examples/tuning/lora/run_qwen3_30b_a3b_megatron.sh diff --git a/verl/docs/advance/reward_loop.rst b/verl/docs/advance/reward_loop.rst new file mode 100644 index 0000000000000000000000000000000000000000..36dedcd5d16528e9c328b5714c27a4d690407481 --- /dev/null +++ b/verl/docs/advance/reward_loop.rst @@ -0,0 +1,308 @@ +Reward Loop +=========== + +.. _yyding: https://yyding1.github.io + +Author: `Yuyang Ding `_ + +Last updated: 2/10/2026. + +Introduction +------------ + +Reward Loop is the default reward computation implementation in verl. +It is designed to support efficient, flexible, and easy-to-use reward computation. + +This document introduces the usage and architectural design. + +Key features include: + +1. **Distributed reward manager**, enabling scalable and efficient reward computation. +2. **Support for hybrid reward settings**, including both generative and discriminative reward models, as well as more complex reward scenarios. +3. **Simple and extensible interface**, for easily defining customized reward functions. + +Distributed Reward manager +-------------------------- + +.. image:: https://github.com/yyDing1/verl-materials/blob/main/distributed_reward_manager.svg?raw=true + +How distributed +~~~~~~~~~~~~~~~ + +Under the single_controller setup, actor rollout and reward computation can be abstracted as: + +.. code:: python + + # initalize rollout manager and async reward loop manager + async_rollout_manager = AgentLoopManager(config) + async_reward_manager = RewardLoopManager(config) + # actor rollout using `async_rollout_manager` + gen_batch = async_rollout_manager.generate_sequences(batch) + # compute reward using `async_reward_manager` + reward_batch = async_reward_manager.compute_rm_score(gen_batch) + +Within the ``RewardLoopManager``, multiple ``RewardWorker`` are launched across all nodes to enable distributed reward computation. +The number of parallel workers can be configured via ``config.reward.num_workers``. + +Upon receiving a batch reward request, the batch is partitioned into smaller chunks and distributed to each reward worker for parallel execution. +User only need to invoke ``compute_rm_score``. + +.. code:: python + + class RewardLoopManager: + """ + RewardLoopManager run in single controller. + This class will create reward loop workers and manage them. + """ + def _init_reward_loop_workers(self): + self.reward_loop_workers = [...] + + def compute_rm_score(self, data): + chunks = data.chunk(len(self.reward_loop_workers)) + outputs = ray.get( + [ + worker.compute_score_batch.remote(chunk) + for worker, chunk in zip(self.reward_loop_workers, chunks, strict=True) + ] + ) + outputs_flat = [item for sublist in outputs for item in sublist] + ... + +This is how the reward manager is parallelized and distributed across all nodes. + +Streaming Reward with Rollout +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Furthermore, we check whether actor rollout and reward computation can be performed in a streaming manner, +where the reward is calculated as soon as each sample is rolled out. + +.. code:: python + + # agent_reward_loop: streaming reward computation with actor rollout + # two conditions satisfied: (1) rule-based reward, or (2) reward model with extra resource pool + enable_agent_reward_loop = not use_rm or config.reward.reward_model.enable_resource_pool + + # if enable_agent_reward_loop, we directly pass reward_loop_workers to agent loop manager + # to stream reward computation with actor rollout + reward_loop_worker_handles = async_reward_manager.reward_loop_workers if enable_agent_reward_loop else None + async_rollout_manager = AgentLoopManager( + config=config, + worker_group=actor_rollout_wg, + rollout_resource_pool=actor_rollout_resource_pool, + reward_loop_worker_handles=reward_loop_worker_handles, + ) + +Hybrid Reward Scenarios Usage +----------------------------- + +As described above, each ``reward_loop_worker`` is responsible for handling reward requests. +The rewards can be categorized as follows: + +- **Rule-based Reward**: The reward is determined by predefined rules, e.g., checking whether the predicted answer matches the ground truth via string matching. +- **Discriminative Reward Model (DisRM)**: The reward is produced by a specified discriminative reward model, such as ``Skywork/Skywork-Reward-Llama-3.1-8B-v0.2``. +- **Generative Reward Model (GenRM)**: The reward is obtained using a generative reward model, for example ``dyyyyyyyy/FAPO-GenRM-4B``. +- **Hybrid Reward Scenarios**: A combination of the above reward types, e.g., rule + GenRM. + +.. code:: python + + class RewardLoopWorker: + + async def compute_score_batch(self, data: DataProto) -> list[dict]: + tasks = [] + for i in range(len(data)): + tasks.append(asyncio.create_task(self.compute_score(data[i : i + 1]))) + outputs = await asyncio.gather(*tasks) + return outputs + + async def compute_score(self, data: DataProto) -> dict: + if self.config.reward.custom_reward_function.path is not None: + # directly use user-customized reward function + return await self.reward_manager.run_single(data) + else: + if self.config.reward.reward_model.enable: + # we assume the rm is disrm + # genrm must set custom_reward_function + return await self.compute_score_disrm(data[-1:]) # only pass the last output to discriminative reward model + else: + return await self.reward_manager.run_single(data) + +Each ``RewardLoopWorker`` will initalize one ``RewardManager``, splits the batch into individual data items and processes them in parallel using asynchronous tasks. + +Reward Manager +~~~~~~~~~~~~~~ + +The ``RewardManager`` maintains a reward function and defines its computation logic, including: + +- **naive**: The simplest implementation. +- **dapo**: DAPO implementation with an overlong reward penalty. +- **limit**: Restricts the concurrency of the reward function, useful when external API calls are rate-limited. +- **remote**: Runs in a separate process, effective for CPU-intensive tasks such as ``Math-Verify``. + +Users can also customize their own ``RewardManager``, inheriting from ``RewardManagerBase``, and implementing the ``run_single`` function. + +.. code:: python + + @register("user_costomized") + class UserCostomizedRewardManager(RewardManagerBase): + async def run_single(self, data: DataProto) -> dict: + data_item = data[-1] + # your own reward manager + ... + +After defining it, users can specify their custom reward manager by setting ``reward.reward_manager.name=user_costomized``. + +When trajectories consist of multiple output sequences (currently supported only by the ``main_ppo_sync`` trainer) reward managers may consider all outputs when computing their scores. +In that case, the ``data`` argument passed to ``run_single`` will contain all outputs in the trajectory. +However, the default reward managers (e.g. ``naive``, ``dapo``, etc.) will only consider the last sequence by default, as they are typically designed for single-output tasks. +The same is true in the ``UserCostomizedRewardManager`` example above, as indicated by the line ``data_item = data[-1]``. + + +Rule-Based Reward +~~~~~~~~~~~~~~~~~ + +If ``reward.custom_reward_function`` is provided, the user-defined reward function will be used. Otherwise, it falls back to the default reward function. + +Note that The custom function can be either synchronous or asynchronous; the system automatically detects its type and loads it accordingly. + +We recommend **using asynchronous functions** when reward computation need to involve external model API calls or sandboxed execution, as they are significantly more efficient. + +.. code:: python + + async def compute_score(data_source, solution_str, ground_truth, extra_info): + """Compute a score by sending an async request to a remote service.""" + + # prepare request payload + payload = {"messages": [{"role": "user", "content": "check the correcness of the question and response ..."}], ...} + + # send async HTTP request + async with aiohttp.ClientSession() as session: + async with session.post("https://api.openai.com/v1/chat/completions", json=payload) as resp: + result = await resp.json() + + # parse and return score + score = int(result["choices"][0]["message"]["content"].strip().split("\n")[-1]) + return {"score": score} + +Model-Base Reward +~~~~~~~~~~~~~~~~~ + +**For discriminative reward model (DisRM)**, we provide a simple implementation: + +.. code:: python + + class RewardLoopWorker: + async def compute_score_disrm(self, data) -> dict: + disrm_prompt = await self._preprocess_reward_inputs(data) + payloads = { + "model": model_name, + "input": disrm_prompt, + "activation": False, + } + output = await self._post_request(payloads, "classify") + rm_score = output["data"][-1]["probs"][-1] + return {"reward_score": rm_score} + +pass the question and the model rollout as inputs to the reward model and obtain a reward score. This is also the standard practice for most DisRM. + +Users should provide ``reward.reward_model.model_path`` to specify the reward model. + +**For generative reward model (GenRM)** + +For generative reward model scenarios, users need to specify both ``reward.reward_model.model_path`` and ``reward.custom_reward_function``. + +The custom reward function should implement the following components: + +- Convert the question and the model rollout into a GenRM input prompt using a custom prompt template. +- Invoke the GenRM to perform generation with custom sampling parameters. For this purpose, the Reward Loop provides an HTTP interface (i.e., ``reward_router_address``) for interacting with GenRM. +- Parse the GenRM output using a custom parser and extract the reward score. + +As these steps are highly customizable and task-dependent, we offer this flexibility entirely to the user-defined reward function. + +Below we provide an example of a custom reward function using GenRM. + +.. code:: python + + async def compute_score_gsm8k( + data_source: str, + solution_str: str, + ground_truth: str, + extra_info: dict, + reward_router_address: str, # an HTTP router endpoint provided by Reward Loop + reward_model_tokenizer: PreTrainedTokenizer, + ): + """Compute the reward score.""" + + # Step 1: Prepare prompt and request payload + grm_prompt = GRM_PROMPT_TEMPLATE.format(problem=extra_info["question"], solution=solution_str) + messages = [{"role": "user", "content": grm_prompt}] + sampling_params = {"temperature": 0.7, "top_p": 0.8, "max_tokens": 4096} + chat_complete_request = {"messages": messages, **sampling_params} + + # Step 2: Send async request to the reward model + # here, chat_complete sends async http request to the router address + result = await chat_complete( + router_address=reward_router_address, + chat_complete_request=chat_complete_request, + ) + + # Step 3: Parse model response and extract score + grm_response = result.choices[0].message.content.strip() + try: + score_str = grm_response.split("\n\n")[-1].strip() + score = int(score_str) + except Exception: + score = 0 + + return {"score": score} + +**For hybrid reward scenarios**, such as combining rule-based rewards with GenRM similarly as above, + +.. _recipe/fapo: https://github.com/verl-project/verl-recipe/tree/main/fapo + +A runnable and reproducible example that demonstrates how to use a rule-based reward function together with a GenRM is provided in the `recipe/fapo`_ directory for reference. Welcome to use and cite. + +Reward Model Arch Design +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +We support multiple execution modes for reward models during: + +- **Colocate Mode**: The reward model shares the same resource pool as the actor/rollout/reference models. In this setup, all rollouts must complete first, after which the reward model is awakened to perform inference. +- **Standalone Mode**: The reward model runs on a separate resource pool, independent from the actor/rollout/reference models. In this setup, each sample is evaluated by the reward model immediately after its rollout finishes. + +The standalone mode can enable the streaming manner stated above. + +By default, the system runs in colocate mode. Users can enable standalone mode by setting ``reward.reward_model.enable_resource_pool=True`` and allocating the corresponding resources via ``reward.reward_model.nnodes`` and ``reward.reward_model.n_gpus_per_node``. + +.. image:: https://github.com/yyDing1/verl-materials/blob/main/reward_loop.svg?raw=true + + +To support flexible and scalable reward model computation, we implement a reward router that coordinates requests among multiple reward model servers. + +Each reward model runs as an independent server and is registered with the router. +This router will forward the requests to the registered reward servers with load balancing and return the results. +This design allows us to expose a single unified router address to user-defined reward functions, enabling them to access various reward models seamlessly through the same interface. + +.. image:: https://github.com/yyDing1/verl-materials/blob/main/reward_loop_full.svg?raw=true + +.. code:: python + + class RewardModelManager: + """Reward model manager.""" + + def __init__( + self, + config: RewardModelConfig, + resource_pool: RayResourcePool = None, + ): + """ + Initialize the reward model manager. + + Args: + config (RewardModelConfig): Reward model configuration. + resource_pool (RayResourcePool, optional): Resource pool. Defaults to None. + """ + self.config = config + self.resource_pool = resource_pool + self._initialize_llm_servers() + self._initialize_router() + diff --git a/verl/docs/advance/rollout_skip.rst b/verl/docs/advance/rollout_skip.rst new file mode 100644 index 0000000000000000000000000000000000000000..6ed661d2bad1750e60bf7d06f31d64de6bd16c24 --- /dev/null +++ b/verl/docs/advance/rollout_skip.rst @@ -0,0 +1,73 @@ +RolloutSkip Function Usage Documentation +======================================== + +Last updated: 2026-03-25 + +Applicable Scenarios +-------------------- +The RolloutSkip utility accelerates RL training by caching and reusing pre-generated rollout data, +avoiding redundant sequence generation during debugging, replay, or fixed-experiment runs. + +It is suitable for: + +1. Re-running experiments with the same configuration +2. Speeding up training by skipping repeated generation +3. Reproducing rollout results in debugging + + +API and Usage Example +---------------------- + +Trainer Adaptation +~~~~~~~~~~~~~~~~~~ +RolloutSkip is already supported in ``RayDAPOTrainer`` and ``RayPPOTrainer``. + +Example integration: + +.. code-block:: python + + from verl.utils.rollout_skip import RolloutSkip + + # Inside trainer.fit() + rollout_skip = RolloutSkip(self.config, self.async_rollout_manager) + rollout_skip.wrap_generate_sequences() + + +Basic Configuration +~~~~~~~~~~~~~~~~~~~ +Add these parameters to enable RolloutSkip: + +.. code-block:: bash + + actor_rollout_ref.rollout.skip.enable=True + actor_rollout_ref.rollout.skip.dump_dir=/path/to/rollout_dump + actor_rollout_ref.rollout.skip.max_dump_step=10 + + +Configuration Parameters +------------------------ +- **skip.enable**: Enable or disable RolloutSkip. +- **skip.dump_dir**: Root directory to save cached rollout data. +- **skip.max_dump_step**: Maximum number of steps to cache. + + +Cached Directory Structure +-------------------------- +The directory structure is automatically generated to isolate different experiments: + +.. code-block:: text + + {dump_dir}/{exp_name}_{project_name}/ + └── GBS{gbs}_N{n}_in{prompt_len}_out{response_len}/ + ├── train_step__gen_step.txt + ├── genstep_000001/ + │ ├── new_batch.dp + │ ├── gen_batch.dp + │ └── meta.json + └── genstep_000002/ + + +Each ``genstep_*`` folder contains: +- ``new_batch.dp``: Input prompt batch +- ``gen_batch.dp``: Generated response batch +- ``meta.json``: Step metadata \ No newline at end of file diff --git a/verl/docs/advance/rollout_trace.rst b/verl/docs/advance/rollout_trace.rst new file mode 100644 index 0000000000000000000000000000000000000000..5801353cb8c64ed741e0f2ecc54c4d5c0300f260 --- /dev/null +++ b/verl/docs/advance/rollout_trace.rst @@ -0,0 +1,146 @@ +Trace Function Usage Instructions +======================================== + +Last updated: 07/10/2025. + +Applicable Scenarios +-------------------- + +Agentic RL involves multiple turns of conversations, tool invocations, and user interactions during the rollout process. During the Model Training process, it is necessary to track function calls, inputs, and outputs to understand the flow path of data within the application. The Trace feature helps, in complex multi-round conversations, to view the transformation of data during each interaction and the entire process leading to the final output by recording the inputs, outputs, and corresponding timestamps of functions, which is conducive to understanding the details of how the model processes data and optimizing the training results. + +The Trace feature integrates commonly used Agent trace tools, including wandb weave and mlflow, which are already supported. Users can choose the appropriate trace tool according to their own needs and preferences. Here, we introduce the usage of each tool. + + +Trace Parameter Configuration +----------------------------- + +- ``actor_rollout_ref.rollout.trace.backend=mlflow|weave`` # the trace backend type +- ``actor_rollout_ref.rollout.trace.token2text=True`` # To show decoded text in trace view +- ``actor_rollout_ref.rollout.trace.max_samples_per_step_per_worker=N`` # Limit traces per worker (optional) + +Limiting Trace Volume +~~~~~~~~~~~~~~~~~~~~~~ + +By default, all samples are traced, which can generate large amounts of data and incur significant costs with trace backends like Weave or MLflow. To limit trace volume while maintaining representative coverage, use ``max_samples_per_step_per_worker``. + +Example configuration: + +.. code-block:: yaml + + actor_rollout_ref: + rollout: + trace: + backend: weave + token2text: False + max_samples_per_step_per_worker: 5 # Each worker traces 5 random samples + +Each agent loop worker independently selects up to N unique samples to trace per training step. For GRPO (``n > 1``), all rollouts for selected samples are traced. Total traces per step = max_samples_per_step_per_worker * num_workers * n. + +Example: With 4 workers, max_samples_per_step_per_worker=5, and GRPO n=4, you get 4 * 5 * 4 = 80 traces per step instead of tracing all samples. Set to null (default) to trace all samples. + + +Glossary +-------- + ++----------------+------------------------------------------------------------------------------------------------------+ +| Object | Explaination | ++================+======================================================================================================+ +| trajectory | A complete multi-turn conversation includes: | +| | 1. LLM output at least once | +| | 2. Tool Call | ++----------------+------------------------------------------------------------------------------------------------------+ +| step | The training step corresponds to the global_steps variable in the trainer | ++----------------+------------------------------------------------------------------------------------------------------+ +| sample_index | The identifier of the sample, defined in the extra_info.index of the dataset. It is usually a number,| +| | but may also be a uuid in some cases. | ++----------------+------------------------------------------------------------------------------------------------------+ +| rollout_n | In the GROP algorithm, each sample is rolled out n times. rollout_n represents the serial number of | +| | the rollout. | ++----------------+------------------------------------------------------------------------------------------------------+ +| validate | Whether the test dataset is used for evaluation? | ++----------------+------------------------------------------------------------------------------------------------------+ + +Rollout trace functions +----------------------- + +There are 2 functions used for tracing: + +1. ``rollout_trace_op``: This is a decorator function used to mark the functions to trace. In default, only few method has it, you can add it to more functions to trace more infor. +2. ``rollout_trace_attr``: This function is used to mark the entry of a trajectory and input some info to trace. If you add new type of agent, you may need to add it to enable trace. + + +Usage of wandb weave +-------------------- + +1.1 Basic Configuration +~~~~~~~~~~~~~~~~~~~~~~~ + +1. Set the ``WANDB_API_KEY`` environment variable +2. Configuration Parameters + + 1. ``actor_rollout_ref.rollout.trace.backend=weave`` + 2. ``trainer.logger=['console', 'wandb']``: This item is optional. Trace and logger are independent functions. When using Weave, it is recommended to also enable the wandb logger to implement both functions in one system. + 3. ``trainer.project_name=$project_name`` + 4. ``trainer.experiment_name=$experiment_name`` + 5. ``actor_rollout_ref.rollout.mode=async``: Since trace is mainly used for agentic RL, need to enable agent toop using async mode for either vllm or sglang. + +Note: +The Weave Free Plan comes with a default monthly network traffic allowance of 1GB. During the training process, the amount of trace data generated is substantial, reaching dozens of gigabytes per day, so it is necessary to select an appropriate wandb plan. + + +1.2 View Trace Logs +~~~~~~~~~~~~~~~~~~~ + +After executing the training, on the project page, you can see the WEAVE sidebar. Click Traces to view it. + +Each Trace project corresponds to a trajectory. You can filter and select the trajectories you need to view by step, sample_index, rollout_n, and experiment_name. + +After enabling token2text, prompt_text and response_text will be automatically added to the output of ToolAgentLoop.run, making it convenient to view the input and output content. + +.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/weave_trace_list.png?raw=true + +1.3 Compare Trace Logs +~~~~~~~~~~~~~~~~~~~~~~ + +Weave can select multiple trace items and then compare the differences among them. + +.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/weave_trace_compare.png?raw=true + +Usage of mlflow +--------------- + +1. Basic Configuration +~~~~~~~~~~~~~~~~~~~~~~ + +1. Set the ``MLFLOW_TRACKING_URI`` environment variable, which can be: + + 1. Http and https URLs corresponding to online services + 2. Local files or directories, such as ``sqlite:////tmp/mlruns.db``, indicate that data is stored in ``/tmp/mlruns.db``. When using local files, it is necessary to initialize the file first (e.g., start the UI: ``mlflow ui --backend-store-uri sqlite:////tmp/mlruns.db``) to avoid conflicts when multiple workers create files simultaneously. + +2. Configuration Parameters + + 1. ``actor_rollout_ref.rollout.trace.backend=mlflow`` + 2. ``trainer.logger=['console', 'mlflow']``. This item is optional. Trace and logger are independent functions. When using mlflow, it is recommended to also enable the mlflow logger to implement both functions in one system. + 3. ``trainer.project_name=$project_name`` + 4. ``trainer.experiment_name=$experiment_name`` + + +2. View Log +~~~~~~~~~~~ + +Since ``trainer.project_name`` corresponds to Experiments in mlflow, in the mlflow view, you need to select the corresponding project name, then click the "Traces" tab to view traces. Among them, ``trainer.experiment_name`` corresponds to the experiment_name of tags, and tags corresponding to step, sample_index, rollout_n, etc., are used for filtering and viewing. + +For example, searching for ``"tags.step = '1'"`` can display all trajectories of step 1. + +.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/mlflow_trace_list.png?raw=true + +Opening one of the trajectories allows you to view each function call process within it. + +After enabling token2text, prompt_text and response_text will be automatically added to the output of ToolAgentLoop.run, making it convenient to view the content. + +.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/mlflow_trace_view.png?raw=true + +Note: + +1. mlflow does not support comparing multiple traces +2. rollout_trace can not associate the mlflow trace with the run, so the trace content cannot be seen in the mlflow run logs. diff --git a/verl/docs/advance/rope.rst b/verl/docs/advance/rope.rst new file mode 100644 index 0000000000000000000000000000000000000000..9463549e47d055552a273e83a851fc76f93f9d1a --- /dev/null +++ b/verl/docs/advance/rope.rst @@ -0,0 +1,39 @@ +RoPE Scaling override +======================================= + +Last updated: 05/14/2025. + +Some models such as `Qwen/Qwen2.5-7B-Instruct `_ support RoPE Scaling but don't have it defined in their config.json file. +For example, this model supports this configuration: + +.. code:: python + + { + ..., + "rope_scaling": { + "factor": 4.0, + "original_max_position_embeddings": 32768, + "type": "yarn" + } + } + + + +In order to support a longer context for such models, you must override the model configs when starting the trainer. + +PPO example: + +.. code:: bash + + +actor_rollout_ref.model.override_config.rope_scaling.type=yarn \ + +actor_rollout_ref.model.override_config.rope_scaling.factor=4.0 \ + +actor_rollout_ref.model.override_config.rope_scaling.original_max_position_embeddings=32768 \ + + +And for the critic model + +.. code:: bash + + +critic.model.override_config.rope_scaling.type=yarn \ + +critic.model.override_config.rope_scaling.factor=4.0 \ + +critic.model.override_config.rope_scaling.original_max_position_embeddings=32768 \ diff --git a/verl/docs/algo/baseline.md b/verl/docs/algo/baseline.md new file mode 100644 index 0000000000000000000000000000000000000000..a907c8f52ccea7ed348ad72326cc22707c3e2e30 --- /dev/null +++ b/verl/docs/algo/baseline.md @@ -0,0 +1,73 @@ +# Algorithm Baselines + +Last updated: 06/18/2025. + +## Math related datasets + +### GSM8k + +Assuming GSM8k/math dataset is preprocessed via: + +```bash +python3 examples/data_preprocess/*.py +``` + +Refer to the table below to reproduce RL training from different pre-trained checkpoints. Below is the performance on the GSM8k dataset if not specified otherwise. More comprehensive benchmark results areavailable in the recipe folder. + +| Hardware | Model | Method | Test score | Details | +| ---------- | -------------------------------- | --------------- | ------------ | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| NVIDIA GPU | google/gemma-2-2b-it | hf checkpoint | 23.9 | [Huggingface](https://huggingface.co/google/gemma-2-2b-it#benchmark-results) | +| NVIDIA GPU | google/gemma-2-2b-it | SFT | 52.06 | [command and logs](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/gemma-2-2b-it-sft-0.411.log) | +| NVIDIA GPU | google/gemma-2-2b-it | SFT + PPO | 64.02 | [command and logs](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/gemma-2-2b-it-ppo-bsz512_4-prompt1024-resp-512-0.640.log), [wandb](https://api.wandb.ai/links/verl-team/h7ux8602) | +| NVIDIA GPU | Qwen/Qwen2.5-0.5B-Instruct | hf checkpoint | 49.6 | [Qwen blog](https://qwen.ai/blog?id=qwen2.5-llm) | +| NVIDIA GPU | Qwen/Qwen2.5-0.5B-Instruct | PPO | 56.7 | [command and log](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/Qwen2.5-0.5B-bsz256_2-prompt1024-resp512-0.567.log) | +| NVIDIA GPU | Qwen/Qwen2.5-0.5B-Instruct | PRIME | 58.7 | [script](https://github.com/verl-project/verl-recipe/blob/main//prime/run_prime_qwen.sh), [wandb](https://api.wandb.ai/links/zefan-wang-thu-tsinghua-university/rxd1btvb) | +| NVIDIA GPU | Qwen/Qwen2.5-0.5B-Instruct | GRPO-LoRA | 54.3 | [command and logs](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/Qwen2.5-0.5B-bsz64_2-prompt512-resp1024-lorarank32-score0.543.log) | +| NVIDIA GPU | Qwen/Qwen2.5-1.5B-Instruct | GRPO-LoRA | 77.9 | [command and logs](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/Qwen2.5-1.5B-bsz64_2-prompt512-resp1024-lorarank32-score0.779.log) | +| NVIDIA GPU | Qwen/Qwen2.5-3B-Instruct | GRPO-LoRA | 86.1 | [command and logs](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/Qwen2.5-3B-bsz64_2-prompt512-resp1024-lorarank32-score0.861.log) | +| NVIDIA GPU | deepseek-ai/deepseek-llm-7b-chat | PPO (Megatron) | 69.5 [1] | [log](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/deepseek-llm-7b-chat-megatron-bsz256_4-prompt512-resp512-0.695.log), [wandb](https://wandb.ai/verl-team/verl_megatron_gsm8k_examples/runs/10fetyr3) | +| NVIDIA GPU | Qwen/Qwen2-7B-Instruct | GRPO | 89 | [script](https://github.com/verl-project/verl/blob/a65c9157bc0b85b64cd753de19f94e80a11bd871/examples/grpo_trainer/run_qwen2-7b_seq_balance.sh) | +| NVIDIA GPU | Qwen/Qwen2-7B-Instruct | GRPO (FSDP2) | 89.8 | [log](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/qwen2-7b-fsdp2.log) | +| NVIDIA GPU | Qwen/Qwen2-7B-Instruct | GRPO (Megatron) | 89.6 | [log](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/qwen2-7b_math_megatron.log) | +| NVIDIA GPU | Qwen/Qwen2.5-7B-Instruct | ReMax | 97 | [script](https://github.com/eric-haibin-lin/verl/blob/main/examples/remax_trainer/run_qwen2.5-3b_seq_balance.sh), [wandb](https://wandb.ai/liziniu1997/verl_remax_example_gsm8k/runs/vxl10pln) | +| NVIDIA GPU | Qwen/Qwen2.5-7B-Instruct | SPPO | 65.6 (MATH) | [SPPO script](https://github.com/verl-project/verl-recipe/tree/main/sppo/README.md) | +| NVIDIA GPU | Qwen/Qwen2.5-7B-Instruct | GRPO-LoRA | 93.4 | [command and logs](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/Qwen2.5-7B-bsz64_8-prompt512-resp1024-lorarank32-score0.934.log) | +| NVIDIA GPU | Mixtral-8x22B-Instruct-v0.1 | Instruct model | 83.7 | [Qwen Blog](https://qwen.ai/blog?id=qwen2.5-llm) | +| NVIDIA GPU | Mixtral-8x22B-Instruct-v0.1 | RLOO (Megatron) | 92.3 | [wandb](https://api.wandb.ai/links/ppo_dev/sbuiuf2d) | +| NVIDIA GPU | Qwen/Qwen2.5-7B-Instruct | SPIN | 92 | [script](https://github.com/verl-project/verl-recipe/tree/main/spin/README.md) | +| NVIDIA GPU | Qwen/Qwen2-7B-Instruct | GPG | 88 | [log](https://github.com/diqiuzhuanzhuan/verldata/blob/main/run_logs/qwen2-7b_math.log), [wandb](https://wandb.ai/diqiuzhuanzhuan/verl_gpg_example_gsm8k_math/runs/ab86c4va) | +| NVIDIA GPU | Qwen/Qwen2-7B-Instruct | GPG (Megatron) | 88 | [log](https://github.com/diqiuzhuanzhuan/verldata/blob/main/run_logs/qwen2-7b_math_megatron.log), [wandb](https://wandb.ai/diqiuzhuanzhuan/verl_gpg_example_gsm8k_math/runs/yy8bheu8) | +| NVIDIA GPU | Qwen/Qwen2.5-VL-7B-Instruct | GRPO (Megatron) | 65.4 (GEO3k) | [script](https://github.com/verl-project/verl/blob/main/examples/grpo_trainer/run_qwen2_5_vl_7b_megatron.sh), [wandb](https://api.wandb.ai/links/megatron-core-moe-dev/1yngvkek) | +| AMD MI300 | deepseek-ai/deepseek-llm-7b-chat | PPO | 70.5 [1] | [log](https://github.com/yushengsu-thu/verl_training_log/blob/main/gsm8k/ppo_run_deepseek7b_llm.log) | +| AMD MI300 | deepseek-ai/deepseek-llm-7b-chat | GRPO | 71.4 [1] | [log](https://github.com/yushengsu-thu/verl_training_log/blob/main/gsm8k/grpo_run_deepseek7b_llm.log) | +| NVIDIA GPU | Qwen/Qwen2.5-14B-Instruct | GRPO-LoRA | 94.6 | [command and logs](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/Qwen2.5-14B-bsz64_8-prompt512-resp1024-lorarank32-score0.946.log) | +| NVIDIA GPU | Qwen/Qwen2.5-32B-Instruct | GRPO-LoRA | 95.8 | [command and logs](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/Qwen2.5-32B-bsz64_8-prompt512-resp1024-lorarank32-score0.958.log) | +| NVIDIA GPU | Qwen/Qwen2.5-72B-Instruct | GRPO-LoRA | 96.0 | [command and logs](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/Qwen2.5-72B-bs64_8-prompt512-resp1024-lorarank32-score0.960.log) | + +### DAPO math-17k + +- Training DAPO math-17k dataset: https://huggingface.co/datasets/BytedTsinghua-SIA/DAPO-Math-17k +- Testing: AIME'24: https://huggingface.co/datasets/BytedTsinghua-SIA/AIME-2024 + +Note: + +- For Qwen/Qwen2.5-Math-7B, we directly modify the max_position_embeddings to 32768 without observing performance degradation in order to train longer response length. + +| Hardware | Model | Method | Test score | Details | +| ---------- | -------------------------- | ----------------------- | ---------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| NVIDIA GPU | Qwen/Qwen2.5-Math-7B (32k) | DAPO | 36.3 | [command](https://github.com/verl-project/verl-recipe/blob/main//dapo/test_dapo_7b_math.sh), [logs](https://wandb.ai/verl-org/DAPO%20Reproduction%20on%20verl/runs/ow47vvon?nw=nwusertongyuxuan361) | +| NVIDIA GPU | Qwen/Qwen2.5-7B-Instruct | DAPO + Code Interpreter | 40.0 | [command](https://github.com/verl-project/verl-recipe/blob/main//retool/run_qwen2_7b_dapo.sh) | + +## Coding related datasets + +Below is the result on leetcode if not specified otherwise. + +| Hardware | Model | Method | Test score | Details | +| ---------- | ----------------------- | ------ | ---------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| NVIDIA GPU | PRIME-RL/Eurus-2-7B-SFT | RPIME | 36.1 | [script](https://github.com/verl-project/verl-recipe/blob/main//prime/run_prime_qwen_code.sh), [swanlab](https://swanlab.cn/@wangzefan/prime_example/runs/7f541qhspgmy8nmhdlx35/chart) | + +### Notes + +[1] During evaluation, we have only extracted answers following the format `"####"`. A more flexible answer extraction, longer response length, and better prompt engineering may lead to a higher score. + +[2] The default value of `actor_rollout_ref.actor.entropy_coeff` is set to `0.0` since verl 0.3.x on 2025-05-30, which is different from previous versions. diff --git a/verl/docs/algo/dapo.md b/verl/docs/algo/dapo.md new file mode 100644 index 0000000000000000000000000000000000000000..09ae88971d0e4f1a5a796c821823973eb5816bfb --- /dev/null +++ b/verl/docs/algo/dapo.md @@ -0,0 +1,187 @@ +# Recipe: Decoupled Clip and Dynamic Sampling Policy Optimization (DAPO) + +Last updated: 06/19/2025. + +> Open-Source Algorithm Implementation & Expriement Running: [Yuxuan Tong](https://tongyx361.github.io/), [Guangming Sheng](https://hk.linkedin.com/in/guangming-sheng-b50640211) + +🏠 [Homepage](https://dapo-sia.github.io/) | 📝 [Paper@arXiv](https://arxiv.org/abs/2503.14476) | 🤗 [Datasets&Models@HF](https://huggingface.co/collections/BytedTsinghua-SIA/dapo-67d7f1517ee33c8aed059da0) | 🐱 [Code@GitHub](https://github.com/verl-project/verl-recipe/tree/main/dapo/recipe/dapo) | 🐱 [Repo@GitHub](https://github.com/BytedTsinghua-SIA/DAPO) + +> We propose the **D**ecoupled Clip and Dynamic s**A**mpling **P**olicy **O**ptimization (DAPO) algorithm. By making our work publicly available, we provide the broader research community and society with practical access to scalable reinforcement learning, enabling all to benefit from these advancements. Our system is based on the awesome [verl](https://github.com/verl-project/verl) framework. Thanks for their great work! Applying DAPO training to Qwen2.5-32B base model proves to outperform the previous state-of-the-art DeepSeek-R1-Zero-Qwen-32B on AIME 2024, achieving **50%** accuracy with **50%** less training steps. +> +> ![dapo-main-result](https://dapo-sia.github.io/static/images/score.png) + +## Quickstart + +1. Prepare the datasets **on the Ray cluster**: + +```bash +bash prepare_dapo_data.sh # This downloads the datasets to ${HOME}/verl/data by default +``` + +2. Submit the job to the Ray cluster **from any machine**: + +```bash +cd verl # Repo root +export RAY_ADDRESS="http://${RAY_IP:-localhost}:8265" # The Ray cluster address to connect to +export WORKING_DIR="${PWD}" # The local directory to package to the Ray cluster +# Set the runtime environment like env vars and pip packages for the Ray cluster in yaml +export RUNTIME_ENV="./recipe/dapo/runtime_env.yaml" # This sets environment variables for the Ray cluster +bash recipe/dapo/run_dapo_qwen2.5_32b.sh # or other scripts +``` + +## Reproduction Runs + +| Setup | AIME 2024 Acc. | Hardware | Image | Commit | Environment Variables | Training Script | Training Record | +| -------------------------------------------- | -------------- | --------- | -------------------------------------------------------------------- | -------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------- | +| DAPO | 52% | 16x8xH800 | `hiyouga/verl:ngc-th2.6.0-cu126-vllm0.8.3-flashinfer0.2.2-cxx11abi0` | [`4f80e4`](https://github.com/verl-project/verl/tree/4f80e465c2ec79ab9c3c30ec74b9745de61d0490) | [runtime_env.yaml](https://github.com/verl-project/verl/blob/4f80e465c2ec79ab9c3c30ec74b9745de61d0490/recipe/dapo/runtime_env.yaml) | [run_dapo_qwen2.5_32b.sh](https://github.com/verl-project/verl/blob/4f80e465c2ec79ab9c3c30ec74b9745de61d0490/recipe/dapo/run_dapo_qwen2.5_32b.sh) | [W&B](https://wandb.ai/verl-org/DAPO%20Reproduction%20on%20verl/workspace?nw=wmb4qxfht0n) | +| DAPO w/o Dynamic Sampling | 50% | 16x8xH800 | `hiyouga/verl:ngc-th2.6.0-cu126-vllm0.8.3-flashinfer0.2.2-cxx11abi0` | [`4f80e4`](https://github.com/verl-project/verl/tree/4f80e465c2ec79ab9c3c30ec74b9745de61d0490) | [runtime_env.yaml](https://github.com/verl-project/verl/blob/4f80e465c2ec79ab9c3c30ec74b9745de61d0490/recipe/dapo/runtime_env.yaml) | [run_dapo_wo_ds_qwen2.5_32b.sh](https://github.com/verl-project/verl/blob/4f80e465c2ec79ab9c3c30ec74b9745de61d0490/recipe/dapo/run_dapo_wo_ds_qwen2.5_32b.sh) | [W&B](https://wandb.ai/verl-org/DAPO%20Reproduction%20on%20verl/workspace?nw=wmb4qxfht0n) | +| DAPO w/o Token-level Loss & Dynamic Sampling | 44% | 16x8xH20 | `hiyouga/verl:ngc-th2.5.1-cu120-vllm0.7.4-hotfix` | [`4f80e4`](https://github.com/verl-project/verl/tree/4f80e465c2ec79ab9c3c30ec74b9745de61d0490) | [runtime_env.yaml](https://github.com/verl-project/verl/blob/4f80e465c2ec79ab9c3c30ec74b9745de61d0490/recipe/dapo/runtime_env.yaml) | [run_dapo_early_qwen2.5_32b.sh](https://github.com/verl-project/verl/blob/4f80e465c2ec79ab9c3c30ec74b9745de61d0490/recipe/dapo/run_dapo_early_qwen2.5_32b.sh) | [W&B](https://wandb.ai/verl-org/DAPO%20Reproduction%20on%20verl/workspace?nw=wmb4qxfht0n) | + +> [!IMPORTANT] +> +> **📢 Call for Contribution!** +> +> Welcome to submit your reproduction runs and setups! + +## Configuration + +### Separated Clip Epsilons (-> Clip-Higher) + +An example configuration: + +```yaml +actor_rollout_ref: + actor: + clip_ratio_low: 0.2 + clip_ratio_high: 0.28 +``` + +`clip_ratio_low` and `clip_ratio_high` specify the $\varepsilon_{\text {low }}$ and $\varepsilon_{\text {high }}$ in the DAPO objective. + +Core relevant code: + +```python +pg_losses1 = -advantages * ratio +pg_losses2 = -advantages * torch.clamp(ratio, 1 - cliprange_low, 1 + cliprange_high) +pg_losses = torch.maximum(pg_losses1, pg_losses2) +``` + +### Dynamic Sampling (with Group Filtering) + +An example configuration: + +```yaml +data: + gen_batch_size: 1536 + train_batch_size: 512 +algorithm: + filter_groups: + enable: True + metric: acc # score / seq_reward / seq_final_reward / ... + max_num_gen_batches: 10 # Non-positive values mean no upper limit +``` + +Setting `filter_groups.enable` to `True` will filter out groups whose outputs' `metric` are all the same, e.g., for `acc`, groups whose outputs' accuracies are all 1 or 0. + +The trainer will repeat sampling with `gen_batch_size` until there are enough qualified groups for `train_batch_size` or reaching the upper limit specified by `max_num_gen_batches`. + +Core relevant code: + +```python +prompt_bsz = self.config.data.train_batch_size +if num_prompt_in_batch < prompt_bsz: + print(f'{num_prompt_in_batch=} < {prompt_bsz=}') + num_gen_batches += 1 + max_num_gen_batches = self.config.algorithm.filter_groups.max_num_gen_batches + if max_num_gen_batches <= 0 or num_gen_batches < max_num_gen_batches: + print(f'{num_gen_batches=} < {max_num_gen_batches=}. Keep generating...') + continue + else: + raise ValueError( + f'{num_gen_batches=} >= {max_num_gen_batches=}. Generated too many. Please check your data.' + ) +else: + # Align the batch + traj_bsz = self.config.data.train_batch_size * self.config.actor_rollout_ref.rollout.n + batch = batch[:traj_bsz] +``` + +### Flexible Loss Aggregation Mode (-> Token-level Loss) + +An example configuration: + +```yaml +actor_rollout_ref: + actor: + loss_agg_mode: "token-mean" # / "seq-mean-token-sum" / "seq-mean-token-mean" + # NOTE: "token-mean" is the default behavior +``` + +Setting `loss_agg_mode` to `token-mean` will mean the (policy gradient) loss across all the tokens in all the sequences in a mini-batch. + +Core relevant code: + +```python +if loss_agg_mode == "token-mean": + loss = verl_F.masked_mean(loss_mat, loss_mask) +elif loss_agg_mode == "seq-mean-token-sum": + seq_losses = torch.sum(loss_mat * loss_mask, dim=-1) # token-sum + loss = torch.mean(seq_losses) # seq-mean +elif loss_agg_mode == "seq-mean-token-mean": + seq_losses = torch.sum(loss_mat * loss_mask, dim=-1) / torch.sum(loss_mask, dim=-1) # token-mean + loss = torch.mean(seq_losses) # seq-mean +else: + raise ValueError(f"Invalid loss_agg_mode: {loss_agg_mode}") +``` + +### Overlong Reward Shaping + +An example configuration: + +```yaml +data: + max_response_length: 20480 # 16384 + 4096 +reward_model: + overlong_buffer: + enable: True + len: 4096 + penalty_factor: 1.0 +``` + +Setting `overlong_buffer.enable` to `True` will penalize the outputs whose lengths are overlong but still within the hard context limit. + +Specifically, the penalty increases linearly from `0` to `overlong_buffer.penalty_factor` when the length of the output exceeds the `max_response_length - overlong_buffer.len` by `0` to `overlong_buffer.len` tokens. + +Core relevant code: + +```python +if self.overlong_buffer_cfg.enable: + overlong_buffer_len = self.overlong_buffer_cfg.len + expected_len = self.max_resp_len - overlong_buffer_len + exceed_len = valid_response_length - expected_len + overlong_penalty_factor = self.overlong_buffer_cfg.penalty_factor + overlong_reward = min(-exceed_len / overlong_buffer_len * overlong_penalty_factor, 0) + reward += overlong_reward +``` + +## FAQ + +### Where is the "Overlong Filtering" in the paper? + +Most experiments in the paper, including the best-performant one, are run without Overlong Filtering because it's somehow overlapping with Overlong Reward Shaping in terms of properly learning from the longest outputs. So we don't implement it here. + +### What's the difference between [the `recipe/dapo` directory in the `main` branch](https://github.com/verl-project/verl-recipe/tree/main/dapo) and the [`recipe/dapo` branch](https://github.com/verl-project/verl-recipe/tree/main/dapo/recipe/dapo)? + +[The `recipe/dapo` branch](https://github.com/verl-project/verl-recipe/tree/main/dapo/recipe/dapo) is for **as-is reproduction** and thus won't be updated with new features. + +[The `recipe/dapo` directory in the `main` branch](https://github.com/verl-project/verl-recipe/tree/main/dapo) works as an example of how to extend the latest `verl` to implement an algorithm recipe, which will be maintained with new features. + +### Why can't I produce similar results after modifications? + +RL infrastructures nowadays still have inherent unrobustness, on which we are still working hard to improve. + +We strongly recommend to only modify one thing at a time. + +We also list some known problems here: + +1. Enabling CUDA graph (`enforce_eager=False`) might cause model performance degradation, whose cause is still under investigation. diff --git a/verl/docs/algo/dppo.md b/verl/docs/algo/dppo.md new file mode 100644 index 0000000000000000000000000000000000000000..d4cd798638c18e179802c6c2ce41698e0bf6028c --- /dev/null +++ b/verl/docs/algo/dppo.md @@ -0,0 +1,96 @@ +# Divergence Proximal Policy Optimization (DPPO) + +Last updated: 02/25/2026. + + +
+ +## Rethinking the Trust Region in LLM Reinforcement Learning + +[![Paper](https://img.shields.io/badge/paper-A42C25?style=for-the-badge&logo=arxiv&logoColor=white )](https://arxiv.org/pdf/2602.04879) +[![Github](https://img.shields.io/badge/Stable_RL-000000?style=for-the-badge&logo=github&logoColor=000&logoColor=white)](https://github.com/sail-sg/Stable-RL) +[![Twitter](https://img.shields.io/badge/Twitter-%23000000.svg?style=for-the-badge&logo=twitter&logoColor=white)](https://x.com/QPHutu/status/2019435642539897303) + +
+ + +## ✨Getting started + +1. Prepare the datasets by running [prepare_dapo_data.sh](https://github.com/verl-project/verl-recipe/blob/3490a22a0a3adeb7e4787fe70b1060b642efbae4/dapo/prepare_dapo_data.sh): + +```bash +bash prepare_dapo_data.sh # This downloads the datasets to ${HOME}/verl/data by default +``` + +2. Prepare the model: + +```bash +hf download Qwen/Qwen3-30B-A3B-Base --local-dir ${HOME}/verl/models/Qwen3-30B-A3B-Base +``` + +3. Run the script: + +```bash +# run DPPO-Binary-KL +LOSS_MODE=dppo_kl bash examples/dppo_trainer/run_qwen3_30b_a3b_megatron.sh + +# run DPPO-Binary-TV +LOSS_MODE=dppo_tv bash examples/dppo_trainer/run_qwen3_30b_a3b_megatron.sh + +# run GRPO baseline +LOSS_MODE=vanilla CLIP_LOW=0.2 CLIP_HIGH=0.2 bash examples/dppo_trainer/run_qwen3_30b_a3b_megatron.sh +# or GRPO with clip higher +LOSS_MODE=vanilla CLIP_LOW=0.2 CLIP_HIGH=0.28 bash examples/dppo_trainer/run_qwen3_30b_a3b_megatron.sh +``` + +## 📖Introduction + +
+ issue +
+ +Comparison of **PPO** and the proposed **DPPO** (the Binary-TV variant). **(Left)** The surrogate objective and corresponding masks for PPO and DPPO. PPO (and variants like GRPO) employs a heuristic mask based on the probability ratio. In contrast, DPPO utilizes a more principled mask based on a direct approximation of policy divergence (e.g., Total Variation), ensuring updates stay within a theoretically grounded trust region. **(Right)** Experimental results on the AIME24 using Qwen3-30B-A3B-Base. DPPO significantly outperforms GRPO baselines, achieving superior training stability and final performance even without rollout routing replay (R3). + +
+ issue +
+ +DPPO variants achieve stable training while controlling the training-inference mismatch at a low level. In contrast, methods without a trust region (PG-IS, CISPO) or with a misspecified one (MiniRL) suffer from growing mismatch and eventual collapse. + +
+ issue +
+ +The plots show numerical differences between a training and an inference engine for Qwen3-30B-A3B-Base with identical parameters. **(Left)** The probability ratio (used in PPO) is highly volatile for low-probability tokens. **(Right)** In contrast, the TV divergence is more stable. This highlights a key flaw of PPO's clipping mechanism: it **over-penalizes low-probability tokens**, which can slow down learning; and **under-penalizes high-probability tokens**, which can permit large, destabilizing updates. + + +
+ issue +
+ +The most frequently clipped tokens (by GRPO) are important to the reasoning task! +They are dominated by: +- numbers, like 1, 4 +- mathematical symbols, like +, -, = +- reasoning and structural Words: Wait, Thus, Next + +## Top-K divergence approximation + +We only implement the DPPO-Binary-TV/DPPO-Binary-KL here due to their simplicity. + +For the TopK divergence approximation, please refer to the [the original repo](https://github.com/sail-sg/Stable-RL) for a complete implementation. + +## Citation +If you find our works useful for your research, please consider citing: + +```bibtex +@article{qi2026dppo, + title={Rethinking the Trust Region in LLM Reinforcement Learning}, + author={Qi, Penghui and Zhou, Xiangxin and Liu, Zichen and Pang, Tianyu and Du, Chao and Lin, Min and Lee, Wee Sun}, + journal={arXiv preprint arXiv:2602.04879}, + year={2026} +} +``` + +## 🌻Acknowledgement +We implement our reinforcement learning algorithm extending from [verl](https://github.com/verl-project/verl). We utilize [vLLM](https://github.com/vllm-project/vllm) and [sglang](https://github.com/sgl-project/sglang) for inference. Our models are trained primarily on [Qwen3 family](https://huggingface.co/collections/Qwen/qwen3). Our training data is built from [DAPO-MATH](https://huggingface.co/datasets/BytedTsinghua-SIA/DAPO-Math-17k). Thanks for their great contributions! diff --git a/verl/docs/algo/entropy.md b/verl/docs/algo/entropy.md new file mode 100644 index 0000000000000000000000000000000000000000..091190f79e6521cd755e9f3fab54565851da8afe --- /dev/null +++ b/verl/docs/algo/entropy.md @@ -0,0 +1,115 @@ +# Recipe: Entropy Mechanism + +Last updated: 06/27/2025. + + +
+ + The Entropy Mechanism of Reinforcement Learning for Large Language Model Reasoning. + +[![Paper](https://img.shields.io/badge/paper-A42C25?style=for-the-badge&logo=arxiv&logoColor=white)](https://arxiv.org/pdf/2505.22617) [![Github](https://img.shields.io/badge/PRIME-000000?style=for-the-badge&logo=github&logoColor=000&logoColor=white)](https://github.com/PRIME-RL/Entropy-Mechanism-of-RL) [![alphaXiv](https://img.shields.io/badge/discussion-A42C25?style=for-the-badge&logo=arxiv&logoColor=white&color=blue +)](https://www.alphaxiv.org/abs/2505.22617) [![Twitter](https://img.shields.io/badge/Twitter-%23000000.svg?style=for-the-badge&logo=twitter&logoColor=white)](https://x.com/stingning/status/1928088554166505667) [![Twitter](https://img.shields.io/badge/Twitter-%23000000.svg?style=for-the-badge&logo=twitter&logoColor=white)](https://x.com/charlesfornlp/status/1928089451080585283) [![Twitter-ak](https://img.shields.io/badge/Twitter-%23000000.svg?style=for-the-badge&logo=twitter&logoColor=white)](https://x.com/_akhaliq/status/1928077929105268861) + + + + +
+ + +## 🎉News + +- **[2025/05/29]** 🎉 Ranked **#1** of the day on [Huggingface Daily Papers](https://huggingface.co/papers?date=2025-05-29). +- **[2025/05/29]** Released our Paper on arXiv. See [here](https://arxiv.org/pdf/2505.22617). We provide insights into the entropy mechanism of RL for LLMs and propose two simple yet effective strategies to alleviate the entropy collapse. + + + +## ✨Getting started + +After preparing the training data, for training Qwen2.5-7B on a single node, taking the KL-Cov approach as an example, you can simply run: + +``` +cd verl +conda activate your_env +bash recipe/dapo/7b_kl_cov.sh +``` + +While for training Qwen2.5-32B on multi nodes, you can run the following commands: + +``` +cd verl +conda activate your_env +bash recipe/dapo/32b_kl_cov.sh +``` + +## 📖Introduction + +
+ issue +
+ +This paper addresses the entropy collapse issue in scaling reinforcement learning (RL) for large language models (LLMs), where policy entropy drops sharply during training, leading to overconfidence and performance saturation. We empirically establish a relationship between entropy ($H$) and performance ($R$): $R=−aexp(H)+b$, showing performance is bottlenecked by entropy exhaustion. + +
+ issue +
+ +Theoretically, we find entropy changes are driven by the covariance between action probability and logit updates, which correlates with advantage in Policy Gradient methods. High-probability, high-advantage actions reduce entropy, while rare, high-advantage actions increase it. Empirically, the covariance term remains positive, explaining entropy’s monotonic decline. To mitigate this, we propose ​​Clip-Cov​​ and ​​KL-Cov​​, which restrict updates for high-covariance tokens. These methods effectively prevent entropy collapse, and improve performance. + +## 📃Evaluation + +
+ issue +
+ + +Our method is able to maintain a considerably higher level of entropy throughout training. For example, when the baseline's entropy reaches a plateau and can no longer be consumed, the KL-Cov method still sustains an entropy level over 10 times higher. Meanwhile, the response length of the policy model steadily increases, and its performance on the test set consistently surpasses that of the baseline. This indicates that our model is able to explore more freely during training, learning better policy through RL. +| **Method** | **AIME24** | **AIME25** | **AMC** | **MATH-500** | **OMNI-MATH** | **OlympiadBench** | **Minerva** | **Avg.** | +| ----------------- | ---------: | ---------: | -------: | -----------: | ------------: | ----------------: | ----------: | -------: | +| *Qwen2.5-7B* | | | | | | | | | +| GRPO | 21.2 | 9.6 | 58.7 | 78.8 | 27.9 | 40.7 | 36.7 | 38.6 | +| w. Clip-higher | 18.1 | 11.5 | 56.6 | 79.2 | 29.8 | 43.3 | 40.4 | 38.8 | +| w. **`CLIP-Cov`** | 22.1 | **15.8** | 58.2 | 80.4 | **30.5** | **44.1** | **41.1** | 40.4 | +| w. **`KL-Cov`** | **22.6** | 12.9 | **61.4** | **80.8** | 29.1 | 42.6 | 38.2 | **40.6** | +| *Qwen2.5-32B* | | | | | | | | | +| GRPO | 21.8 | 16.2 | 69.7 | 84.2 | 35.2 | 43.6 | 45.5 | 45.8 | +| w. Clip-higher | 35.6 | 22.3 | 69.5 | 77.2 | 35.1 | 42.5 | 43.0 | 47.2 | +| w. **`CLIP-Cov`** | 32.3 | 22.7 | 67.2 | **87.0** | **42.0** | **57.2** | 46.0 | 50.3 | +| w. **`KL-Cov`** | **36.8** | **30.8** | **74.5** | 84.6 | 39.1 | 49.0 | **46.3** | **52.2** | + +Our two approaches both achieve non-trivial improvements across all benchmarks. Compared to GRPO, our method outperforms it by 2.0% on average for the 7B model and by 6.4% for the 32B model. Moreover, we observe that our method yields more substantial gains on the larger Qwen2.5-32B. Specifically, our method achieves improvements of 15.0% and 14.6% compared to GRPO on the most challenging benchmarks, AIME24 and AIME25, respectively. + + +## 🎈Citation +If you find this paper or repo helpful, please cite us. + +```bibtex +@article{cui2025entropy, + title={The Entropy Mechanism of Reinforcement Learning for Reasoning Language Models}, + author={Cui, Ganqu and Zhang, Yuchen and Chen, Jiacheng and Yuan, Lifan and Wang, Zhi and Zuo, Yuxin and Li, Haozhan and Fan, Yuchen and Chen, Huayu and Chen, Weize and others}, + journal={arXiv preprint arXiv:2505.22617}, + year={2025} +} +``` +## 🌻Acknowledgement +We implement our reinforcement learning algorithm extending from [verl](https://github.com/verl-project/verl). We utilize [vLLM](https://github.com/vllm-project/vllm) for inference. Our models are trained primarily on [Qwen2.5 family](https://github.com/QwenLM/Qwen2.5). Our training data is built from [DAPO-MATH](https://huggingface.co/datasets/BytedTsinghua-SIA/DAPO-Math-17k). Thanks for their great contributions! + +## 📬 Contact + +For questions, discussion, or collaboration opportunities, feel free to contact: +- Ganqu Cui: cuiganqu@pjlab.org.cn +- Yuchen Zhang: yuchen.zhang2003@gmail.com +- Jiacheng Chen: jackchan9345@gmail.com +- Ning Ding: ningding.cs@gmail.com + diff --git a/verl/docs/algo/gpg.md b/verl/docs/algo/gpg.md new file mode 100644 index 0000000000000000000000000000000000000000..36bede8c319040ae713ef335372f2caa40ce44a3 --- /dev/null +++ b/verl/docs/algo/gpg.md @@ -0,0 +1,36 @@ +# GPG: Group Policy Gradient + +Last updated: 07/03/2025. + +Group Policy Gradient (GPG) is a minimalist reinforcement learning (RL) method that enhances the reasoning ability of large language models without relying on supervised fine-tuning or complex tricks. GPG revisits traditional policy gradients and directly optimizes the RL objective—no surrogate losses, no KL penalties, no critic, and no reference model. Compared to GRPO, GPG is simpler, more efficient, and achieves better results on many tasks. For more details, please refer to the original paper [GPG: A Simple and Strong Reinforcement Learning Baseline for Model Reasoning +](https://arxiv.org/abs/2504.02546). + +## Key Components +- Use a corrected advantage function to improve policy gradient accuracy and training efficiency. +- By eliminating the critic and reference models, avoiding KL divergence constraints, significantly simplifies the training process compared to Group Relative Policy Optimization (GRPO) + +## Configuration +To configure GPG within the framework, use the following YAML settings. + +```yaml +algorithm: + adv_estimator: gpg +actor_rollout_ref: + actor: + policy_loss: + loss_mode: "gpg" +``` + +## Advanced Extensions +GPG is a simple and strong baseline for model reasoning. Although it avoids using KL loss in its original form, you can still use KL loss to further improve the performance. + +```yaml +algorithm: + adv_estimator: gpg +actor_rollout_ref: + actor: + use_kl_loss: True # enable kl regularization + kl_loss_coef: 0.01 + policy_loss: + loss_mode: "gpg" +``` \ No newline at end of file diff --git a/verl/docs/algo/grpo.md b/verl/docs/algo/grpo.md new file mode 100644 index 0000000000000000000000000000000000000000..0a1fb1d324ae89ad97eb49dbc2e9331aaf2da314 --- /dev/null +++ b/verl/docs/algo/grpo.md @@ -0,0 +1,72 @@ +# Group Relative Policy Optimization (GRPO) + +Last updated: 05/31/2025. + +In reinforcement learning, classic algorithms like PPO rely on a "critic" model to estimate the value of actions, guiding the learning process. However, training this critic model can be resource-intensive. + +GRPO simplifies this process by eliminating the need for a separate critic model. Instead, it operates as follows: +- Group Sampling: For a given problem, the model generates multiple possible solutions, forming a "group" of outputs. +- Reward Assignment: Each solution is evaluated and assigned a reward based on its correctness or quality. +- Baseline Calculation: The average reward of the group serves as a baseline. +- Policy Update: The model updates its parameters by comparing each solution's reward to the group baseline, reinforcing better-than-average solutions and discouraging worse-than-average ones. + +This approach reduces computational overhead by avoiding the training of a separate value estimation model, making the learning process more efficient. For more details, refer to the original paper [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://arxiv.org/pdf/2402.03300) + +## Key Components + +- No Value Function (Critic-less): unlike PPO, GRPO does not train a separate value network (critic) +- Group Sampling (Grouped Rollouts): instead of evaluating one rollout per input, GRPO generates multiple completions (responses) from the current policy for each prompt. This set of completions is referred to as a group. +- Relative Rewards: within each group, completions are scored (e.g., based on correctness), and rewards are normalized relative to the group. + +## Configuration + +Note that all configs containing `micro_batch_size` are used to configure the maximum sample or token count per forward or backward pass to avoid GPU OOMs, whose value should not change algorithmic/convergence behavior. + +Despite that many configurations start with the `ppo_` prefix, they work across different RL algorithms in verl, as the GRPO training loop is similar to that of PPO (without critic). + +![image](https://github.com/user-attachments/assets/16aebad1-0da6-4eb3-806d-54a74e712c2d) + +- `actor_rollout.ref.rollout.n`: For each prompt, sample n times. Default to 1. For GRPO, please set it to a value larger than 1 for group sampling. + +- `data.train_batch_size`: The global batch size of prompts used to generate a set of sampled trajectories/rollouts. The number of responses/trajectories is `data.train_batch_size * actor_rollout.ref.rollout.n` + +- `actor_rollout_ref.actor.ppo_mini_batch_size`: The set of sampled trajectories is split into multiple mini-batches with batch_size=ppo_mini_batch_size for PPO actor updates. The ppo_mini_batch_size is a global size across all workers. + +- `actor_rollout_ref.actor.ppo_epochs`: Number of epochs for GRPO updates on one set of sampled trajectories for actor + +- `actor_rollout_ref.actor.clip_ratio`: The GRPO clip range. Default to 0.2 + +- `algorithm.adv_estimator`: Default is gae. Please set it to grpo instead + +- `actor_rollout_ref.actor.loss_agg_mode`: Default is "token-mean". Options include "token-mean", "seq-mean-token-sum", "seq-mean-token-mean". The original GRPO paper takes the sample-level loss (seq-mean-token-mean), which may be unstable in long-CoT scenarios. All GRPO example scripts provided in verl uses the default configuration "token-mean" for loss aggregation instead. + +Instead of adding KL penalty in the reward, GRPO regularizes by directly adding the KL divergence between the trained policy and the reference policy to the loss: + +- `actor_rollout_ref.actor.use_kl_loss`: To use kl loss in the actor. When used, we are not applying KL in the reward function. Default is False. Please set it to True for GRPO. + +- `actor_rollout_ref.actor.kl_loss_coef`: The coefficient of kl loss. Default is 0.001. + +- `actor_rollout_ref.actor.kl_loss_type`: Support kl(k1), abs, mse(k2), low_var_kl(k3) and full. Appending "+" in the end (e.g., 'k1+' and 'k3+') would apply straight through to employ k2 for unbiased gradient estimation, regardless of the kl value estimation (see https://github.com/verl-project/verl/pull/2953#issuecomment-3162113848 for more details). How to calculate the kl divergence between actor and reference policy. See this blog post for detailed analysis: http://joschu.net/blog/kl-approx.html + +## Advanced Extensions + +### DrGRPO + +[Understanding R1-Zero-Like Training: A Critical Perspective](https://arxiv.org/pdf/2503.20783) claims there's optimization bias in GRPO, which leads to artificially longer responses, especially for incorrect outputs. This inefficiency stems from the way GRPO calculates advantages using group-based reward normalization. Instead, DrGRPO aggregates token-level losses by normalizing with a global constant to eliminate length bias. + +Configure the following to enable DrGRPO, with all other parameters the same as GRPO's: + +- `actor_rollout_ref.actor.loss_agg_mode`: "seq-mean-token-sum-norm", which turns off seq-dim averaging +- `actor_rollout_ref.actor.loss_scale_factor`: (Optional) Set to a constant integer (e.g., max response length) to ensure consistent normalization throughout training. If not set, uses the current batch's response length. +- `actor_rollout_ref.actor.use_kl_loss`: Please set it to False for DrGRPO +- `algorithm.norm_adv_by_std_in_grpo`: False, which turns off standard deviation norm + +## Reference Example + +Qwen2.5 GRPO training log and commands: [link](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/qwen2-7b-fsdp2.log) + +```bash +bash examples/grpo_trainer/run_qwen3_8b_fsdp.sh +``` + +For more reference performance, please see https://verl.readthedocs.io/en/latest/algo/baseline.html diff --git a/verl/docs/algo/opd.md b/verl/docs/algo/opd.md new file mode 100644 index 0000000000000000000000000000000000000000..f852a0622ffc47892b92dbf45efa1bb036335d7b --- /dev/null +++ b/verl/docs/algo/opd.md @@ -0,0 +1,751 @@ +# On-Policy Distillation (OPD) + +**Author:** [Jacob Helwig](https://jacobhelwig.github.io/) + +Last updated: 05/14/2026. + +## Background + +### Summary + +1. OPD distills knowledge from teacher model(s) into a student model on states sampled from the student policy. +2. Compared with SFT or standard KD, OPD reduces exposure bias by aligning training-time states with inference-time states. +3. Compared with RLVR, OPD provides dense, continuous, token-level supervision rather than sparse outcome-level rewards. + +### Knowledge Distillation + +Knowledge distillation (KD) transfers behavior from a teacher model to a student model. In mathematical reasoning, for example, standard KD samples reasoning traces and final answers from the teacher, then trains the student with a next-token prediction objective against the teacher distribution. + +This can introduce exposure bias. During training, the student observes states sampled from the teacher. At inference time, however, states are sampled from the student. Unless the teacher and student induce the same state distribution, the student may not learn how the teacher would act in the states the student actually visits. + +For example, the student may prefer algebraic proofs while the teacher prefers geometric proofs. Standard KD primarily distills the teacher's behavior along geometric-proof trajectories, even if the student continues to generate algebraic-proof trajectories at inference time. + +### On-Policy RL + +RLVR has no train/inference state mismatch by construction: rollouts are sampled from the student policy, so the states the student trains on are exactly the states it would visit at inference time. If a rollout produces a correct final answer, the policy is updated to increase the likelihood of the sampled solution. + +This aligns training and inference states, but the reward is sparse and outcome-based. A rollout typically contributes a binary success signal at the sequence level rather than dense token-level feedback. + +### On-Policy Distillation + +On-policy distillation (OPD) [1,2,3] combines the state alignment of on-policy RL with the dense supervision of KD. The student samples rollouts from its own policy. Given each student-generated state, the teacher provides next-token log-probabilities, and the student is trained to match the teacher distribution at those states. + +Intuitively, the teacher provides guidance conditioned on the trajectory the student actually chose. If the student follows an algebraic proof path, the teacher supplies supervision for what it would do from that algebraic state. + +Formally, let $x \sim p_{\mathrm{data}}$ be a prompt, $y \sim \pi_{\theta}(\cdot \mid x)$ be a student rollout, and $s_t = (x, y_{.*` — one entry per teacher ([`DistillationTeacherModelConfig`](../../verl/workers/config/distillation.py)) +- `distillation.distillation_loss.*` — loss-mode and aggregation settings ([`DistillationLossConfig`](../../verl/workers/config/distillation.py)) + +Defaults below are the YAML defaults from +[`verl/trainer/config/distillation/distillation.yaml`](../../verl/trainer/config/distillation/distillation.yaml). + +--- + +### `distillation.enabled` (bool) + +Whether on-policy distillation is enabled. Default: `false`. + +When `true`, `main_ppo` allocates a separate teacher resource pool and spins up +one or more teacher inference servers; the actor loss switches from `ppo_loss` +to `distillation_ppo_loss`. + +### `distillation.n_gpus_per_node` (int) + +Number of GPUs per node in the teacher resource pool. Default: `8`. + +### `distillation.nnodes` (int) + +Number of nodes in the teacher resource pool. Default: `0` (effectively +disables the pool — must be set to `≥ 1` when `enabled=True`). + +**Constraint:** the total teacher pool size (`n_gpus_per_node × nnodes`) must +exactly equal the sum of `(num_replicas × per_replica_world_size)` across all +configured teachers, or `DistillationConfig.__post_init__` raises. + +### `distillation.teacher_key` (str) + +Field on each sample's data proto used to route the sample to the right +teacher in multi-teacher setups. Default: `"data_source"`. + +- **Single-teacher**: ignored (everything goes to the sole teacher). +- **Multi-teacher**: the value of `sample[teacher_key]` must match the `key` + of one of the configured teachers, or + `AsyncTeacherLLMServerManager._resolve_teacher_key` raises. + +### `distillation.teacher_models` (dict) + +Map of teacher entries. Each value is a `DistillationTeacherModelConfig`. + +The single-teacher entry is named `teacher_model` by convention. **Pitfall:** +when adding more named teachers, the `teacher_model` entry is silently popped +— so do **not** keep `teacher_model` as one entry alongside other named +teachers. Either rely on it alone, or rename it (e.g. `teacher_model1`) and +add the others. + +```bash +# WRONG: teacher_model is popped, only teacher_model2 is used +distillation.teacher_models.teacher_model.key=openai/gsm8k +distillation.teacher_models.teacher_model.model_path=Qwen/Qwen3-4B ++distillation.teacher_models.teacher_model2.key=hiyouga/geometry3k ++distillation.teacher_models.teacher_model2.model_path=Qwen/Qwen3-VL-4B-Instruct + +# RIGHT: rename the first teacher ++distillation.teacher_models.teacher_model1.key=openai/gsm8k ++distillation.teacher_models.teacher_model1.model_path=Qwen/Qwen3-4B ++distillation.teacher_models.teacher_model2.key=hiyouga/geometry3k ++distillation.teacher_models.teacher_model2.model_path=Qwen/Qwen3-VL-4B-Instruct +``` + +--- + +### `distillation.teacher_models..key` (str) + +Identifier used to route samples to this teacher in multi-teacher mode. Must +match the value of `sample[distillation.teacher_key]`. Default: `null` +(required for multi-teacher; auto-set to `"default"` for single-teacher). + +### `distillation.teacher_models..model_path` (str) + +Local path or Hugging Face model id for the teacher. **Required.** + +The teacher must share the student's tokenizer/vocab — typically satisfied by +picking a teacher in the same model family (e.g. `Qwen3-32B` teacher for a +`Qwen3-8B` student). + +### `distillation.teacher_models..num_replicas` (int) + +Number of inference replicas of this teacher to launch. Default: `0`. + +Each replica occupies +`per_replica_world_size = inference.tensor_model_parallel_size * inference.data_parallel_size * inference.pipeline_model_parallel_size` +GPUs, so the teacher's total footprint is `num_replicas × per_replica_world_size`. + +For a **single teacher**, you may leave this at `0`: `_resolve_teacher_models` +auto-fills it as `pool_size // per_replica_world_size`. + +### `distillation.teacher_models..inference.*` + +Inference-engine config for this teacher; see [`RolloutConfig`](../../verl/workers/config/rollout.py). Same shape as +`actor_rollout_ref.rollout.*`. Notable defaults inherited from the YAML: + +- `inference.name` — e.g. `vllm` or `sglang`. +- `inference.tensor_model_parallel_size` — default `2`. +- `inference.gpu_memory_utilization` — default `0.5`. +- `inference.max_model_len` — must accommodate `student_prompt_length + + student_response_length + 1`; otherwise + `validate_and_prepare_for_distillation` raises. +- `inference.engine_kwargs.vllm.max_logprobs` — auto-bumped to `≥ + distillation.distillation_loss.topk` whenever the active loss mode requires + top-k. + +`validate_and_prepare_for_distillation` rewrites +`inference.prompt_length := prompt_length + response_length` and +`inference.response_length := 1`, since the teacher only scores the +(prompt + response) prefix and emits one dummy token. + +--- + +### `distillation.distillation_loss.loss_mode` (str) + +Distillation divergence to use. Default: `"k3"`. + +Two registered families: + +- **Top-k** (`forward_kl_topk`): forward KL using the teacher's top-k logits. +- **Single-sample KL estimators** (`kl`, `k1`, `abs`, `mse`, `k2`, + `low_var_kl`, `k3`): per-token Monte Carlo estimators of reverse KL + computed from the student's `log_probs` and the teacher's single + `log_prob` at the sampled token. + +### `distillation.distillation_loss.topk` (int, optional) + +`k` for top-k distillation losses. Default: `32`. + +Only used when `loss_mode` requires top-$k$ (e.g. `forward_kl_topk`). Drives both +the teacher's `prompt_logprobs` request size and (for vLLM) the engine's +`max_logprobs` cap. + +### `distillation.distillation_loss.use_task_rewards` (bool) + +Whether to add the standard PPO/GRPO task-reward loss on top of the +distillation loss. Default: `true`. + +- `true`: final loss is `policy_loss + distillation_loss_coef × distill_loss`. +- `false`: the PPO term is zeroed and only the distillation loss contributes. + +Orthogonal to `use_policy_gradient` (which controls how the *distillation +signal itself* is applied). + +### `distillation.distillation_loss.distillation_loss_coef` (float) + +Coefficient on the distillation loss when combined with task rewards. +Default: `1.0`. Only takes effect when `use_task_rewards=true`. + +### `distillation.distillation_loss.loss_max_clamp` (float, optional) + +Per-token clamp on the distillation loss to `[-clamp, +clamp]`. Default: +`null` (no clamp). + +### `distillation.distillation_loss.log_prob_min_clamp` (float, optional) + +Lower clamp on log probabilities used inside divergence computations, to +prevent `log q − log p` from blowing up when `p` or `q` are near zero. +Default: `null`. + +### `distillation.distillation_loss.use_policy_gradient` (bool) + +How the distillation signal is applied. `true` corresponds to PG OPD, `false` to GKD OPD. Default: `false`. + + +**Validation:** + +- `use_policy_gradient=False` + `loss_mode="k1"` $\to$ `ValueError`. The k1 loss + has no gradient through the teacher logprob, so backpropagating it directly + is meaningless. +- `use_policy_gradient=True` + `loss_mode="forward_kl_topk"` $\to$ warning. The + PG update only moves $\nabla_\theta\log\pi_\theta(y_t|s_t)$ for the sampled token $y_t$, so the top-$k$ + distributional signal is largely unused. + +### `distillation.distillation_loss.policy_loss_mode` (str) + +Name of the policy loss to use when `use_policy_gradient=True`. Default: +`"vanilla"`. **Currently only `"vanilla"` is supported**; anything else raises +`NotImplementedError`. + +### `distillation.distillation_loss.clip_ratio` (float) + +PPO clip ratio used by the policy-gradient update when +`use_policy_gradient=True`. Default: `0.2`. + +### `distillation.distillation_loss.clip_ratio_low` (float) + +Lower bound of the PPO clip range. Default: `0.2`. + +### `distillation.distillation_loss.clip_ratio_high` (float) + +Upper bound of the PPO clip range. Default: `0.2`. + +### `distillation.distillation_loss.global_batch_info` / `loss_settings` + +Internal fields populated at runtime — **do not set from the user side.** +`loss_settings` is auto-populated from `loss_mode` via +`get_distillation_loss_settings`; `global_batch_info` is filled by the actor +worker before the loss runs. + +## Usage + +Example scripts are available in `examples/on_policy_distillation_trainer`. This section shows how to configure different OPD recipes. + +### Quick start + +For single-teacher OPD, first enable distillation, allocate a teacher resource pool, and specify the teacher model and inference server settings: + +```yaml +distillation: + enabled: true + + n_gpus_per_node: 2 + nnodes: 1 + + teacher_models: + teacher_model: + model_path: Qwen/Qwen3-32B + inference: + name: vllm + gpu_memory_utilization: 0.8 +``` + +The teacher must share the student's tokenizer and vocabulary. This is usually true for models from the same family, such as a `Qwen3-8B` student and a `Qwen3-32B` teacher. + +In most OPD runs, disable the standard PPO/GRPO reference-policy KL. Otherwise, the student is simultaneously regularized toward the reference policy and distilled from the teacher: + +```yaml +actor_rollout_ref: + actor: + use_kl_loss: false +algorithm: + use_kl_in_reward: false +``` + +### GKD OPD + +For efficiency, the current implementation of GKD OPD uses a top-$k$ approximation to forward KL using the top-$k$ teacher logits and the forward KL: + +$$ +\mathcal{L}_{\mathrm{GKD}}^{(k)}(s_t) += +\sum_{v \in \operatorname{TopK}(\nu(\cdot \mid s_t))} +\nu(v \mid s_t) +\bigl[ +\log \nu(v \mid s_t) +- +\log \pi_\theta(v \mid s_t) +\bigr]. +$$ + +The reason GKD OPD is implemented only over the teacher top-$k$ logits is because current inference servers return log-probabilities for the sampled token and the teacher top-$k$ tokens, but do not support gathering log-probabilities at arbitrary token IDs. Therefore, the implementation supports teacher-top-$k$ forward KL, but not student-top-$k$ reverse KL. + +To use GKD OPD, set `loss_mode=forward_kl_topk`, choose `topk`, and disable policy-gradient distillation: + +```yaml +distillation: + distillation_loss: + loss_mode: forward_kl_topk + topk: 128 + use_policy_gradient: false +``` + +Do not use `forward_kl_topk` with `use_policy_gradient=true`. The top-$k$ loss contains distributional information for many teacher-preferred tokens, but a policy-gradient update only acts through the sampled token: + +$$ +\nabla_\theta \mathcal{L}_{\mathrm{PG}} +\propto +- A_t \nabla_\theta \log \pi_\theta(y_t \mid s_t). +$$ + +Thus, the update cannot directly assign credit to the non-sampled top-$k$ tokens. This discards most of the distributional signal and can produce misleading updates. For example, if the student already matches the teacher on the sampled token but overestimates other teacher-top-$k$ tokens, the forward KL is still positive; using it as a policy-gradient reward would incorrectly push on the sampled token. + + +### PG OPD + +PG OPD treats the negative reverse-KL estimate as a reward and applies a policy-gradient update. To use PG OPD with the `k1` estimator, set `loss_mode=k1`, enable policy-gradient distillation, and configure the PPO clipping range: + +```yaml +distillation: + distillation_loss: + loss_mode: k1 + use_policy_gradient: true + policy_loss_mode: vanilla + clip_ratio_low: 0.2 + clip_ratio_high: 0.28 +``` + +Currently, only `policy_loss_mode=vanilla` is supported. Other policy-loss modes, such as `dppo_tv`, require additional parameters and are not implemented for OPD. + +### Task rewards + +OPD can be optimized alone or combined with the standard PPO/GRPO task-reward loss. + +When `use_task_rewards=true`, the final loss is + +$$ +\mathcal{L} += +\mathcal{L}_{\mathrm{policy}} ++ +\lambda_{\mathrm{distill}} +\mathcal{L}_{\mathrm{distill}}, +$$ + +where $\mathcal{L}_{\mathrm{policy}}$ is the PPO/GRPO task-reward loss, $\mathcal{L}_{\mathrm{distill}}$ is the distillation loss, and $\lambda_{\mathrm{distill}}$ is set by `distillation_loss_coef`. + +To combine task rewards with distillation: + +```yaml +distillation: + distillation_loss: + use_task_rewards: true + distillation_loss_coef: 1.5 +``` + +When `use_task_rewards=false`, the PPO/GRPO task-reward loss is zeroed and the update optimizes only the distillation loss. + +### Multi-teacher OPD + +Multiple teachers can be configured by adding one entry under `distillation.teacher_models` per teacher. Each teacher has a routing `key`, model path, replica count, and inference configuration. + +```yaml +distillation: + n_gpus_per_node: 8 + nnodes: 2 + teacher_key: data_source + + teacher_models: + gsm8k: + key: "openai/gsm8k" + model_path: Qwen/Qwen3-32B + num_replicas: 2 + inference: + name: vllm + tensor_model_parallel_size: 2 + gpu_memory_utilization: 0.6 + + geo3k: + key: "hiyouga/geometry3k" + model_path: Qwen/Qwen3-VL-32B-Instruct + num_replicas: 3 + inference: + name: vllm + tensor_model_parallel_size: 4 + gpu_memory_utilization: 0.8 + +data: + shuffle: true + reward_fn_key: data_source +``` + +In this example, the teacher pool has `8 * 2 = 16` GPUs. Assuming `data_parallel_size=1` and `pipeline_model_parallel_size=1`, the teacher footprints are: + +$$ +\text{gsm8k}: 2 \text{ replicas} \times 2 \text{ GPUs} = 4 \text{ GPUs} +$$ + +$$ +\text{geo3k}: 3 \text{ replicas} \times 4 \text{ GPUs} = 12 \text{ GPUs} +$$ + +so the total teacher footprint is $4 + 12 = 16$ GPUs, matching the resource pool. + +Teacher replicas are assigned by linearly splitting the teacher resource pool into contiguous GPU bundles. Each individual replica must occupy the expected number of nodes implied by its `per_replica_world_size`: + +```python +per_replica_world_size = tensor_model_parallel_size * data_parallel_size * pipeline_model_parallel_size +``` + +With `n_gpus_per_node=8`, the example above aligns cleanly: + +```text +node 0: [gsm8k replica 0: 2 GPUs] [gsm8k replica 1: 2 GPUs] [geo3k replica 0: 4 GPUs] +node 1: [geo3k replica 1: 4 GPUs] [geo3k replica 2: 4 GPUs] +``` + +No replica crosses a node boundary unless its `per_replica_world_size` requires multiple nodes. + +A similar-looking configuration can fail if a replica's GPU bundle does not fall entirely within a single node. For example, with the same `n_gpus_per_node=8`, `nnodes=2` (pool size 16) but two teachers configured as + +- `a`: `tensor_model_parallel_size=3`, `num_replicas=2` $\to$ 6 GPUs total +- `b`: `tensor_model_parallel_size=5`, `num_replicas=2` $\to$ 10 GPUs total + +the pool size still matches (`6 + 10 = 16`), but the linear bundle layout becomes + +```text +node 0: [a_0: 0,1,2] [a_1: 3,4,5] [b_0: 6,7,... +node 1: ...,8,9,10] [b_1: 11,12,13,14,15] +``` + +Replica `b_0` spans bundles `[6, 11)` — straddling node 0 (bundles 6, 7) and node 1 (bundles 8, 9, 10). A 5-GPU replica with `n_gpus_per_node=8` is expected to fit on a single node (`ceil(5 / 8) = 1`), so `_validate_replica_node_alignment` raises. To fix it, adjust `num_replicas` and the per-teacher inference parallelism. + +#### Teacher routing + +The `teacher_key` controls routing. It must name a top-level field on each sample's `non_tensor_batch`. `data_source` is one such field, set by the dataset loader. In the example above, `teacher_key=data_source`, so samples with `data_source="openai/gsm8k"` are routed to the `gsm8k` teacher, and samples with `data_source="hiyouga/geometry3k"` are routed to the `geo3k` teacher. + +When routing by data source, enable data shuffling. Without shuffling, a dataset created by concatenating other datasets may activate only one teacher for long contiguous stretches. For example, if GSM8K examples are followed by Geo3K examples, then training will use only the GSM8K teacher for the first portion of the epoch and only the Geo3K teacher for the remaining portion. + +## Metrics + +OPD logs metrics under `actor/distillation/*`. + +### Core metrics + +- `actor/distillation/loss` + Unscaled distillation loss. When `use_task_rewards=true`, compare this with `actor/pg_loss` to choose `distillation_loss_coef`. + +- `actor/distillation/abs_loss` + Absolute value of the distillation loss. This is mainly useful for signed estimators such as `k1`, where divergences can be negative. + +- `actor/distillation/loss_min` / `actor/distillation/loss_max` + Minimum and maximum per-token distillation loss in the batch. + +### Top-$k$ metrics + +These metrics are logged for top-$k$ loss modes such as `forward_kl_topk`. + +- `actor/distillation/student_mass` + Average student probability mass assigned to the teacher top-$k$ tokens. + +- `actor/distillation/teacher_mass` + Average teacher probability mass assigned to its own top-$k$ tokens. + +- `actor/distillation/student_mass_min` / `actor/distillation/student_mass_max` + Minimum and maximum student mass on the teacher top-$k$ tokens within the batch. + +- `actor/distillation/teacher_mass_min` / `actor/distillation/teacher_mass_max` + Minimum and maximum teacher mass on the teacher top-$k$ tokens within the batch. + +`teacher_mass` indicates how much of the teacher distribution is covered by the selected top-$k$. Low `teacher_mass` means the top-$k$ approximation is truncating substantial teacher probability mass. This can happen either by selecting too small $k$, or due to unstable optimization leading to the student generating low probability sequences. + +`student_mass` indicates how much probability the student assigns to the teacher-preferred tokens. + +## Debugging + +A useful technique for debugging modifications and additions to the distillation pipeline is to set the student to be the same model as the teacher. The loss should be approximately zero (not exact, due to differences between train/inference engines). + +## Architecture + +OPD has two components: + +1. **Teacher logprob computation** — runs on a dedicated teacher resource pool as part of the agent loop. +2. **Student optimization** — runs on the train workers, the same actor workers + that handle PPO/GRPO updates. + +### Teacher logprob computation + +Teacher logprob computation is interleaved with rollouts inside the **Agent +Loop**. Each sample's teacher call fires as soon as its rollout finishes (no batch-wide barrier) so teacher work overlaps with the still-running +rollouts on other samples. + +1. **Input.** `AgentLoopManager.generate_sequences(prompts: DataProto)` receives + a batch of prompts. + +2. **Chunking across workers.** The manager splits the batch evenly across its + `AgentLoopWorker` actors, then dispatches each + chunk via `worker.generate_sequences.remote(chunk)`. + +3. **Per-sample fan-out inside a worker.** Inside + `AgentLoopWorker.generate_sequences`, each sample in the chunk is launched as + its own asyncio task running `_run_agent_loop`. The agent loop runs on the + rollout GPUs and produces a rollout (prompt + response token ids). + +4. **Postprocess hook.** `_run_agent_loop` calls + `self._agent_loop_postprocess(...)`. This is where teacher logprob + computation is triggered, per sample, as soon as that sample's rollout is + ready. + +5. **Worker-side teacher dispatch.** `_agent_loop_postprocess` calls + `self._compute_teacher_logprobs(...)`, which extracts the routing value from + the sample's non-tensor fields using `sample_kwargs[self.teacher_key]` + (default `teacher_key="data_source"`), then calls + `self.teacher_server_manager.compute_teacher_logprobs_single(...)`. + +6. **Teacher selection.** + `AsyncTeacherLLMServerManager.compute_teacher_logprobs_single` resolves the + teacher via `_resolve_teacher_key`: + + - **Single-teacher**: routing key is ignored; the sole configured teacher + is used. + - **Multi-teacher**: `routing_key` must match a configured teacher in + `distillation.teacher_models`; otherwise an error is raised. + + The resolved key indexes into the per-teacher `LLMServerClient` dict to pick + the right client. + +7. **Sampling params for logprob computation.** The manager builds + sampling params with `max_tokens=1` plus `prompt_logprobs=topk` (or `0`), + so the teacher computes logprobs for the (prompt + response) sequence rather than + generating new tokens. `topk` is set to `distillation.distillation_loss.topk` + when the loss mode requires top-$k$ (e.g. `forward_kl_topk`); otherwise `0` + (single-sample logprob only). + +8. **Server-side load balancing.** The manager calls `client.generate(...)`, + which acquires a backing server through the shared `GlobalRequestLoadBalancer` + actor. + +9. **Backend execution.** With the vLLM backend the selected server is a + `vLLMHttpServer` actor; its `generate` method runs the forward pass and + returns a `TokenOutput` containing `prompt_ids` and `prompt_logprobs` for + the full (prompt + response) sequence. The SGLang backend has an analogous + server class. + +10. **Return path.** `compute_teacher_logprobs_single` packs the response into + two tensors `teacher_ids` and `teacher_logprobs`, each of shape `(S, 1 or K)` + where `S` is the sequence length and `K` is `topk` (or `1`). These are + stashed on the rollout output and later concatenated into the per-batch + `DataProto` for the student training step. + + +### Student training + +Using the `DataProto` produced by the Agent Loop (rollouts + teacher logprobs in +`teacher_logprobs`), the student step proceeds as follows. + +1. **Train entry.** `TrainingWorker.train_batch` invokes + `self.engine.train_batch(data, loss_function=self.loss_fn)`. When + distillation is enabled, `self.loss_fn` is bound to `distillation_ppo_loss` + at worker init; otherwise it is the standard `ppo_loss`. + +2. **Forward pass and (optional) inline top-k loss.** The training engine's + forward step runs the model forward and, for top-$k$ loss modes + (`distillation_use_topk=True`), invokes `distillation_ppo_loss` **as a + logits processor** while the full logits tensor is still in memory; this is + the `student_logits is not None` branch of `distillation_ppo_loss`. The + logits-processor branch dispatches to `compute_forward_kl_topk`, which has a + separate implementation per training engine (FSDP and Megatron). Per-token + `distillation_losses`, `student_mass`, and `teacher_mass` tensors are + written back into `model_output` so the full logits can be freed before the + final loss step. + +3. **Final loss.** After the forward, the engine calls the loss function with + `model_output` (full logits already freed); this is the + `student_logits is None` branch of `distillation_ppo_loss`, where: + + 1. **Per-token distillation loss** is produced by `distillation_loss(...)`, + which dispatches via `get_distillation_loss_fn(loss_mode)` to one of the registered distillation losses. + + 2. **Optional clamp.** If `loss_max_clamp` is set, per-token losses are + clamped to `[-clamp, +clamp]` (k1 in particular can be negative). + + 3. **Aggregation mode** — controlled by `use_policy_gradient`: + + - `False` (GKD OPD): straight backprop on `distillation_losses`. + - `True` (PG OPD): treat `-distillation_losses` as + advantages and run PPO-style clipped importance sampling. + + 4. **Combine with task rewards.** A standard PPO policy loss is computed + from the rollout's task rewards via `ppo_loss(...)`. If + `use_task_rewards=False` it is zeroed; otherwise the final loss is + `policy_loss + distillation_loss_coef * distill_loss`. + +The returned scalar loss is what `engine.train_batch` backpropagates. + +## Files + +### **Core Implementation** + +- `verl/experimental/teacher_loop/teacher_model.py` — `MultiTeacherModelManager` and `TeacherModelManager`; spin up teacher inference replicas on the dedicated teacher resource pool and expose per-teacher `LLMServerClient` factories +- `verl/experimental/teacher_loop/teacher_manager.py` — `AsyncTeacherLLMServerManager`; routes per-sample teacher calls (single- or multi-teacher) and builds teacher sampling params +- `verl/experimental/agent_loop/agent_loop.py` — `AgentLoopWorker._compute_teacher_logprobs`; per-sample teacher dispatch from `_agent_loop_postprocess`, packs `teacher_logprobs` into the rollout output +- `verl/trainer/distillation/fsdp/losses.py` — FSDP backend `compute_forward_kl_topk` +- `verl/trainer/distillation/megatron/losses.py` — Megatron backend `compute_forward_kl_topk` +- `verl/workers/engine_workers.py` — `ActorRolloutRefWorker.init_model`; binds `distillation_ppo_loss` as the actor's `loss_fn` when distillation is enabled +- `verl/workers/engine/{fsdp,megatron}/transformer_impl.py` — training-engine forward steps; invoke `distillation_ppo_loss` first as a logits processor (top-$k$ modes) and again as the final loss +- `verl/trainer/main_ppo.py` — `is_distillation_enabled` gate; allocates the dedicated `teacher_pool` resource pool +- `verl/trainer/ppo/ray_trainer.py` — constructs `MultiTeacherModelManager` and hands its `get_client()` dict to `AgentLoopWorker(... teacher_client=...)` +- `verl/workers/rollout/llm_server.py` — `LLMServerClient` and `GlobalRequestLoadBalancer` used for both student rollout and teacher logprob computation + +### **Configuration Files** + +- `verl/trainer/config/distillation/distillation.yaml` — YAML defaults for the `distillation.*` config tree +- `verl/workers/config/distillation.py` — dataclass schema (`DistillationConfig`, `DistillationLossConfig`, `DistillationTeacherModelConfig`) + +### **Documentation** + +- `docs/algo/opd.md` — this document + +### **Example Scripts** + +- `examples/on_policy_distillation_trainer/README.md` — script index +- `examples/on_policy_distillation_trainer/run_qwen3_8b_fsdp.sh` — text, vLLM rollout, FSDP student, single teacher +- `examples/on_policy_distillation_trainer/run_qwen3_8b_megatron.sh` — text, vLLM rollout, Megatron student, single teacher +- `examples/on_policy_distillation_trainer/run_qwen3_vl_8b_fsdp.sh` — VL student/teacher, vLLM rollout, FSDP student +- `examples/on_policy_distillation_trainer/run_qwen3_8b_mopd_fsdp.sh` — multi-teacher (gsm8k text + geo3k VL), routed by `data_source` + +### **Tests** + +- `tests/workers/test_distillation_topk_symmetry_on_cpu.py` — top-k loss symmetry checks +- `tests/utils/test_special_megatron_kl_loss_tp.py` — Megatron KL loss under tensor parallelism +- `tests/special_e2e/run_fully_async_policy_opd.sh` — end-to-end OPD with the fully-async rollouter \ No newline at end of file diff --git a/verl/docs/algo/opo.md b/verl/docs/algo/opo.md new file mode 100644 index 0000000000000000000000000000000000000000..338f3a762d9585c608af28cdf4e75837dbfe11e4 --- /dev/null +++ b/verl/docs/algo/opo.md @@ -0,0 +1,33 @@ +# On-Policy RL with Optimal Reward Baseline (OPO) + +Last updated: 06/02/2025. + +Loose on-policy constraints and suboptimal baselines in reinforcement learning often lead to training instability such as large policy shifts and entropy collapse. OPO addresses these challenges by using exact on-policy training with the theretically optimal reward baseline for advantage estimation. It achieves lower policy shifts and higher output entropy, encouraging more diverse and less repetitive responses. + +OPO uses group sampling to generate multiple outputs for each input like GRPO. Unlike group-based algorithms which typically use the mean reward of a group as its baseline, OPO employs a theoretically optimal baseline: the length-weighted reward of the group. It also omits the standard deviation normalization. By adopting these two key components, OPO enables the training of a single policy model with the objective of maximizing only the expected reward. For more detailes, refer to the original paper [On-Policy RL with Optimal Reward Baseline](https://arxiv.org/pdf/2505.23585). + +## Key Components + +- Exact On-Policy Training: always generates responses from the current policy, without using any pre-generated data or off-policy data. +- Optimal Reward Baseline: uses a length-weighted reward of the group as the baseline for normalizing the rewards. + +## Configuration + +To configure OPO within the framework, use the following YAML settings. These parameters are crucial for enabling exact on-policy training and activating the optimal reward baseline. + +```yaml +algorithm: + adv_estimator: opo # Use OPO for optimal reward baseline +data: + train_batch_size: 1024 +actor_rollout_ref: + actor: + ppo_mini_batch_size: 1024 # ppo_mini_batch_size should equal to train_batch_size to enable exact on-policy training + entropy_coeff: 0 # disable entropy regularization + use_kl_loss: False # disable kl regularization + kl_loss_coef: 0 +``` + +## Advanced Extensions + +OPO can also be extended to other algorithms like RLOO and Reinforce++. It just needs to adjust their configurations to enable exact on-policy training and incorporate the optimal length-weighted reward baseline with minimal modifications to their advantage estimation functions. diff --git a/verl/docs/algo/otb.md b/verl/docs/algo/otb.md new file mode 100644 index 0000000000000000000000000000000000000000..84b943986dd0ffe3eb0e040e3fcc32608aa2b80e --- /dev/null +++ b/verl/docs/algo/otb.md @@ -0,0 +1,99 @@ +# Optimal Token Baseline (OTB) + +Last updated: 02/23/2026. + +📝 [ArXiv](https://www.arxiv.org/abs/2602.07078) | 📒 [Blog](https://richardli.xyz/optimal-token-baseline) | 🤗 [Datasets](https://huggingface.co/datasets/Jiawei415/DPAO_filter) + +Optimal Token Baseline (OTB) is a dynamic token-level baseline for gradient variance reduction in policy-gradient reinforcement learning. It weights updates with the "Realized Energy" statistic that tracks how much uncertainty has accumulated up to each token, so noisy regions get downweighted while confident regions carry more weight. + +## Key properties + +- _Token-level baselines:_ OTB adapts per token by tracking realized energy, avoiding the padding artifacts that appear when group means dilute the signal with `EOS` tokens. +- _Forward-only overhead:_ The realized-energy statistic is computed via the **Logit-Gradient Proxy**, so OTB requires no extra backward passes or gradient-norm kernels. + +## Logit-Gradient Proxy + +Computing true uncertainty per token would normally mandate per-token backward passes. OTB sidesteps this by estimating realized energy entirely from forward probabilities, so it introduces negligible runtime overhead in practice. + +## Mechanics at a glance + +For each prompt group of size `N`, OTB computes rewards-to-go `G_t` and cumulative variance weights `W_t`. The optimal baseline per token is + +``` +B*_t = (Σ_i G_t^{(i)} · W_t^{(i)}) / (Σ_i W_t^{(i)} + ε), +W_t = Σ_{j=1}^t (1 - 2π_j + Σπ_j²), +Σπ_j² = exp(logsumexp(2·logits_j) - 2·logsumexp(logits_j)). +``` + +The final advantage is `(G_t - B*_t) · mask_t`, so padding tokens stay at zero. + +## Integration in VERL + +- `AdvantageEstimator.OPTIMAL_TOKEN_BASELINE` registers `compute_optimal_token_baseline_advantage`, invoked whenever `algorithm.adv_estimator` is set to `optimal_token_baseline`. +- `ActorRolloutRefWorker.compute_log_prob` emits an additional tensor `sum_pi_squared` (Σπ² per token) when `actor.calculate_sum_pi_squared=True`. This requires disabling fused log-prob kernels, because they do not surface logits. +- Trainers assert `sum_pi_squared` exists, regroup trajectories by `non_tensor_batch["uid"]`, and run the OTB calculation. If rollout IS is active, they rescale the weights by `rollout_is_weights**2` before aggregating. +- In Ulysses sequence-parallel setups, the actor gathers, unpads, and returns Σπ² in the same way it handles log-probabilities, so OTB supports sharded sequence-parallel models out of the box. +## Configuration checklist + +- `actor_rollout_ref.actor.calculate_sum_pi_squared: true` (mandatory). +- `actor_rollout_ref.model.use_fused_kernels: false` (required until fused kernels emit logits). +- `algorithm.adv_estimator: optimal_token_baseline` for single-turn RL and `tir_optimal_token_baseline` for multi-turn RL. +- Group sampling (`actor_rollout_ref.rollout.n > 1`) to unlock OTB’s variance reduction; with `n=1` the baseline collapses to returns. + +Example OmegaConf overlay: + +```yaml +algorithm: + adv_estimator: optimal_token_baseline + +actor_rollout_ref: + actor: + calculate_sum_pi_squared: true + rollout: + n: 8 +``` + +## Example script + +See `examples/otb_trainer/run_qwen3_8b_fsdp.sh` for a reference training loop. + +## Gradient Variance Proxy Metrics + +All gradient-variance analysis in the Optimal Token Baseline work starts from the variance identity + +``` +Var(ĝ) = E[||ĝ||²] - ||E[ĝ]||², +``` + +which states that the variance of any stochastic gradient equals the mean squared magnitude minus the squared norm of its expectation. + +For a trajectory `τ`, the policy-gradient estimator is + +``` +ĝ(τ) = ∇ log π_θ(τ) · A(τ), A(τ) = R(τ) - B. +``` + +The logit-gradient proxy approximates the squared gradient norm without an extra backward pass: + +``` +||ĝ(τ)||² ≈ Ŵ(τ) · A(τ)², +``` + +where `Ŵ(τ)` is the realized energy built. Given a mini-batch `{τ_i}` of size `N`, we decompose its statistics into three diagnostics: + +- **Signal strength (squared norm of the mean gradient)** + ``` + S = || (1/N) · Σ ĝ(τ_i) ||² + ``` +- **Total power (signal + noise)** + ``` + P_total = (1/N) · Σ Ŵ(τ_i) · A(τ_i)² + ``` +- **Pure noise (estimated variance of the batch mean)** + ``` + Var_proxy = (1/(N-1)) · (P_total - S) + ``` + +`verl/trainer/ppo/metric_utils.py#L306` implements these diagnostics via `compute_variance_proxy_metrics`, emitting `variance_proxy/proxy1_signal_strength`, `variance_proxy/proxy2_total_power`, and `variance_proxy/proxy3_pure_noise`. + +Tracking these metrics provides a forward-only, low-overhead view of gradient health for any advantage estimator that supplies `sum_pi_squared`. diff --git a/verl/docs/algo/ppo.md b/verl/docs/algo/ppo.md new file mode 100644 index 0000000000000000000000000000000000000000..e17620aee6303907a50d777053620b26fea6f1eb --- /dev/null +++ b/verl/docs/algo/ppo.md @@ -0,0 +1,105 @@ +# Proximal Policy Optimization (PPO) + +Last updated: 06/19/2025. + +Proximal Policy Optimization (PPO) is a family of policy gradient methods for reinforcement learning, proposed by OpenAI in 2017. PPO strikes a balance between simplicity, stability, and performance, making it one of the most widely used algorithms in modern RL applications, including large-scale language model fine-tuning. + +Traditional policy gradient methods like REINFORCE or Vanilla Policy Gradient suffer from: + +- High variance and sample inefficiency. +- Instability due to large policy updates. + +PPO addresses this problem using a clipped surrogate objective that avoids overly large updates without requiring second-order derivatives. + +For more technical details regarding PPO, we suggest reading the introduction in the [OpenAI spinning up tutorial](https://spinningup.openai.com/en/latest/algorithms/ppo.html), and the paper [Proximal Policy Optimization Algorithms](https://arxiv.org/abs/1707.06347). + +## Key Components + +- Actor-Critic Architecture: PPO requires both an actor model (policy) and a critic model (value function). This differs from other algorithms like GRPO and RLOO that don't require a critic model. + +- Generalized Advantage Estimation (GAE): PPO uses GAE for computing advantage values, which helps reduce variance in policy gradient estimates while maintaining low bias. + +- Clipped Surrogate Objective: The core of PPO is implemented through the clipped surrogate objective function that limits policy updates. + +## Configuration + +Note that all configs containing `micro_batch_size` are used to configure the maximum sample or token count per forward or backward pass to avoid GPU OOMs, whose value should not change algorithmic/convergence behavior. + +Most critic configs are similar to those of actors. Note that the critic model is omitted from the figure below. + +![image](https://github.com/user-attachments/assets/16aebad1-0da6-4eb3-806d-54a74e712c2d) + +- `data.train_batch_size`: The global batch size of prompts used to generate a set of sampled trajectories/rollouts. The number of responses/trajectories is `data.train_batch_size * actor_rollout.ref.rollout.n` + +- `actor_rollout_ref.actor.ppo_mini_batch_size`: The set of sampled trajectories is split into multiple mini-batches with batch_size=ppo_mini_batch_size for PPO actor updates. The ppo_mini_batch_size is a global size across all workers + +- `critic.ppo_mini_batch_size`: The set of sampled trajectories is split into multiple mini-batches with batch_size=ppo_mini_batch_size for PPO critic updates. The ppo_mini_batch_size is a global size across all workers + +- `actor_rollout_ref.actor.clip_ratio`: The PPO clip range. Default to 0.2 + +- `actor_rollout_ref.actor.ppo_epochs`: Number of epochs for PPO updates on one set of sampled trajectories for actor + +- `critic.ppo_epochs`: Number of epochs for PPO updates on one set of sampled trajectories for critic. Defaults to `actor_rollout_ref.actor.ppo_epochs` + +- `algorithm.gemma`: discount factor + +- `algorithm.lam`: The lambda term that trades off between bias and variance in the GAE estimator + +- `algorithm.adv_estimator`: Support gae, grpo, reinforce_plus_plus, reinforce_plus_plus_baseline, rloo + +## Advanced Extensions + +### KL Divergence Control + +Options to prevent the policy from diverging too far from a reference policy. Two mechanisms are available: KL reward penalty and KL loss. For more technical details, see [Training language models to follow instructions with human feedback](https://arxiv.org/abs/2203.02155) + +Options to use KL loss for KL divergence control: + +- `actor_rollout_ref.actor.use_kl_loss`: to use kl loss in the actor. When used, we are not applying KL in the reward function. Default is False + +- `actor_rollout_ref.actor.kl_loss_coef`: The coefficient of kl loss. Default is 0.001. + +- `actor_rollout_ref.actor.kl_loss_type`: Support kl(k1), abs, mse(k2), low_var_kl(k3) and full. Appending "+" in the end (e.g., 'k1+' and 'k3+') would apply straight through to employ k2 for unbiased gradient estimation, regardless of the kl value estimation (see https://github.com/verl-project/verl/pull/2953#issuecomment-3162113848 for more details). How to calculate the kl divergence between actor and reference policy. See this blog post for detailed analysis: http://joschu.net/blog/kl-approx.html + +Options to use KL penalty in the reward: + +- `algorithm.use_kl_in_reward`: Whether to enable in-reward kl penalty. Default is False. + +- `algorithm.kl_penalty`: Support kl(k1), abs, mse(k2), low_var_kl(k3) and full. This defines the way to calculate the kl divergence between actor and reference policy. For specific options, refer to `kl_penalty` in core_algos.py. See this blog post for detailed analysis: http://joschu.net/blog/kl-approx.html + +- `algorithm.kl_ctrl.kl_coef`: The (initial) coefficient of in-reward kl_penalty. Default is 0.001. +- `algorithm.kl_ctrl.type`: 'fixed' for FixedKLController and 'adaptive' for AdaptiveKLController. +- `algorithm.kl_ctrl.horizon`: See source code of AdaptiveKLController for details. +- `algorithm.kl_ctrl.target_kl`: See source code of AdaptiveKLController for details. + +### Dual-clip PPO + +The Dual-Clip PPO introduces a approach by applying a lower bound to the policy ratio when the advantage is less than zero, when multiplied by a large raito, does not exceed a specified lower bound. + +![image](https://github.com/user-attachments/assets/fc232181-d8b0-4307-8dd2-4dc0a4c1c139) + +- `actor_rollout_ref.actor.clip_ratio_c`: lower bound of the value for Dual-clip PPO, defaults to 3.0 + +## Reference Example + +Qwen2.5 training log and commands: [link](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/Qwen2.5-0.5B-bsz256_2-prompt1024-resp512-0.567.log) + +```bash +bash run_gemma.sh + trainer.n_gpus_per_node=1 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=1 \ + trainer.logger=console \ + critic.model.path=Qwen/Qwen2.5-0.5B-Instruct \ + actor_rollout_ref.model.path=Qwen/Qwen2.5-0.5B-Instruct \ + data.train_batch_size=256 \ + actor_rollout_ref.actor.ppo_mini_batch_size=64 \ + actor_rollout_ref.actor.ppo_micro_batch_size=2 \ + critic.ppo_micro_batch_size=2 +``` + +Reference performance with verl v0.2: + +| Model | Method | Score | Link | +|-------------------------------|------------------|-------|------------------------------------------------------------------------------------------------| +| Qwen/Qwen2.5-0.5B-Instruct | pretrained model | 36.4 | [Qwen Blog](https://qwenlm.github.io/blog/qwen2.5-llm/) | +| Qwen/Qwen2.5-0.5B-Instruct | PPO | 56.7 | [PPO Command and Logs](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/Qwen2.5-0.5B-bsz256_2-prompt1024-resp512-0.567.log) | diff --git a/verl/docs/algo/rollout_corr.md b/verl/docs/algo/rollout_corr.md new file mode 100644 index 0000000000000000000000000000000000000000..5bb91657db862c667e75dab0cbdbc87bcdafa50a --- /dev/null +++ b/verl/docs/algo/rollout_corr.md @@ -0,0 +1,1327 @@ +# Rollout Correction + +**Author:** [Yingru Li](https://richardli.xyz/) + +Last updated: 10/30/2025. + +--- + +> **📖 Documentation Structure** +> +> - **This document** - Practical usage guide: configurations, presets, troubleshooting +> - **[Mathematical Formulations](rollout_corr_math.md)** - Theoretical foundations, derivations, and algorithmic details +> +> Start here for implementation, refer to the math doc for theory and design rationale. + +--- + +This document provides a comprehensive overview of the Rollout Correction implementation in verl. + +**Note on Naming**: This feature is called "Rollout Correction" to reflect the complete functionality: importance sampling (IS) weights and rejection sampling (RS). The internal variable `rollout_is_weights` retains its name as it specifically refers to the IS weights component. + +### BibTeX Citation + +```bibtex +@online{liu-li-2025-rl-collapse, + title = {When Speed Kills Stability: Demystifying {RL} Collapse from the Training-Inference Mismatch}, + author = {Liu, Jiacai and Li, Yingru and Fu, Yuqian and Wang, Jiawei and Liu, Qian and Shen, Yu}, + year = {2025}, + month = sep, + url = {https://richardli.xyz/rl-collapse} +} + +@article{li2025trust, + title={Trust Region Masking for Long-Horizon LLM Reinforcement Learning}, + author={Li, Yingru and Liu, Jiacai and Xu, Jiawei and Tong, Yuxuan and Li, Ziniu and Liu, Qian and Wang, Baoxiang}, + journal={arXiv preprint arXiv:2512.23075}, + year={2025} +} +``` + +### Blog Series + +- Main blog post: https://richardli.xyz/rl-collapse +- [Part 1: Why Mismatch Breaks LLM-RL](https://richardli.xyz/rl-collapse-1) (analytical framework using TV distance for bias and χ²-divergence for variance) +- [Part 2: The Gradient Estimator Trials](https://richardli.xyz/rl-collapse-2) (token-level vs sequence-level correction bias-variance tradeoff) +- [Part 3: When Math Meets Reality—Toxic Tails and Length Traps](https://richardli.xyz/rl-collapse-3) (why rejection over clipping, and geometric-level RS) +- Latest Paper: https://arxiv.org/abs/2512.23075 + +## Overview + +Rollout Correction provides a unified framework to handle **general off-policy problems** in RL training. Any scenario where the data collection distribution differs from the training distribution can benefit from these methods. + +**Common off-policy scenarios:** + +1. **Policy Mismatch** (Implementation Differences) + + - Different precision: FP8 vs FP16 vs BF16 vs FP32 + - Different backends: vLLM vs SGLang vs FSDP vs Megatron + - Different implementations even with identical weights + +2. **Temporal Lag** (Model Staleness) + + - Rollout uses older checkpoint while training has progressed + - Asynchronous rollout workers with stale parameters + - Common in distributed/async RL systems + +3. **Replay Buffers** + + - Training on historical trajectories from earlier iterations + - Experience replay from different policy versions + - Data augmentation or resampling strategies + +4. **Off-Policy Algorithms** + + - Behavioral cloning from expert demonstrations + - DAPO (data from auxiliary policies) + - Any algorithm using trajectories from a different policy + +5. **Data Quality Filtering** + - Reweighting or filtering collected data + - Preference learning with modified distributions + - Curriculum learning with distribution shifts + +These off-policy gaps can cause training instability and policy collapse. Rollout Correction uses importance sampling (IS) weights and rejection sampling (RS) to correct for any distribution shift between data collection and training. + +**Important Note on Common Implementation Mistakes:** + +Many LLM-RL implementations incorrectly apply PPO by **ignoring the actual rollout policy** π_rollout and assuming the training reference policy π_old is the behavior policy. This is mathematically incorrect when π_rollout ≠ π_old (which is typical in LLM-RL due to precision/backend differences between rollout and training). + +**This is not PPO's fault** - PPO itself is mathematically correct. The issue is the incorrect assumption that π_old = π_rollout in naive implementations. + +This critical implementation mistake that leads to RL training collapse was identified in the blog post ["When Speed Kills Stability: Demystifying RL Collapse from the Training-Inference Mismatch"](https://richardli.xyz/rl-collapse) and motivated the development of this rollout correction framework. + +**Mathematically correct approaches:** + +- **Decoupled mode**: Three policies (π_rollout, π_old, π_θ) with IS correction from π_rollout to π_old +- **Bypass mode**: Two policies (π_rollout = π_old, π_θ) using actual rollout policy as PPO anchor +- **Bypass + Policy Gradient mode**: Two policies (π_rollout, π_θ) with IS/RS correction and no PPO clipping + +See [Mathematical Formulations](rollout_corr_math.md#37-common-implementation-mistake) for detailed explanation. + +### Key Design Principle: Separation of IS Weights and Rejection Sampling + +The implementation cleanly separates two orthogonal mechanisms: + +1. **IS Weights** (`rollout_is_weights`): Continuous reweighting for gradient correction + + - Policy ratio: π_old/π_rollout (decoupled) or π_θ/π_rollout (bypass) + - **Safety-bounded**: Clamped to [exp(-20), exp(20)] ≈ [2e-9, 5e8] to prevent overflow + - Token level: Bounds per-token ratios + - Sequence level: Bounds product of ratios (broadcast to all tokens) + - **Truncated**: Upper clamped via `.clamp(max=rollout_is_threshold)` (TIS: Truncated Importance Sampling) + - **Zeroed at padding**: Multiplied by response_mask to zero out padding positions + - Used to weight policy gradients (variance reduction) + +2. **Rejection Sampling** (`modified_response_mask`): Binary filtering for outlier exclusion + - Creates binary mask: 1 = keep, 0 = reject + - Rejects tokens/sequences with IS ratios outside [lower_threshold, upper_threshold] + - Modifies response_mask to exclude rejected samples from training + +This separation ensures: + +- ✅ IS weights provide continuous reweighting (reduce variance) +- ✅ Rejection sampling provides hard filtering (remove extreme outliers) +- ✅ Both mechanisms can be enabled independently or together +- ✅ Safety bounds prevent numerical overflow in all cases + +## Quick Start: Using Verified Presets + +**NEW**: We now provide typed configuration with verified presets for common scenarios. These presets have been validated with tens of thousands of GPU hours across various models and training scenarios. + +### Python API + +```python +from verl.trainer.config.algorithm import RolloutCorrectionConfig + +# === Decoupled PPO mode (3 policies: π_rollout, π_old, π_θ) === +# IS weights correct for gap between π_old and π_rollout +config = RolloutCorrectionConfig.decoupled_token_is() # Token-TIS +config = RolloutCorrectionConfig.decoupled_seq_is() # Seq-TIS +config = RolloutCorrectionConfig.decoupled_seq_is_rs() # Seq-MIS +config = RolloutCorrectionConfig.decoupled_geo_rs() # Geo-RS (ratio mode) +config = RolloutCorrectionConfig.decoupled_geo_rs_token_tis() # Geo-RS + Token-TIS + +# === K3 KL Estimator presets (more stable for small KL) === +config = RolloutCorrectionConfig.decoupled_k3_rs() # K3-RS only +config = RolloutCorrectionConfig.decoupled_k3_rs_token_tis() # K3-RS + Token-TIS + +# === Bypass PPO mode (2 policies: π_rollout = π_old, π_θ) - fast === +# PPO ratio handles IS, so no explicit IS weights needed +config = RolloutCorrectionConfig.bypass_ppo_clip() # PPO-clip only +config = RolloutCorrectionConfig.bypass_ppo_clip_geo_rs() # PPO-clip + Geo-RS +config = RolloutCorrectionConfig.bypass_ppo_clip_k3_rs() # PPO-clip + K3-RS + +# === Bypass PG mode (2 policies, no PPO clipping) - fast === +# IS weights computed on-the-fly as π_θ / π_rollout +config = RolloutCorrectionConfig.bypass_pg_is() # Seq-TIS + PG +config = RolloutCorrectionConfig.bypass_pg_geo_rs() # Geo-RS + PG +config = RolloutCorrectionConfig.bypass_pg_geo_rs_token_tis() # Geo-RS + Token-TIS + PG + +# === Other === +config = RolloutCorrectionConfig.disabled() # Metrics only (no correction) +``` + +### YAML Configuration (Advanced) + +For advanced customization or YAML-based configs: + +```yaml +algorithm: + rollout_correction: + rollout_is: token # IS weights: "token", "sequence", or null + rollout_is_threshold: 2.0 # TIS upper bound, or "0.5_5.0" for IcePop + rollout_is_batch_normalize: false # Batch normalize IS weights to mean=1.0 + rollout_rs: null # Rejection sampling: comma-separated canonical options (e.g. "token_k1,seq_max_k2") + rollout_rs_threshold: null # Threshold spec: float(s) or "lower_upper" string(s) + bypass_mode: false # Skip old_log_prob computation (sets π_old = π_rollout) + loss_type: ppo_clip # Loss type in bypass mode: "ppo_clip" (default) or "reinforce" + +# REQUIRED: Enable log prob calculation +actor_rollout_ref: + rollout: + calculate_log_probs: true +``` + +## Files + +### **Core Implementation** + +- `verl/trainer/ppo/rollout_corr_helper.py` - Contains `compute_rollout_correction_and_rejection_mask()` and `compute_offpolicy_metrics()` +- `verl/trainer/ppo/core_algos.py` - Rollout Correction integration with PPO and REINFORCE modes (`compute_policy_loss_bypass_mode()`, `compute_policy_loss_reinforce()`) +- `verl/trainer/ppo/ray_trainer.py` - Bypass mode implementation (skips `old_log_prob` computation) +- `verl/workers/utils/losses.py` - `ppo_loss` loss function wired to actor `TrainingWorker` via `verl.workers.engine_workers.ActorRolloutRefWorker.init_model` +- `verl/trainer/ppo/core_algos.py` - `@register_policy_loss("bypass_mode")` policy loss that invokes `compute_rollout_correction_and_rejection_mask` and emits off-policy metrics + +### **Configuration Files** + +- `verl/trainer/config/algorithm.py` - Rollout Correction parameters in `RolloutCorrectionConfig` +- `verl/workers/config/actor.py` - Rollout Correction parameters in `PolicyLossConfig` +- `verl/trainer/config/actor/actor.yaml` - Rollout Correction configuration section +- `verl/trainer/config/ppo_trainer.yaml` - Algorithm config with Rollout Correction + +### **Documentation** + +- `docs/examples/config.rst` - Configuration parameter descriptions + +### **Example Scripts** + +- `recipe/dapo/run_dapo_qwen2.5_32b_rollout_corr.sh` - DAPO example with Rollout Correction +- `examples/rollout_correction/run_qwen2_5_7b_fsdp.sh` - Basic example +- `examples/rollout_correction/run_qwen2_5_7b_fsdp_multi_rs.sh` - Multi-RS example + +### **Tests** + +- `tests/trainer/ppo/test_rollout_corr.py` - Unit tests for IS/RS mechanisms +- `tests/trainer/ppo/test_rollout_corr_integration.py` - Integration tests + +## Configuration Parameters + +All parameters are under `algorithm.rollout_correction`: + +### `rollout_is` (str or null) + +Importance sampling weights aggregation level: + +- `null` = No IS weights computed (metrics-only mode) +- `"token"`: Per-token IS weights + - **Decoupled mode**: ρ_t = π_old(t)/π_rollout(t) + - **Bypass/Pure IS mode**: ρ_t = π_θ(t)/π_rollout(t) + - Independent truncation per token + - Typical threshold: 1.5 - 5.0 +- `"sequence"`: Per-sequence weight ρ_seq = ∏_t ρ_t + - Multiplicative aggregation across sequence + - Typical threshold: 2.0 - 10.0 + +All IS weights are safety-bounded to [exp(-20), exp(20)] ≈ [2e-9, 5e8] + +### `rollout_is_threshold` (str or float) + +Threshold specification for IS weighting. Default: `2.0` + +- Single float or float-like string: TIS via `.clamp(max=rollout_is_threshold)` +- `"lower_upper"` string such as `"0.5_5.0"`: IcePop, zero weights outside `[lower, upper]` +- Applied to IS weights for variance reduction +- Separate from rejection sampling (controlled by `rollout_rs` parameters) +- Unlike `rollout_rs`, IcePop does not modify `response_mask`; it only changes the IS coefficients + +### `rollout_is_batch_normalize` (bool) + +Apply batch normalization to IS weights. Default: `False` + +- `True`: Normalize IS weights to have mean=1.0 within each batch + - **Token-level IS**: Normalizes over all token weights + - **Sequence-level IS**: Normalizes over sequence means (one weight per sequence) +- `False`: Use raw (truncated) IS weights +- Reduces variance by ensuring average weight is 1.0 per batch +- Applied AFTER truncation to preserve truncation semantics +- Only affects IS weight values, not rejection sampling + +### `rollout_rs` (str or null) + +Rejection sampling aggregation modes. Supply a comma-separated string (spaces optional) using the canonical options implemented in `rollout_corr_helper`: + +- `token_k1`: Token-level rejection with `-log r` bounds (ratio thresholds supplied as `lower_upper`). Example: `"0.6_1.4"` +- `token_k2`: Token-level rejection with `0.5 * (log r)^2` (upper bound only) +- `token_k3`: Token-level rejection with `exp(log r) - 1 - log r` (upper bound only) +- `seq_sum_k1`: Sequence-level rejection with sum of `-log r` (ratio bounds) +- `seq_sum_k2`: Sequence-level rejection with sum of `0.5 * (log r)^2` (upper bound only) +- `seq_sum_k3`: Sequence-level rejection with sum of `exp(log r) - 1 - log r` (upper bound only) +- `seq_mean_k1`: Sequence-level rejection with mean of `-log r` (ratio bounds) +- `seq_mean_k2`: Sequence-level rejection with mean of `0.5 * (log r)^2` (upper bound only) +- `seq_mean_k3`: Sequence-level rejection with mean of `exp(log r) - 1 - log r` (upper bound only) +- `seq_max_k2`: Sequence-level rejection with max of `0.5 * (log r)^2` (upper bound only) +- `seq_max_k3`: Sequence-level rejection with max of `exp(log r) - 1 - log r` (upper bound only) + +### `rollout_rs_threshold` (str, float, or null) + +Threshold specification for rejection sampling. + +- Provide **one entry per option**, separated by commas. A single entry is broadcast to every option. +- **K1 KL modes (`*k1`)**: Use `"lower_upper"` strings (e.g. `"0.7_1.3"`). Supplying a float implies only the upper bound; the lower bound defaults to its reciprocal. +- **K2/K3 KL modes (`*k2`/`*k3`)**: Supply positive upper bounds (float or numeric string). +- Set to `null` to disable thresholds entirely (only valid when `rollout_rs` is null). + +## Understanding the Framework: Components and Combinations + +The rollout correction framework is built from **orthogonal components** that can be combined flexibly. Understanding these components helps you choose the right configuration for your scenario. + +### Key Components + +1. **Operating Mode** (Section: [Operation Modes](#operation-modes)) + + - **Decoupled**: Three policies (π_rollout, π_old, π_θ) with separate π_old computation + - **Bypass**: Two policies (π_rollout = π_old, π_θ), skips π_old computation + +2. **Loss Function** (in bypass mode, controlled by `loss_type`) + + - **PPO-clip** (`loss_type="ppo_clip"`, default): PPO clipped objective (IS handled by ratio) + - **REINFORCE** (`loss_type="reinforce"`): Policy gradient with explicit IS weights (no clipping) + +3. **IS/RS Aggregation Level** + - **Token**: Per-token IS weights/rejection + - **Sequence**: Sequence-level IS weights/rejection + +See [Mathematical Formulations](rollout_corr_math.md#3-algorithmic-components-and-combinations) for detailed theory. + +--- + +## Preset Configuration Guide + +This section provides detailed guidance on choosing and using the verified presets. Each preset is a specific combination of components optimized for common scenarios. + +### Understanding the Presets + +#### Available Preset Methods + +| Preset Method | Estimator | Mode | IS Level | RS Level | Properties | +| ------------------------------------------------------------------------------ | ---------------- | ------------------ | -------- | -------- | --------------------------------------- | +| **Decoupled PPO Mode** (3 policies: π_rollout, π_old, π_θ) | +| `decoupled_token_is()` | Token-TIS | Decoupled | token | - | Token-level IS weights | +| `decoupled_seq_is()` | Seq-TIS | Decoupled | sequence | - | Sequence-level IS weights | +| `decoupled_seq_is_rs()` | Seq-MIS | Decoupled | sequence | sequence | Sequence IS + seq_sum_k1 RS | +| `decoupled_geo_rs()` | Geo-RS | Decoupled | - | sequence | Geometric RS (seq_mean_k1) | +| `decoupled_geo_rs_token_tis()` | Geo-RS-Token-TIS | Decoupled | token | sequence | Geometric RS + token IS | +| **K3 KL Estimator** (more stable for small KL values) | +| `decoupled_k3_rs()` | K3-RS | Decoupled | - | sequence | seq_mean_k3 RS | +| `decoupled_k3_rs_token_tis()` | K3-RS-Token-TIS | Decoupled | token | sequence | seq_mean_k3 RS + token IS | +| **Bypass Mode (PPO-clip)** (2 policies; ratio handles IS, RS masks outliers) | +| `bypass_ppo_clip()` | - | Bypass (PPO-clip) | - | - | PPO-clip only | +| `bypass_ppo_clip_geo_rs()` | Geo-RS | Bypass (PPO-clip) | - | sequence | PPO-clip + Geo-RS | +| `bypass_ppo_clip_k3_rs()` | K3-RS | Bypass (PPO-clip) | - | sequence | PPO-clip + K3-RS | +| **Bypass Mode (REINFORCE)** (2 policies; explicit IS weights, no PPO clipping) | +| `bypass_pg_is()` | Seq-TIS | Bypass (REINFORCE) | sequence | - | REINFORCE with explicit IS | +| `bypass_pg_geo_rs()` | Geo-RS | Bypass (REINFORCE) | - | sequence | REINFORCE with Geo-RS | +| `bypass_pg_geo_rs_token_tis()` | Geo-RS-Token-TIS | Bypass (REINFORCE) | token | sequence | REINFORCE + Geo-RS + token IS | +| **Other** | +| `disabled()` | - | - | - | - | Metrics only, no correction | + +**Note:** + +- **Bypass mode** sets π_old = π_rollout and uses `loss_type` to select the loss function: + - `"ppo_clip"` (default): PPO clipped objective where ratio = π_θ/π_rollout already handles IS + - `"reinforce"`: REINFORCE with explicit IS weights as π_θ/π_rollout +- Both loss types benefit from rejection sampling (RS) which masks out-of-distribution samples. +- All estimators (Token-TIS, Seq-TIS, Seq-MIS, Geo-RS, ...) are compatible with Decoupled and Bypass modes. + +#### Other Supported Combinations (Manual Configuration Required) + +**Other supported combinations without preset methods:** + +- Token IS + Token RS: Token-level IS weights + Token-level RS mask +- Pure token RS: Token-level RS only, no IS weights +- Pure sequence RS: Sequence-level RS only, no IS weights + +See [detailed configuration examples below](#additional-useful-configurations-not-exposed-as-presets) for manual configurations. + +**Key properties:** + +- Any aggregation level (token/sequence) works in either decoupled or bypass mode +- All combinations are fully supported by the implementation +- Rejection sampling is independent of IS weighting +- Pure RS (`bypass_pg_rs`) uses bypass + geometric RS with `loss_type="reinforce"` (no IS weights) + +--- + +### 1. Decoupled Mode with Token-level Importance Sampling (`decoupled_token_is`) + +**Configuration:** + +```python +config = RolloutCorrectionConfig.decoupled_token_is(threshold=2.0) +``` + +**Components:** + +- **Operating Mode**: Decoupled (3 policies) +- **Loss**: PPO with clipping (only for the second drift correction) +- **IS Aggregation**: Token-level +- **RS**: None (can be added separately) + +**Equivalent YAML:** + +```yaml +algorithm: + rollout_correction: + rollout_is: token + rollout_is_threshold: 2.0 + rollout_rs: null + bypass_mode: false # Decoupled mode +``` + +**Properties:** + +- Independent truncation per token +- Lower variance than sequence-level (product of ratios bounded individually) +- Typical threshold: 1.5 - 5.0 + +**Theory:** See [rollout_corr_math.md §3.3.1](rollout_corr_math.md#331-token-level-aggregation) + +--- + +### 2. Decoupled Mode with Sequence-level Importance Sampling (`decoupled_seq_is`) + +**Also known as: Seq-TIS (Sequence-Level Truncated IS)** + +**Configuration:** + +```python +config = RolloutCorrectionConfig.decoupled_seq_is(threshold=2.0) +``` + +**Components:** + +- **Operating Mode**: Decoupled (3 policies) +- **Loss**: PPO with clipping (only for the second drift correction) +- **IS Aggregation**: Sequence-level (Seq-TIS) +- **RS**: None (can be added separately) + +**Equivalent YAML:** + +```yaml +algorithm: + rollout_correction: + rollout_is: sequence + rollout_is_threshold: 2.0 + rollout_rs: null + bypass_mode: false # Decoupled mode +``` + +**Properties:** + +- Multiplicative aggregation across sequence +- More sensitive to outliers than token-level +- Typical threshold: 2.0 - 10.0 (higher than token-level) + +**Theory:** See [rollout_corr_math.md §3.3.2](rollout_corr_math.md#332-sequence-level-aggregation) + +--- + +### 3. Decoupled Mode with Sequence-level IS + Rejection Sampling (`decoupled_seq_is_rs`) + +**Also known as: Seq-MIS (Sequence-Level Masked IS)** + +**Configuration:** + +```python +config = RolloutCorrectionConfig.decoupled_seq_is_rs(is_threshold=2.0, rs_threshold="0.5_2.0") +``` + +**Components:** + +- **Operating Mode**: Decoupled (3 policies) +- **Loss**: PPO with clipping (only for the second drift correction) +- **IS Aggregation**: Sequence-level (Seq-TIS) +- **RS**: Sequence-level rejection (Seq-MIS) + +**Equivalent YAML:** + +```yaml +algorithm: + rollout_correction: + rollout_is: sequence + rollout_is_threshold: 2.0 + rollout_rs: seq_sum_k1 + rollout_rs_threshold: 0.5_2.0 + bypass_mode: false # Decoupled mode +``` + +**Properties:** + +- Double mechanism: IS reweighting (Seq-TIS) + rejection filtering (Seq-MIS) +- Lower effective sample size (rejects outliers) +- For severe off-policy gaps or when the distribution tail is "toxic" (garbage/adversarial samples) + +**When to use Seq-MIS over Seq-TIS:** + +- **Seq-TIS (clipping only)**: Maximizes information efficiency; extracts signal from all samples. Use when data is clean and mismatch is moderate. +- **Seq-MIS (rejection)**: Maximizes safety; acts as a hard trust region filter. Use when mismatch is severe or when high-weight samples are likely garbage rather than signal. + +**Theory:** See [rollout_corr_math.md §3.5](rollout_corr_math.md#35-rejection-sampling-rs) + +--- + +### 6. Bypass Mode with PPO-clip (`bypass_ppo_clip`) + +**Configuration:** + +```python +config = RolloutCorrectionConfig.bypass_ppo_clip() +``` + +**Components:** + +- **Operating Mode**: Bypass (2 policies: π_rollout = π_old, π_θ) +- **Loss**: PPO-clip (IS handled by ratio, no explicit IS weights) +- **IS Aggregation**: None (PPO ratio handles it) +- **RS**: None + +**Equivalent YAML:** + +```yaml +rollout_correction: + rollout_is: null + rollout_rs: null + bypass_mode: true + loss_type: ppo_clip +``` + +**Properties:** + +- PPO clipped objective in bypass mode +- The PPO ratio = π_θ/π_rollout already handles IS (no explicit IS weights needed) +- Skips `actor.compute_log_prob()` forward pass (2 policies instead of 3) +- No rejection sampling - use `bypass_ppo_clip_geo_rs()` for RS + +**Configuration requirement:** + +- Set `actor_rollout_ref.rollout.calculate_log_probs: true` + +**Additional requirements for bypass mode:** + +- Set `actor_rollout_ref.actor.use_rollout_log_probs: true` +- Set `actor_rollout_ref.actor.policy_loss.loss_mode: bypass_mode` +- Set rollout correction config via `actor_rollout_ref.actor.policy_loss.rollout_correction` + +**Theory:** See [rollout_corr_math.md §3.1.2](rollout_corr_math.md#312-bypass-mode-two-policies) + +--- + +### 7. REINFORCE with IS (`bypass_pg_is`) + +**Configuration:** + +```python +config = RolloutCorrectionConfig.bypass_pg_is(threshold=2.0) +``` + +**Components:** + +- **Operating Mode**: Bypass (2 policies: π_rollout, π_θ) +- **Loss**: REINFORCE (policy gradient with explicit IS weights, no PPO clipping) +- **IS Aggregation**: Sequence-level +- **RS**: None + +**Equivalent YAML:** + +```yaml +rollout_correction: + rollout_is: sequence + rollout_is_threshold: 2.0 + rollout_rs: null + bypass_mode: true + loss_type: reinforce # REINFORCE with explicit IS weights +``` + +**Properties:** + +- REINFORCE loss with explicit IS weights (no PPO clipping) +- Single forward pass (skips old_log_prob computation) +- IS weights computed on-the-fly in loss function + +**Theory:** See [rollout_corr_math.md §3.2.2](rollout_corr_math.md#322-policy-gradient-loss-with-isrs-correction) + +--- + +## Additional Useful Configurations (Not Exposed as Presets) + +These configurations are **fully supported** but don't have convenience preset methods yet. + +### 1. Token IS + Token RS (`token_is_rs`) + +Token-level IS weights with token-level RS mask. + +**Python:** + +```python +config = RolloutCorrectionConfig( + rollout_is="token", + rollout_is_threshold=2.0, + rollout_rs="token_k1", + rollout_rs_threshold=2.0, +) +``` + +**Properties:** Per-token IS weights + per-token RS mask. + +### 2. Pure Token RS (`token_rs`) + +Token-level RS only, no IS weights. + +**Python:** + +```python +config = RolloutCorrectionConfig( + rollout_is=None, + rollout_rs="token_k1", + rollout_rs_threshold=2.0, +) +``` + +**Properties:** Token-level RS mask, no IS reweighting. + +### 3. Pure Sequence RS (`seq_rs`) + +Sequence-level RS only, no IS weights. + +**Python:** + +```python +config = RolloutCorrectionConfig( + rollout_is=None, + rollout_rs="seq_sum_k1", + rollout_rs_threshold="0.5_2.0", +) +``` + +**Properties:** Sequence-level RS mask, no IS reweighting. + +--- + +### Summary: How IS Weights are Processed + +IS weights (`rollout_is_weights`) go through a fixed processing pipeline: + +**Stage 1: Safety Bound (Prevent Overflow)** + +- Token level: `exp(clamp(log_ratio, -20, 20))` per token → bounds each token to [2e-9, 5e8] +- Sequence level: `exp(clamp(sum(log_ratio), -20, 20))` → bounds product to [2e-9, 5e8], broadcast to all tokens + +**Stage 2: Truncation (Reduce Variance)** + +- `.clamp(max=rollout_is_threshold)` → caps weights at upper threshold (TIS: Truncated Importance Sampling) +- No lower truncation (preserves unbiasedness for small weights) + +**Stage 3: Padding Zeroing (Correct Aggregation)** + +- `weights * response_mask` → zeros out padding positions + +**Stage 4: Optional Batch Normalization** + +- If `rollout_is_batch_normalize=True`: Normalize weights to mean=1.0 within batch +- Applied after truncation to preserve truncation semantics + +**Rejection Sampling (Separate Mechanism)** + +Rejection sampling modifies `response_mask` (NOT weights) through `compute_rollout_rejection_mask()`: + +- Computes safety-bounded ratios independently +- Creates binary mask: tokens/sequences outside [lower_threshold, upper_threshold] → 0 (rejected) +- Modified mask used for loss aggregation + +## Operation Modes + +The framework provides **two operating modes** for computing π_old, which can be combined with different loss functions. + +### Operating Modes and Configuration + +| Configuration | `bypass_mode` | `loss_type` | Operating Mode | Loss Function | Description | +| ---------------------- | ------------- | ---------------------- | -------------- | ------------- | ----------------------------------------------------------------- | +| **Decoupled** | `false` | N/A | Decoupled | PPO | Computes `old_log_prob` separately via `actor.compute_log_prob()` | +| **Bypass + PPO-clip** | `true` | `"ppo_clip"` (default) | Bypass | PPO-clip | PPO clipped objective (IS handled by ratio) | +| **Bypass + REINFORCE** | `true` | `"reinforce"` | Bypass | REINFORCE | Policy gradient with explicit IS weights (no PPO clipping) | + +### Operating Mode Details + +#### Decoupled Mode (Three Policies) + +**Policy setup:** + +- π_rollout: Behavior policy (data collection) +- π_old: Proximal policy (computed via `actor.compute_log_prob()` at start of training epoch) +- π_θ: Current policy (being updated) + +**Configuration:** `bypass_mode = false` + +**Properties:** + +- ✅ Achieves batch size invariance +- ✅ Separately corrects Drift 1 (rollout→old) and Drift 2 (old→current) +- ✅ Efficient stale data utilization +- ❌ Extra forward pass needed (`actor.compute_log_prob()`) + +**Theory:** See [rollout_corr_math.md §3.1.1](rollout_corr_math.md#311-decoupled-mode-three-policies) + +#### Bypass Mode (Two Policies) + +**Policy setup:** + +- π_rollout: Behavior policy (data collection) +- π_old = π_rollout: Proximal policy equals behavior policy +- π_θ: Current policy (being updated) + +**Configuration:** `bypass_mode = true` + +**Properties:** + +- ✅ Skips `actor.compute_log_prob()` call (faster) +- ✅ Handles off-policy correction via IS/RS (when using policy gradient with IS/RS) +- ✅ Uses two policies instead of three (π_rollout = π_old) +- ⚠️ Does not separate proximal policy from behavior policy (unlike decoupled mode) + +**Theory:** See [rollout_corr_math.md §3.1.2](rollout_corr_math.md#312-bypass-mode-two-policies) + +--- + +### IS/RS Aggregation Levels (Orthogonal to Operating Mode) + +The aggregation level can be chosen **independently** of the operating mode. Any aggregation level works in either decoupled or bypass mode. + +| `rollout_is` | `rollout_rs` | Behavior | +| ------------------------- | ------------------------------------------------------------------ | --------------------------------------------------------------------------------- | +| `null` | `null` | **Disabled**: No computation, no metrics, no rejection | +| `null` | `"token_k1"`, `"seq_sum_k1"`, `"seq_mean_k1"`, `"seq_max_k2"`, etc | **Rejection only**: Compute metrics, NO weight correction, YES rejection sampling | +| `"token"` or `"sequence"` | `null` | **IS weights only**: Weight correction enabled, NO rejection sampling | +| `"token"` or `"sequence"` | `"token_k1"`, `"seq_sum_k1"`, `"seq_mean_k1"`, `"seq_max_k2"`, etc | **Full correction**: Both weight correction and rejection sampling enabled | + +### Key Insights + +- ✅ Any IS/RS aggregation level (token/sequence/geometric) can be used in **either** decoupled or bypass mode +- ✅ You can use **rejection sampling alone** without IS weight correction (`rollout_is=null, rollout_rs="token_k1"`) +- ✅ You can use **IS weights alone** without outlier rejection (`rollout_is="token", rollout_rs=null`) +- ✅ You can use **both together** (`rollout_is="token", rollout_rs="token_k1"`) +- ✅ You can **monitor metrics only** without any correction by setting both to `null` but still providing rollout_log_probs + +**Theory:** See [rollout_corr_math.md §3.3](rollout_corr_math.md#33-isrs-aggregation-levels) for details on aggregation levels. + +### Example Workflow + +**Recommended: Bypass Mode** + +This workflow uses bypass mode for efficiency. + +1. **Start with metrics only** to understand the off-policy gap: + + ```yaml + rollout_correction: + rollout_is: null + rollout_rs: null + bypass_mode: true # Bypass mode (recommended) + loss_type: ppo_clip # Default: PPO clipped objective + ``` + + Monitor `rollout_corr/kl`, `rollout_corr/log_ppl_abs_diff`, `rollout_corr/chi2_token` to assess off-policy gap. + +2. **Enable rejection sampling** if you see high outlier fractions: + + ```yaml + rollout_correction: + rollout_is: null + rollout_rs: sequence # or "geometric" for higher sensitivity + rollout_rs_threshold: 2.0 + bypass_mode: true # Bypass mode + loss_type: ppo_clip # or "reinforce" for explicit IS weights + ``` + + This excludes outliers from training without modifying gradients. + +3. **Enable full IS correction** (with REINFORCE loss) once comfortable with metrics: + ```yaml + rollout_correction: + rollout_is: sequence # Recommended: unbiased, suitable for most cases + rollout_is_threshold: 2.0 + rollout_rs: sequence # or "geometric" for more aggressive filtering + rollout_rs_threshold: 2.0 + bypass_mode: true # Bypass mode + loss_type: reinforce # REINFORCE with explicit IS weights + ``` + +**Benefits of bypass mode:** + +- ✅ Skips expensive `actor.compute_log_prob()` forward pass (faster) +- ✅ `loss_type` controls the loss function: "ppo_clip" (default) or "reinforce" +- ✅ PPO-clip: IS handled by ratio (no explicit weights), RS mask applied +- ✅ REINFORCE: Explicit IS weights computed on-the-fly (π_θ / π_rollout) +- ✅ Both loss types work with all IS/RS combinations + +## Usage + +### Basic Setup + +```yaml +algorithm: + rollout_correction: + rollout_is: token # Enable IS weights at token level + rollout_is_threshold: 2.0 # Threshold for IS weights + rollout_rs: null # No rejection sampling + +actor_rollout_ref: + rollout: + calculate_log_probs: true # Required! +``` + +### Additional Configurations for Bypass Mode + +- Set `actor_rollout_ref.actor.use_rollout_log_probs: true` +- Set `actor_rollout_ref.actor.policy_loss.loss_mode: bypass_mode` +- Set rollout correction config via `actor_rollout_ref.actor.policy_loss.rollout_correction` + +### Metrics + +All metrics are prefixed with `rollout_corr/` in logs. For example, `rollout_is_mean` appears as `rollout_corr/rollout_is_mean`. + +These metrics cover both: + +- **Diagnostic metrics**: KL divergence, perplexity differences (measuring off-policy gap) +- **Correction statistics**: IS weights, rejection rates (measuring correction applied) + +#### **Core IS Weight Metrics** + +- **`rollout_is_mean`**: Mean importance sampling weight across all valid tokens + + - Value close to 1.0 indicates minimal off-policy gap + +- **`rollout_is_std`**: Standard deviation of IS weights + + - Higher values indicate greater variance in IS weights + +- **`rollout_is_min`**: Minimum IS weight observed + + - Shows the most underweighted token/sequence + - For sequence/geometric: computed from unclamped log-space ratios (true minimum) + - For token: computed from safety-bounded weights + +- **`rollout_is_max`**: Maximum IS weight observed + - Shows the most overweighted token/sequence + - For sequence/geometric: computed from unclamped log-space ratios (true maximum before safety bound) + - For token: computed from safety-bounded weights (before threshold clamping) + - Compare with `rollout_is_threshold` to see truncation impact + +#### **Effective Sample Size** + +- **`rollout_is_eff_sample_size`**: Effective sample size after IS weighting + - **Formula**: `1 / mean(weights²)` where weights are normalized + - **Range**: 0.0 to 1.0 (as fraction of original batch) + - Lower values indicate weight concentration on fewer samples + +#### **Threshold Exceedance Metrics** + +- **`rollout_is_ratio_fraction_high`**: Fraction of weights exceeding upper threshold + + - Shows how often truncation/masking occurs on high end + - For sequence/geometric: computed from unclamped log-space ratios (true exceedance) + - For token: computed from safety-bounded weights (before threshold clamping) + +- **`rollout_is_ratio_fraction_low`**: Fraction of weights below lower threshold (1/upper_threshold) + - Diagnostic metric showing how many weights are below the reciprocal threshold + - For sequence/geometric: computed from unclamped log-space ratios (true exceedance) + - For token: computed from safety-bounded weights (before truncation) + +#### **Sequence-Level Metrics** (for sequence aggregation) + +- **`rollout_is_seq_mean`**: Mean IS weight at sequence level + + - Should match `rollout_is_mean` for sequence-level aggregation + +- **`rollout_is_seq_std`**: Standard deviation of sequence-level IS weights + +- **`rollout_is_seq_min`**: Minimum sequence-level IS weight + +- **`rollout_is_seq_max`**: Maximum sequence-level IS weight + +- **`rollout_is_seq_max_deviation`**: Maximum absolute deviation from 1.0 at sequence level + + - Shows worst-case sequence off-policy gap + +- **`rollout_is_seq_fraction_high`**: Fraction of sequences exceeding upper threshold + +- **`rollout_is_seq_fraction_low`**: Fraction of sequences below lower threshold + +#### **Rejection Sampling Metrics** (when `rollout_rs` is enabled) + +- **`rollout_rs_masked_fraction`**: Fraction of tokens rejected via rejection sampling + + - **Important**: Rejection sampling modifies `response_mask` (sets rejected tokens to 0) + - **Separate from IS weights**: IS weights are still truncated; rejection is an independent filtering step + - Only present when `rollout_rs` is enabled (token/sequence/geometric) + +- **`rollout_rs_seq_masked_fraction`**: Fraction of sequences with at least one rejected token + - Shows sequence-level impact of rejection sampling + - Token-level RS: sequence rejected if ANY token is outside [lower, upper] + - Sequence-level RS: entire sequence rejected or accepted based on sequence-level ratio + - Geometric RS: entire sequence rejected or accepted based on geometric mean + +#### **Off-Policy Diagnostic Metrics** (Training vs Rollout Policy) + +**Note on terminology:** These metrics use "training" to refer to the training reference policy and "rollout" to refer to π_rollout (the behavior policy used for data collection). + +- **Decoupled mode**: "training" = π_old (computed at start of training epoch) +- **Bypass/Pure IS mode**: "training" = π_θ (current policy being trained) + +In bypass/pure IS mode, metrics measure the drift between π_θ and π_rollout directly. + +- **`training_ppl`**: Perplexity of training reference policy (π_old in decoupled mode, π_θ in bypass/pure IS mode) + + - **Formula**: `exp(-mean(log_probs))` + - Lower values indicate higher model confidence + +- **`rollout_ppl`**: Perplexity of rollout policy π_rollout (e.g., vLLM BF16) + +- **`ppl_ratio`**: Ratio of training PPL to rollout PPL + + - **Formula**: `exp(mean(log(training_ppl / rollout_ppl)))` + - **Meaning**: > 1.0 means training is less confident than rollout + +- **`training_log_ppl`**: Log perplexity of training policy + + - Useful for identifying trends (linear scale) + +- **`rollout_log_ppl`**: Log perplexity of rollout policy + +- **`log_ppl_diff`**: Mean difference in log perplexities + + - **Formula**: `mean(log_ppl_rollout - log_ppl_training)` + - Sign indicates which policy is more confident + +- **`log_ppl_abs_diff`**: Mean absolute log perplexity difference + + - Magnitude of off-policy gap regardless of direction + +- **`log_ppl_diff_max`**: Maximum log perplexity difference across sequences + + - Identifies worst-case sequence + +- **`log_ppl_diff_min`**: Minimum log perplexity difference across sequences + +- **`kl`**: KL divergence KL(π_rollout || π_training) + + - **Formula**: `mean(log_prob_rollout - log_prob_training)` + - **Note**: Can be negative (rollout is less confident) + +- **`k3_kl`**: K3 divergence (equals KL(π_rollout || π_training) in expectation) + + - **Formula**: `mean(exp(log_ratio) - log_ratio - 1)` + - More stable than direct KL (non-negative per token) + - Always >= 0 + +- **`chi2_token`**: Chi-squared divergence at token level + + - **Formula**: `mean(ratio²) - 1` where ratio = π_training/π_rollout + - Measures second moment of IS weight distribution + - Always non-negative + +- **`chi2_seq`**: Chi-squared divergence at sequence level + - **Formula**: `mean((∏_t ratio_t)²) - 1` + - Sequence-level second moment of IS weights + - More sensitive than token-level chi-squared + +#### **Example: Accessing Metrics in Code** + +```python +# Metrics are returned from compute_rollout_correction_and_rejection_mask +from verl.trainer.ppo.rollout_corr_helper import compute_rollout_correction_and_rejection_mask + +# Returns 3 values (weights, modified_response_mask, metrics) +weights_proto, modified_response_mask, metrics = compute_rollout_correction_and_rejection_mask( + old_log_prob=training_log_probs, # from training policy + rollout_log_prob=rollout_log_probs, # from rollout policy + response_mask=response_mask, + rollout_is="token", # Enable IS weights at token level + rollout_is_threshold=2.0, + rollout_rs="token_k1", + rollout_rs_threshold="0.5_2.0", +) + +# Extract IS weights (processed, zeroed at padding) +is_weights = weights_proto.batch["rollout_is_weights"] + +# IS weights processing (with IS enabled at token level): +# 1. Safety-bounded: exp(clamp(log_ratio, -20, 20)) per token +# 2. Truncated: .clamp(max=2.0) to cap extreme weights +# 3. Zeroed at padding positions +# Note: Truncation is ALWAYS applied to IS weights (TIS: Truncated Importance Sampling) + +# modified_response_mask has rejection applied (since rollout_rs="token_k1"): +# 1. RS rejection: tokens outside [0.5, 2.0] masked to 0 via response_mask +# Note: RS and IS are separate mechanisms - both can be enabled independently + +# All metrics have 'rollout_corr/' prefix +print(f"Mean IS weight: {metrics['rollout_corr/rollout_is_mean']:.3f}") +print(f"Effective sample size: {metrics['rollout_corr/rollout_is_eff_sample_size']:.3f}") +print(f"RS masked fraction: {metrics['rollout_corr/rollout_rs_masked_fraction']:.3f}") +print(f"KL divergence: {metrics['rollout_corr/kl']:.3f}") + +# Check IS weights for valid tokens (non-padding) +valid_weights = is_weights[response_mask.bool()] +print(f"\n✓ IS weights min (valid tokens): {valid_weights.min():.4f}") +print(f"✓ IS weights max (valid tokens): {valid_weights.max():.4f}") +print(f"✓ All valid IS weights > 0: {(valid_weights > 0).all()}") +print(f"✓ IS weights are capped at threshold: {(valid_weights <= 2.0).all()}") + +# Check rejection via response_mask +rejected_tokens = (response_mask == 1) & (modified_response_mask == 0) +print(f"\n✓ Rejected {rejected_tokens.sum()} tokens via response_mask") +print(f"✓ Rejection sampling modifies response_mask (separate from IS weight truncation)") +print(f"✓ IS weights are always truncated to [0, threshold] after safety bounding") + +# Check for warning conditions +if metrics['rollout_corr/rollout_is_mean'] < 0.5 or metrics['rollout_corr/rollout_is_mean'] > 2.0: + print("⚠️ Warning: Mean IS weight far from 1.0, significant off-policy gap detected") + +if metrics['rollout_corr/rollout_is_eff_sample_size'] < 0.3: + print("⚠️ Warning: Low effective sample size, high weight concentration") +``` + +#### **Example: Monitoring Metrics During Training** + +```python +# In your training loop +for epoch in range(num_epochs): + for batch_idx, batch in enumerate(dataloader): + # ... rollout phase ... + + # Compute IS weights and get metrics + rollout_corr_config = config.algorithm.get("rollout_correction", None) + if rollout_corr_config is not None: + weights_proto, modified_response_mask, metrics = compute_rollout_correction_and_rejection_mask( + old_log_prob=batch.old_log_prob, + rollout_log_prob=batch.rollout_log_prob, + response_mask=batch.response_mask, + rollout_is=rollout_corr_config.get("rollout_is", None), + rollout_is_threshold=rollout_corr_config.get("rollout_is_threshold", 2.0), + rollout_rs=rollout_corr_config.get("rollout_rs", None), + rollout_rs_threshold=rollout_corr_config.get("rollout_rs_threshold", None), + ) + + # Log to tensorboard/wandb + for metric_name, metric_value in metrics.items(): + logger.log_scalar(metric_name, metric_value, step=global_step) + + # IMPORTANT: Update batch response_mask with rejection applied + batch.response_mask = modified_response_mask + + # Use IS weights in training (always safety-bounded, zeroed at padding) + is_weights = weights_proto.batch["rollout_is_weights"] + # ... apply weights to policy gradient ... +``` + +#### **Example: Conditional Alerting Based on Metrics** + +```python +def check_rollout_correction_health(metrics, config): + """Check if Rollout Correction metrics indicate healthy training.""" + warnings = [] + + # Check mean IS weight + mean_weight = metrics['rollout_corr/rollout_is_mean'] + if mean_weight < 0.5 or mean_weight > 2.0: + warnings.append(f"Mean IS weight {mean_weight:.3f} is far from 1.0") + + # Check effective sample size + ess = metrics['rollout_corr/rollout_is_eff_sample_size'] + if ess < 0.3: + warnings.append(f"Effective sample size {ess:.3f} is too low") + + # Check standard deviation + std = metrics['rollout_corr/rollout_is_std'] + if std > 1.0: + warnings.append(f"IS weight std {std:.3f} is too high") + + # Check KL divergence + kl = metrics['rollout_corr/kl'] + if abs(kl) > 0.1: + warnings.append(f"KL divergence {kl:.3f} indicates significant off-policy gap") + + # Check chi-squared divergence + if 'rollout_corr/chi2_token' in metrics: + chi2_token = metrics['rollout_corr/chi2_token'] + if chi2_token > 1.0: + warnings.append(f"Chi-squared divergence (token) {chi2_token:.3f} indicates severe distribution shift") + + if warnings: + print("⚠️ Rollout Correction Health Warnings:") + for warning in warnings: + print(f" - {warning}") + return False + else: + print("✅ Rollout Correction metrics look healthy") + return True + +# Use in training +_, _, metrics = compute_rollout_correction_and_rejection_mask(...) +is_healthy = check_rollout_correction_health(metrics, config) + +if not is_healthy: + # Consider adjusting config or investigating issues + print("Consider:") + print(" - Tightening rollout_is_threshold") + print(" - Switching to geometric aggregation level") + print(" - Checking if rollout and training policies are too different") +``` + +### Running Examples + +Start with the basic token-level truncate configuration: + +```bash +bash examples/rollout_correction/run_qwen2_5_7b_fsdp.sh +``` + +Monitor metrics for 1-2 epochs before adjusting parameters. + +## Configuration Examples + +### Example 1: IS Weights Only (Token Level) + +```yaml +algorithm: + rollout_correction: + rollout_is: token + rollout_is_threshold: 2.0 + rollout_rs: null # No rejection sampling +``` + +### Example 2: Rejection Sampling Only (No IS Weights) + +```yaml +algorithm: + rollout_correction: + rollout_is: null # No IS weights + rollout_rs: token_k1 + rollout_rs_threshold: "0.5_2.0" +``` + +### Example 3: Both IS and RS (Token RS) + +```yaml +algorithm: + rollout_correction: + rollout_is: token + rollout_is_threshold: 2.0 + rollout_rs: token_k1 + rollout_rs_threshold: "0.5_2.0" +``` + +### Example 5: Bypass Mode with PPO-clip (Default) + +```yaml +algorithm: + rollout_correction: + rollout_is: token + rollout_is_threshold: 2.0 + rollout_rs: token_k1 + rollout_rs_threshold: "0.5_2.0" + bypass_mode: true # Skip old_log_prob computation + loss_type: ppo_clip # PPO clipped objective (default) +``` + +**Skips expensive `actor.compute_log_prob()` forward pass. PPO ratio = π_θ/π_rollout handles IS.** + +### Example 6: Bypass Mode with REINFORCE + +```yaml +rollout_correction: + rollout_is: sequence # Explicit IS correction in loss + rollout_is_threshold: 2.0 + rollout_rs: null # Optional: can add rejection sampling + bypass_mode: true + loss_type: reinforce # REINFORCE with explicit IS weights +``` + +**No PPO clipping, pure policy gradient with IS correction** + +### Example 7: Bypass Mode with PPO-clip + Rejection Sampling + +```yaml +rollout_correction: + rollout_is: sequence # Computed for metrics + rollout_is_threshold: 2.0 + rollout_rs: seq_max_k2 # Sequence max χ²/2 guard + rollout_rs_threshold: 2.5 + bypass_mode: true + loss_type: ppo_clip # PPO clipped objective (IS handled by ratio) +``` + +**PPO clipping with rejection sampling. IS handled by PPO ratio (no explicit IS weights).** + +## Troubleshooting + +### Issue: High spread in IS weights + +**Symptoms:** `rollout_is_std` > 1.0, `rollout_is_eff_sample_size` < 0.3 + +**Solutions:** + +1. Switch from `sequence` to `geometric` level +2. Tighten thresholds +3. Verify rollout and training aren't too different + +### Issue: Mean IS weight far from 1.0 + +**Symptoms:** `rollout_is_mean` < 0.5 or > 2.0 + +**Solutions:** + +1. Verify `calculate_log_probs=True` is set +2. Check rollout_log_probs are correctly passed +3. Check for systematic distribution shift + +### Debugging: Visualizing Metrics + +**Example: Plot IS weight distribution** + +```python +import matplotlib.pyplot as plt +import numpy as np + +def plot_is_metrics(metrics_history): + """Plot rollout IS metrics over training steps.""" + fig, axes = plt.subplots(2, 3, figsize=(15, 10)) + + # Plot 1: Mean IS weight over time + axes[0, 0].plot(metrics_history['rollout_corr/rollout_is_mean']) + axes[0, 0].axhline(y=1.0, color='r', linestyle='--', label='Ideal') + axes[0, 0].set_title('Mean IS Weight') + axes[0, 0].set_xlabel('Step') + axes[0, 0].legend() + + # Plot 2: Effective sample size + axes[0, 1].plot(metrics_history['rollout_corr/rollout_is_eff_sample_size']) + axes[0, 1].axhline(y=0.5, color='g', linestyle='--', label='Good') + axes[0, 1].axhline(y=0.3, color='r', linestyle='--', label='Warning') + axes[0, 1].set_title('Effective Sample Size') + axes[0, 1].set_xlabel('Step') + axes[0, 1].legend() + + # Plot 3: KL divergence over time + axes[1, 0].plot(metrics_history['rollout_corr/kl'], label='KL') + axes[1, 0].plot(metrics_history['rollout_corr/k3_kl'], label='K3 KL') + axes[1, 0].axhline(y=0, color='g', linestyle='--', alpha=0.3) + axes[1, 0].set_title('KL Divergence') + axes[1, 0].set_xlabel('Step') + axes[1, 0].legend() + + # Plot 4: PPL ratio over time + axes[1, 1].plot(metrics_history['rollout_corr/ppl_ratio']) + axes[1, 1].axhline(y=1.0, color='r', linestyle='--', label='Ideal') + axes[1, 1].set_title('PPL Ratio (Training/Rollout)') + axes[1, 1].set_xlabel('Step') + axes[1, 1].legend() + + # Plot 5: Chi-squared divergence + if 'rollout_corr/chi2_token' in metrics_history: + axes[1, 2].plot(metrics_history['rollout_corr/chi2_token'], label='Token-level') + if 'rollout_corr/chi2_seq' in metrics_history: + axes[1, 2].plot(metrics_history['rollout_corr/chi2_seq'], label='Seq-level') + axes[1, 2].axhline(y=1.0, color='r', linestyle='--', label='Warning') + axes[1, 2].set_title('Chi-squared Divergence') + axes[1, 2].set_xlabel('Step') + axes[1, 2].legend() + else: + axes[1, 2].axis('off') + + plt.tight_layout() + plt.savefig('rollout_is_metrics.png', dpi=150) + print("Saved plot to rollout_is_metrics.png") +``` + +**Example: Metric collection during training** + +```python +# Collect metrics over time +metrics_history = { + 'rollout_corr/rollout_is_mean': [], + 'rollout_corr/rollout_is_eff_sample_size': [], + 'rollout_corr/kl': [], + 'rollout_corr/k3_kl': [], + 'rollout_corr/ppl_ratio': [], + 'rollout_corr/chi2_token': [], + 'rollout_corr/chi2_seq': [], +} + +# In training loop +for step in range(num_steps): + # ... compute IS weights and rejection mask ... + _, _, metrics = compute_rollout_correction_and_rejection_mask(...) + + # Store metrics + for key in metrics_history.keys(): + if key in metrics: + metrics_history[key].append(metrics[key]) + + # Plot every 100 steps + if step % 100 == 0: + plot_is_metrics(metrics_history) +``` + +## Performance Impact + +- **Memory overhead**: ~1% of model memory +- **Computational overhead**: 1-3% depending on level +- **Training stability**: Significantly improved when off-policy gap exists + +## Testing + +Run the test suite to verify everything works: + +```bash +# Basic unit tests +python tests/trainer/ppo/test_rollout_corr.py + +# Integration tests (if pytest is available) +pytest tests/trainer/ppo/test_rollout_corr_integration.py -v +``` + +Expected output: All tests pass ✓ + +## Additional Resources + +- **Implementation**: `verl/trainer/ppo/rollout_corr_helper.py` +- **Examples**: `examples/rollout_correction/` +- **DAPO Example**: `recipe/dapo/run_dapo_qwen2.5_32b_rollout_corr.sh` + +## Summary + +Rollout Correction provides a unified framework for handling general off-policy problems in RL: + +- ✅ Corrects ANY distribution shift between data collection and training +- ✅ Supports diverse scenarios: policy mismatch, staleness, replay buffers, off-policy algorithms +- ✅ Numerical stability with safety bounds and rejection mechanisms +- ✅ Comprehensive diagnostics: KL, perplexity, χ² divergence +- ✅ Flexible methods from token-level to sequence-level aggregation +- ✅ Memory-efficient implementation + +## References + +- **[Mathematical Formulations](rollout_corr_math.md)** - Detailed mathematical theory and derivations for all rollout correction methods +- [Your Efficient RL Framework Secretly Brings You Off-Policy RL Training](https://fengyao.notion.site/off-policy-rl) diff --git a/verl/docs/algo/rollout_corr_math.md b/verl/docs/algo/rollout_corr_math.md new file mode 100644 index 0000000000000000000000000000000000000000..7118c1ddb4987d4211659af14ca5f6567bce08fe --- /dev/null +++ b/verl/docs/algo/rollout_corr_math.md @@ -0,0 +1,961 @@ +# Mathematical Formulations of Rollout Correction Methods in `verl` + +**Author:** [Yingru Li](https://richardli.xyz) +**Last updated:** 2025-11-04 + +--- + +> **📖 Documentation Structure** +> - **This document** - Mathematical theory: formulations, derivations, and algorithmic foundations +> - **[Rollout Correction Usage Guide](rollout_corr.md)** - Practical implementation: configurations, presets, troubleshooting +> +> Start here for theory and design rationale, refer to the usage guide for implementation. + +--- + +### BibTeX Citation + +```bibtex +@online{liu-li-2025-rl-collapse, + title = {When Speed Kills Stability: Demystifying {RL} Collapse from the Training-Inference Mismatch}, + author = {Liu, Jiacai and Li, Yingru and Fu, Yuqian and Wang, Jiawei and Liu, Qian and Shen, Yu}, + year = {2025}, + month = sep, + url = {https://richardli.xyz/rl-collapse} +} + + +@article{li2025trust, + title={Trust Region Masking for Long-Horizon LLM Reinforcement Learning}, + author={Li, Yingru and Liu, Jiacai and Xu, Jiawei and Tong, Yuxuan and Li, Ziniu and Liu, Qian and Wang, Baoxiang}, + journal={arXiv preprint arXiv:2512.23075}, + year={2025} +} +``` + +### Blog Series + +- Main blog post: https://richardli.xyz/rl-collapse +- [Part 1: Why Mismatch Breaks LLM-RL](https://richardli.xyz/rl-collapse-1) (analytical framework using TV distance for bias and χ²-divergence for variance) +- [Part 2: The Gradient Estimator Trials](https://richardli.xyz/rl-collapse-2) (token-level vs sequence-level correction bias-variance tradeoff) +- [Part 3: When Math Meets Reality—Toxic Tails and Length Traps](https://richardli.xyz/rl-collapse-3) (why rejection over clipping, and geometric-level RS) +- Latest Paper: https://arxiv.org/abs/2512.23075 + +## Abstract + +This document provides the definitive mathematical formulations for rollout correction methods in `verl`, following the natural progression from **REINFORCE** to **PPO** to **Decoupled PPO**. + +Rollout correction provides a unified framework to handle **general off-policy problems** in RL training - any scenario where the data collection distribution differs from the training distribution. + +**Applicable scenarios include:** +- **Policy mismatch**: Different precision (FP8 vs FP16 vs BF16 vs FP32), different backends (vLLM vs SGLang vs FSDP vs Megatron) +- **Temporal lag**: Model staleness, asynchronous rollout workers +- **Replay buffers**: Training on historical trajectories from earlier policy versions +- **Off-policy algorithms**: Behavioral cloning, DAPO, expert demonstrations +- **Data filtering**: Reweighting, preference learning, curriculum learning + +--- + +## Table of Contents + +1. [Theoretical Foundation: From REINFORCE to Decoupled PPO](#1-theoretical-foundation-from-reinforce-to-decoupled-ppo) +2. [Implementation in verl: The Three-Policy Framework](#2-implementation-in-verl-the-three-policy-framework) +3. [Algorithmic Components and Combinations](#3-algorithmic-components-and-combinations) +4. [Off-Policy Diagnostic Metrics](#4-off-policy-diagnostic-metrics) +5. [Summary and Decision Guide](#5-summary-and-decision-guide) +6. [Implementation References](#6-implementation-references) + +--- + +## 1. Theoretical Foundation: From REINFORCE to Decoupled PPO + +This section establishes the theoretical progression that `verl` implements. + +### 1.1 REINFORCE: Policy Gradient Baseline + +The REINFORCE algorithm ([Williams, 1992](https://doi.org/10.1007/BF00992696)) is the foundation of policy gradient methods. + +**Vanilla REINFORCE (On-Policy)** + +For trajectories $\tau = (s_0, a_0, s_1, a_1, \ldots, s_T, a_T)$ sampled from the current policy $\pi_\theta$, the policy gradient is: + +$$ +\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta} \left[ \sum_{t=0}^T \nabla_\theta \log \pi_\theta(a_t|s_t) \cdot A_t \right] +$$ + +where $A_t$ is the advantage function at timestep $t$. + +**Off-Policy REINFORCE** + +When trajectories are sampled from a different behavior policy $\mu$, we apply importance sampling over the **joint trajectory distribution**: + +$$ +\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim \mu} \left[ \frac{P_{\pi_\theta}(\tau)}{P_\mu(\tau)} \sum_{t=0}^T \nabla_\theta \log \pi_\theta(a_t|s_t) \cdot A_t \right] +$$ + +where the trajectory-level importance weight is: + +$$ +\frac{P_{\pi_\theta}(\tau)}{P_\mu(\tau)} = \frac{p(s_0) \prod_{t=0}^T \pi_\theta(a_t|s_t) p(s_{t+1}|s_t, a_t)}{p(s_0) \prod_{t=0}^T \mu(a_t|s_t) p(s_{t+1}|s_t, a_t)} = \prod_{t=0}^T \frac{\pi_\theta(a_t|s_t)}{\mu(a_t|s_t)} +$$ + +The transition dynamics $p(s_{t+1}|s_t, a_t)$ and initial state $p(s_0)$ cancel out, leaving only the product of per-step action probability ratios. + +**Key properties:** +- **Off-policy capable**: Can learn from any behavior policy via importance sampling +- **No trust region**: Policy updates not constrained + +**Implementation in verl:** The `bypass_pg_is` preset implements off-policy REINFORCE with truncated importance sampling. + +### 1.2 PPO: Adding Trust Region Control + +Proximal Policy Optimization ([Schulman et al., 2017](https://arxiv.org/abs/1707.06347)) adds a clipped surrogate objective: + +$$ +L_{\text{PPO}}(\theta) = -\mathbb{E}_{(s,a) \sim \mu} \left[ \min\left( r_t(\theta) A_t, \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon) A_t \right) \right] +$$ + +where $r_t(\theta) = \frac{\pi_\theta(a_t|s_t)}{\mu(a_t|s_t)}$ and $\epsilon$ is the clip range (typically 0.2). + +**Key properties:** +- **Two policies**: $\mu$ (reference for clipping) and $\pi_\theta$ (being updated) +- **Trust region via clipping**: Limits policy update magnitude via ratio $r_t(\theta) = \frac{\pi_\theta}{\mu}$ + +### 1.3 Decoupled PPO: Achieving Batch Size Invariance + +Decoupled PPO ([Hilton et al., 2021](https://arxiv.org/abs/2110.00641)) solves PPO's batch size sensitivity by **decoupling two roles**: +1. **Proximal policy** $\pi_{\text{prox}}$: The anchor policy for PPO clipping (controls policy update size) +2. **Behavior policy** $\mu$: The policy that collected the data (for off-policy correction via importance sampling) + +**The problem**: Standard PPO controls policy update size via the ratio $\frac{\pi_\theta}{\pi_{\text{old}}}$, where $\pi_{\text{old}}$ is assumed to be both the proximal policy *and* the behavior policy. This coupling makes the algorithm sensitive to batch size because aggregating data from multiple workers or using replay buffers changes the effective behavior policy. + +**The solution**: Decouple these two roles, leading to a **three-policy formulation**: + +$$ +L_{\text{DecoupledPPO}}(\theta) = -\mathbb{E}_{(s,a) \sim \mu} \left[ w_t \cdot \min\left( r_t(\theta) A_t, \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon) A_t \right) \right] +$$ + +where: +- $w_t = \frac{\pi_{\text{prox}}(a_t|s_t)}{\mu(a_t|s_t)}$: Importance sampling weight (corrects for behavior policy $\mu$). Here $\pi_{\text{prox}}$ is frozen during training, so $w_t$ is constant (no stopgrad operator needed). +- $r_t(\theta) = \frac{\pi_\theta(a_t|s_t)}{\pi_{\text{prox}}(a_t|s_t)}$: PPO ratio (controls policy update size against proximal policy $\pi_{\text{prox}}$) + +**Key properties**: By decoupling: +- **Batch size invariance**: Policy update control (via $\pi_{\text{prox}}$) is independent of data aggregation +- **Flexible behavior policy**: Any $\mu$ can be used (different workers, replay buffers, or stale checkpoints) +- **Stale data utilization**: Older trajectories can be corrected via importance sampling +- **Clipping preserved**: Clipping against $\pi_{\text{prox}}$ limits update magnitude + +**This is the algorithm that `verl` implements via its three-policy framework.** + +--- + +## 2. Implementation in verl: The Three-Policy Framework + +The `verl` library implements decoupled PPO using three distinct policies, each serving a specific role. + +### 2.1 Policy Roles and Notation + +**$\pi_{\text{rollout}}$ (Behavior Policy $\mu$)** +The policy used for data collection. This is the behavior distribution $\mu$ from theory. + +- **When created**: During rollout/data collection phase +- **Purpose**: Generate trajectories for training +- **Common sources**: + - Policy mismatch: Same weights, different implementation (precision, backend) + - Temporal lag: Stale checkpoint from async workers + - Replay buffer: Historical data from earlier iterations + - Off-policy algorithms: Expert demonstrations, auxiliary policies (DAPO) + - Data filtering: Reweighted or filtered data +- **Fixed**: Frozen during training on a batch + +**$\pi_{\text{old}}$ (Proximal Policy $\pi_{\text{prox}}$)** +The reference policy for PPO clipping. This is the "proximal policy" from decoupled PPO theory. + +- **When created**: + - **Decoupled mode**: Computed at start of training epoch via `actor.compute_log_prob()` + - **Bypass mode**: Set equal to $\pi_{\text{rollout}}$ (skips separate computation) +- **Purpose**: + - Anchor point for PPO clipping (controls policy update size) + - When separate from $\pi_{\text{rollout}}$: Enables batch size invariance and efficient use of stale data +- **Fixed**: Frozen during all PPO update epochs on the same batch + +**$\pi_{\theta}$ (Current Policy)** +The policy being actively optimized during training. + +- **Updated**: Every gradient step +- **Purpose**: The policy we're improving + +### 2.2 Operating Modes + +The three-policy framework can operate in two modes: + +**Decoupled Mode (Three Policies)** +- Computes $\pi_{\text{old}}$ separately at the start of each training epoch +- **Algorithm**: Full decoupled PPO with three policies (mathematically correct) +- **Properties**: Achieves batch size invariance; separately corrects Drift 1 (rollout→old) and Drift 2 (old→current) + +**Bypass Mode (Two Policies)** +- Sets $\pi_{\text{old}} = \pi_{\text{rollout}}$ (skips separate computation) +- **Algorithm**: Uses $\pi_{\text{rollout}}$ as both behavior policy and proximal policy (mathematically correct) +- **Key difference**: Proximal policy equals behavior policy, so no IS correction needed between them +- **Properties**: Faster (skips `actor.compute_log_prob()` call); does not achieve batch size invariance + +### 2.3 Two Distribution Shifts + +The three-policy framework handles two types of distribution drift: + +**Drift 1: $\pi_{\text{rollout}} \to \pi_{\text{old}}$ (Off-Policy Gap)** + +This is the distribution shift between the data collection policy and the training reference policy. + +- **Nature**: Ranges from negligible (same checkpoint, minor differences) to severe (replay buffers, expert data) +- **Correction**: Importance sampling weight $w_t = \frac{\pi_{\text{old}}(a_t|s_t)}{\pi_{\text{rollout}}(a_t|s_t)}$ +- **Optional**: Can be ignored (bypass mode) when negligible + +**Drift 2: $\pi_{\text{old}} \to \pi_{\theta}$ (Policy Update Drift)** + +This is the drift from policy parameter updates during training. + +- **Nature**: Occurs as $\pi_\theta$ is updated via gradient descent +- **Correction**: PPO clipping on ratio $r_t(\theta) = \frac{\pi_\theta(a_t|s_t)}{\pi_{\text{old}}(a_t|s_t)}$ +- **Universal**: Applies to both on-policy and off-policy training + +### 2.4 Notation Summary + +- $\pi_{\text{rollout}}$: Behavior policy (data collection) +- $\pi_{\text{old}}$: Proximal policy (PPO anchor) +- $\pi_{\theta}$: Current policy (being updated) +- $\rho_t = \frac{\pi_{\text{old}}(a_t|s_t)}{\pi_{\text{rollout}}(a_t|s_t)}$: Per-token IS ratio (corrects Drift 1) +- $r_t(\theta) = \frac{\pi_{\theta}(a_t|s_t)}{\pi_{\text{old}}(a_t|s_t)}$: PPO ratio (corrects Drift 2) +- $A_t$: Advantage at token $t$ +- $T$: Set of valid tokens in a sequence +- $C_{\text{IS}}$: Upper threshold for IS weights (e.g., 2.0) +- $C_{\text{RS-upper}}$: Upper threshold for RS mask (e.g., 2.0) +- $C_{\text{RS-lower}}$: Lower threshold for RS mask (typically $1/C_{\text{RS-upper}}$) +- $\epsilon$: PPO clip range (typically 0.2) + +--- + +## 3. Algorithmic Components and Combinations + +The rollout correction framework in `verl` is built from **orthogonal components** that can be combined flexibly: + +1. **Operating Mode**: How $\pi_{\text{old}}$ is computed (Decoupled vs Bypass) +2. **Loss Function**: PPO (with clipping) vs Pure IS (policy gradient only) +3. **IS/RS Aggregation Level**: Token, Sequence, or Geometric + +This section explains each component and their valid combinations. + +### 3.1 Operating Modes: Decoupled vs Bypass + +The operating mode determines how the proximal policy $\pi_{\text{old}}$ is computed. + +#### 3.1.1 Decoupled Mode (Three Policies) + +**Configuration:** `bypass_mode = false` + +**Policy setup:** +- $\pi_{\text{rollout}}$: Behavior policy (data collection) +- $\pi_{\text{old}}$: Proximal policy (computed via `actor.compute_log_prob()` at start of training epoch) +- $\pi_{\theta}$: Current policy (being updated) + +**IS ratio:** $\rho_t = \frac{\pi_{\text{old}}(a_t|s_t)}{\pi_{\text{rollout}}(a_t|s_t)}$ (corrects Drift 1: rollout→old) + +**PPO ratio:** $r_t(\theta) = \frac{\pi_{\theta}(a_t|s_t)}{\pi_{\text{old}}(a_t|s_t)}$ (corrects Drift 2: old→current) + +**Properties:** +- ✅ Achieves batch size invariance +- ✅ Separately corrects two distribution drifts +- ✅ Efficient stale data utilization +- ❌ Extra forward pass needed (`actor.compute_log_prob()`) + +#### 3.1.2 Bypass Mode (Two Policies) + +**Configuration:** `bypass_mode = true` + +**Policy setup:** +- $\pi_{\text{rollout}}$: Behavior policy (data collection) +- $\pi_{\text{old}} = \pi_{\text{rollout}}$: Proximal policy equals behavior policy +- $\pi_{\theta}$: Current policy (being updated) + +**Ratios:** +- **With PPO-clip loss** (`loss_type = "ppo_clip"`, default): PPO ratio $r_t(\theta) = \frac{\pi_{\theta}(a_t|s_t)}{\pi_{\text{rollout}}(a_t|s_t)}$ clips against rollout policy (IS handled by ratio) +- **With REINFORCE loss** (`loss_type = "reinforce"`): IS ratio $\rho_t = \frac{\pi_{\theta}(a_t|s_t)}{\pi_{\text{rollout}}(a_t|s_t)}$ computed on-the-fly in loss function + +**Properties:** +- ✅ Skips `actor.compute_log_prob()` call (faster) +- ✅ Handles off-policy correction via IS/RS (when using policy gradient with IS/RS) +- ✅ Uses two policies instead of three (π_rollout = π_old) +- ⚠️ Does not separate proximal policy from behavior policy (unlike decoupled mode) + +--- + +### 3.2 Loss Functions: PPO vs Policy Gradient + +#### 3.2.1 PPO Loss (with Clipping) + +**Configuration:** `loss_type = "ppo_clip"` (default in bypass mode) + +**Loss function:** + +$$ +L_{\text{PPO}}(\theta) = -\mathbb{E}_t \left[ w_t \cdot \min\left( r_t(\theta) A_t, \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon) A_t \right) \right] +$$ + +where: +- $w_t$: IS weight (depends on aggregation level, see Section 3.3). In decoupled mode, $w_t = \frac{\pi_{\text{old}}}{\pi_{\text{rollout}}}$ where $\pi_{\text{old}}$ is frozen, so $w_t$ is constant (no stopgrad needed). In bypass mode with PPO loss, no separate IS weights are typically computed. +- $r_t(\theta) = \frac{\pi_{\theta}(a_t|s_t)}{\pi_{\text{old}}(a_t|s_t)}$: PPO ratio +- $\epsilon$: Clip range (typically 0.2) + +**Properties:** +- Trust region control via clipping +- Limits policy update magnitude +- Standard in RL training + +#### 3.2.2 Policy Gradient Loss (with IS/RS Correction) + +**Configuration:** `loss_type = "reinforce"` (requires `bypass_mode = true`) + +**Loss function** (example with sequence-level IS): + +$$ +L_{\text{PG}}(\theta) = -\mathbb{E}_{(s,a) \sim \pi_{\text{rollout}}} \left[ \text{stopgrad}(w_{\text{seq}}(\theta)) \cdot \sum_{t \in T} \log \pi_{\theta}(a_t|s_t) \cdot A_t \right] +$$ + +where: +- $w_{\text{seq}}(\theta)$: Sample weight (IS or RS, see §3.3-3.4 for details) +- For IS: $w_{\text{seq}}(\theta) = \min\left( \prod_{t \in T} \frac{\pi_{\theta}(a_t|s_t)}{\pi_{\text{rollout}}(a_t|s_t)}, C_{\text{IS}} \right)$ +- For RS: $w_{\text{seq}}(\theta) \in \{0, 1\}$ (binary rejection mask) +- **stopgrad operator**: The weight $w_{\text{seq}}(\theta)$ is computed using $\pi_\theta$ but treated as a **constant coefficient** when computing $\nabla_\theta L$. This is essential for importance sampling correctness (see theoretical justification below). + +**Effective gradient:** + +$$ +\nabla_\theta L_{\text{PG}} = -\mathbb{E}_{(s,a) \sim \pi_{\text{rollout}}} \left[ \text{stopgrad}(w_{\text{seq}}(\theta)) \cdot \sum_{t \in T} \nabla_\theta \log \pi_{\theta}(a_t|s_t) \cdot A_t \right] +$$ + +**Theoretical Justification for stopgrad:** + +The stopgrad operator is **mathematically required** by importance sampling theory, not an implementation detail. Here's why: + +**The fundamental principle**: Importance sampling is a technique to **change the measure** (reweight samples from one distribution to estimate expectations under another), not to optimize the reweighting function itself. + +**Formal derivation**: + +1. **Original objective**: We want to optimize $J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}[\sum_t A_t]$. + +2. **Off-policy setting**: We only have samples from $\pi_{\text{rollout}}$, so we use importance sampling: + $$ + J(\theta) = \mathbb{E}_{\tau \sim \pi_{\text{rollout}}} \left[ \underbrace{\frac{P_{\pi_\theta}(\tau)}{P_{\pi_{\text{rollout}}}(\tau)}}_{w(\tau;\theta)} \sum_t A_t \right] + $$ + +3. **Computing the policy gradient**: The correct gradient uses the **policy gradient theorem BEFORE importance sampling**: + $$ + \begin{aligned} + \nabla_\theta J(\theta) &= \nabla_\theta \mathbb{E}_{\tau \sim \pi_\theta}\left[\sum_t A_t\right] \\ + &= \mathbb{E}_{\tau \sim \pi_\theta} \left[\sum_t A_t \nabla_\theta \log \pi_\theta(a_t|s_t) \right] \quad \text{(policy gradient theorem)} \\ + &= \mathbb{E}_{\tau \sim \pi_{\text{rollout}}} \left[ w(\tau;\theta) \sum_t A_t \nabla_\theta \log \pi_\theta(a_t|s_t) \right] \quad \text{(change of measure)} + \end{aligned} + $$ + + In the final line, $w(\tau;\theta)$ appears as a **multiplicative coefficient** from the change of measure, not as something we differentiate. + +4. **What goes wrong without stopgrad**: If we naively compute $\nabla_\theta \left[w(\theta) \log \pi_\theta \right]$ in the loss, we get: + $$ + \nabla_\theta \left[w(\theta) \log \pi_\theta \right] = \underbrace{\log \pi_\theta \cdot \nabla_\theta w(\theta)}_{\text{WRONG: bias term}} + \underbrace{w(\theta) \cdot \nabla_\theta \log \pi_\theta}_{\text{CORRECT: IS-weighted gradient}} + $$ + + The first term $\log \pi_\theta \cdot \nabla_\theta w(\theta)$ is an artifact of the computational trick (using loss times log-prob), not part of the true policy gradient. It biases the gradient estimator and optimizes a different objective than $J(\theta)$. + +5. **Implementation requirement**: In PyTorch, to compute only the second term, we must use: + ```python + loss = -advantages * log_prob * rollout_is_weights.detach() # stopgrad on weights + ``` + Without `.detach()`, autograd computes both terms, giving an incorrect gradient. + +**Intuition**: The IS weight $w(\theta)$ tells us "how much to trust this sample" for estimating the gradient under $\pi_\theta$. We update $\theta$ to maximize the reweighted objective, but we don't update $\theta$ to maximize the weight itself—that would be circular reasoning (optimizing the correction factor instead of the actual objective). + +**Properties:** +- **Algorithm**: Off-policy policy gradient with IS/RS correction +- **Loss types** (`loss_type` config option in bypass mode): + - `"ppo_clip"` (default): PPO clipped objective + - $L = -\mathbb{E}[\min(r \cdot A, \text{clip}(r) \cdot A)]$ where $r = \pi_\theta / \pi_{\text{rollout}}$ + - Note: IS weights NOT applied (PPO ratio already handles it; would be double-counting) + - `"reinforce"`: Pure policy gradient with explicit IS weights, no PPO clipping + - $L = -\mathbb{E}[w \cdot \log \pi_\theta(a|s) \cdot A]$ where $w = \pi_\theta / \pi_{\text{rollout}}$ +- **Always uses bypass mode**: Direct $\pi_\theta$ to $\pi_{\text{rollout}}$ comparison +- **Fast**: Single forward pass + +**Implementation:** `compute_policy_loss_bypass_mode()` and `compute_policy_loss_reinforce()` in [core_algos.py](../../verl/trainer/ppo/core_algos.py) + +--- + +### 3.3 IS/RS Aggregation Levels + +The aggregation level determines how per-token probability ratios are combined into IS weights and/or rejection masks. This choice is **orthogonal to the operating mode** - you can use any aggregation level in either decoupled or bypass mode. + +#### 3.3.1 Token-Level Aggregation + +**IS weights:** $w_t = \min(\rho_t, C_{\text{IS}})$ where $\rho_t = \frac{\pi_{\text{old}}(a_t|s_t)}{\pi_{\text{rollout}}(a_t|s_t)}$ (decoupled) or $\rho_t = \frac{\pi_{\theta}(a_t|s_t)}{\pi_{\text{rollout}}(a_t|s_t)}$ (bypass/pure IS) + +**Configuration:** +```python +rollout_is = "token" # IS weights +rollout_rs = "token_k1" # Optional: rejection sampling (ratio bounds) +``` + +**Properties:** +- Independent truncation per token +- Lower variance than sequence-level (product of ratios bounded individually) +- **Bias-variance tradeoff**: Token-level correction has $O(T^2 \Delta_{\max})$ bias where $T$ is sequence length and $\Delta_{\max}$ is maximum per-token policy divergence. This bias becomes significant when the rollout policy deviates substantially from the training policy. Sequence-level correction is unbiased but has higher variance. +- Typical threshold: 1.5 - 5.0 +- Optional batch normalization [§3.4](rollout_corr_math.md#34-batch-normalization): Normalizes over all token weights to ensure $\mathbb{E}[\tilde{w}_t] = 1$ (reduces variance) +- **When to use**: Token-level works well when rollout policy stays within the trust region of training policy. When mismatch is significant, the bias becomes intolerable and sequence-level correction is preferred. + +**Loss function (REINFORCE + Token IS):** + +$$ +L_{\text{REINFORCE+TIS}}(\theta) = -\mathbb{E}_t \left[ \text{stopgrad}(w_t) \cdot \log \pi_\theta(a_t|s_t) \cdot A_t \right] +$$ + +where $w_t = \min(\rho_t, C_{\text{IS}})$ are the truncated token-level IS weights. The stopgrad operator ensures that when computing $\nabla_\theta L$, the weights are treated as constants (see §3.2.2 for theoretical justification). This formulation can also be combined with PPO clipping by replacing the REINFORCE gradient with the clipped surrogate objective. + +**Implementation:** +- IS weights: `compute_rollout_correction_weights()` in [rollout_corr_helper.py](../../verl/trainer/ppo/rollout_corr_helper.py#L325-L402) +- Loss: `compute_policy_loss()` in [core_algos.py](../../verl/trainer/ppo/core_algos.py#L812-L884) + +#### 3.3.2 Sequence-Level Aggregation + +**IS weights:** $w_{\text{seq}} = \min\left( \prod_{t \in T} \rho_t, C_{\text{IS}} \right) = \min\left( \exp\left(\sum_{t \in T} \log \rho_t\right), C_{\text{IS}} \right)$ (broadcast to all tokens) + +**Configuration:** +```python +rollout_is = "sequence" # IS weights +rollout_rs = "seq_sum_k1" # Optional: rejection sampling +``` + +**Properties:** +- Multiplicative aggregation across sequence +- More sensitive to outliers than token-level +- Typical threshold: 2.0 - 10.0 +- Optional batch normalization [§3.4](rollout_corr_math.md#34-batch-normalization): Normalizes over sequence means (one weight per sequence) + +**Terminology Note:** +- **Seq-TIS (Sequence-Level Truncated IS)**: Clips the sequence ratio $\rho(\tau) \to \min(\rho(\tau), C)$. Maximizes information efficiency by extracting signal from all samples. Best for clean data with moderate mismatch. +- **Seq-MIS (Sequence-Level Masked IS)**: Rejects (masks) sequences with $\rho(\tau) > C$ instead of clipping. Acts as a hard trust region filter. Best for severe mismatch or when the distribution tail is "toxic" (contains garbage/adversarial samples rather than signal). + +**Loss function (REINFORCE + Sequence IS):** + +$$ +L_{\text{REINFORCE+SeqIS}}(\theta) = -\mathbb{E}_t \left[ \text{stopgrad}(w_{\text{seq}}) \cdot \log \pi_\theta(a_t|s_t) \cdot A_t \right] +$$ + +where $w_{\text{seq}}$ is broadcast to all tokens in the sequence. The stopgrad operator ensures correct IS gradient computation (see §3.2.2). This formulation can also be combined with PPO clipping. + +#### 3.3.3 Geometric Mean Aggregation (Geo-RS) + +**Geometric mean ratio:** $\rho_{\text{geo}} = \exp\left( \frac{1}{|T|} \sum_{t \in T} \log \rho_t \right) = \left(\prod_{t \in T} \rho_t\right)^{1/|T|}$ (broadcast to all tokens) + +**Configuration:** +```python +rollout_is = null # No IS weights, pure rejection +rollout_rs = "seq_mean_k1" # Geometric mean rejection sampling (ratio bounds) +``` + +**Properties:** +- Length-invariant (normalizes by sequence length) +- Ideal ratio = 1.0 (policies match) +- Typical bounds: `"0.999_1.001"` (~±0.1%) +- **Used for rejection sampling only, not IS weighting** + +**The Length Trap Problem:** + +Standard IS estimators have a systematic **length bias** that penalizes long sequences. The importance ratio $\rho(y)$ is multiplicative: + +$$ +\rho(y) = \prod_{t=1}^T \frac{\pi(y_t|y_{= 0 per token (equals 0 when ρ = 1) +- More stable than geometric ratio checks because each token term is non-negative +- Only upper threshold applies (no lower threshold since K3 >= 0) +- Typical threshold: 0.001 - 0.01 + +**Why K3 over geometric ratio?** +- Geometric ratio uses average log-ratio; small numerical bias can flip sign +- K3 = E[ρ - log ρ - 1] is non-negative per token, offering a smoother detector +- Both estimate the same quantity: KL(π_rollout || π_old) +- For small divergences, K3 ≈ 0.5 × Var(log_ratio) + +**Combined Estimator (K3-RS-Token-TIS):** + +For best results, combine K3 filter with token-level IS weights: + +$$ +\hat{g}_{\text{k3-rs-token-tis}}(y) = \underbrace{\mathbb{I}\left( K3_{\text{seq}} \le C_{\text{k3}} \right)}_{\text{K3 Filter}} \cdot \prod_t \min(\rho_t, C) \cdot f(y) +$$ + +This is implemented by combining `rollout_rs="seq_mean_k3"` with `rollout_is="token"`. + + +--- + +### 3.4 Batch Normalization + +An optional variance reduction technique that normalizes IS weights to have mean 1.0 within each batch. + +**Configuration:** +```python +rollout_is_batch_normalize = True # Default: False +``` + +**Normalization formula (aggregation-aware):** + +For **token-level IS** (§3.3.1): + +$$ +\tilde{w}_t = \frac{w_t}{\frac{1}{\sum_{i,t} m_{i,t}} \sum_{i,t} w_{i,t} \cdot m_{i,t}} +$$ + +where $w_{i,t}$ are truncated token IS weights, $m_{i,t}$ is the response mask, and normalization is over **all tokens**. + +For **sequence-level IS** (§3.3.2): + +$$ +\tilde{w}_i = \frac{w_i}{\frac{1}{B}\sum_{j=1}^B \bar{w}_j} +$$ + +where $\bar{w}_j = \frac{1}{T_j}\sum_{t=1}^{T_j} w_{j,t} \cdot m_{j,t}$ is the per-sequence mean (all tokens in a sequence have the same weight), and normalization is over **sequences**. + +**Properties:** +- Applied **after** truncation to preserve truncation semantics +- Ensures $\mathbb{E}[\tilde{w}] = 1$ within each batch +- **Aggregation-aware**: Token-level normalizes over tokens; sequence-level normalizes over sequences +- Uses `masked_mean` to respect padding tokens +- Reduces gradient magnitude variance by removing random batch-level scale fluctuations + +**Metrics:** +- `rollout_is_batch_norm_factor`: The normalization factor applied (batch mean before normalization) + +**Implementation:** [rollout_corr_helper.py](../../verl/trainer/ppo/rollout_corr_helper.py#L401-L421) + +--- + +### 3.5 Rejection Sampling (RS) + +Rejection sampling can be added to **any combination** of operating mode and aggregation level. It modifies the `response_mask` to exclude outlier tokens/sequences. + +**Configuration examples:** +```python +rollout_rs = "token_k1" # Token-level ratio bounds +rollout_rs_threshold = "0.6_1.6" + +rollout_rs = "seq_sum_k1" # Sequence sum of log ratios +rollout_rs_threshold = "0.5_2.0" + +rollout_rs = "seq_mean_k3" # Sequence mean of K3 divergence +rollout_rs_threshold = 0.01 +``` + +**Acceptance set:** +- **Token-level**: $\mathcal{A}_{\text{token}} = \{ t : C_{\text{RS-lower}} \leq \rho_t \leq C_{\text{RS-upper}} \}$ +- **Sequence-level**: $\mathcal{A}_{\text{seq}} = \{ \text{seq} : C_{\text{RS-lower}} \leq \prod_{t \in T} \rho_t \leq C_{\text{RS-upper}} \}$ +- **Geometric**: $\mathcal{A}_{\text{geo}} = \{ \text{seq} : C_{\text{RS-lower}} \leq \rho_{\text{geo}} \leq C_{\text{RS-upper}} \}$ + +**Properties:** +- Separate from IS weighting (can use RS without IS) +- Reduces effective sample size +- Filters extreme outliers + +**Implementation:** `compute_rollout_rejection_mask()` in [rollout_corr_helper.py](../../verl/trainer/ppo/rollout_corr_helper.py#L80-L188) + +--- + +### 3.6 Combination Matrix + +**Key insight:** Estimators (how IS/RS is computed) and operating modes (decoupled PPO vs bypass PG) are **orthogonal**. Any estimator can be combined with any operating mode. + +#### Estimator × Operating Mode + +| Estimator | Configuration | Compatible Modes | +|-----------|---------------|------------------| +| **Token-TIS** | `rollout_is="token"` | Decoupled PPO, Bypass PG | +| **Seq-TIS** | `rollout_is="sequence"` | Decoupled PPO, Bypass PG | +| **Seq-MIS** | `rollout_is="sequence"` + `rollout_rs="seq_sum_k1"` | Decoupled PPO, Bypass PG | +| **Geo-RS** | `rollout_rs="seq_mean_k1"` (geometric mean) | Decoupled PPO, Bypass PG | +| **Geo-RS-Token-TIS** | `rollout_is="token"` + `rollout_rs="seq_mean_k1"` | Decoupled PPO, Bypass PG | +| **K3-RS** | `rollout_rs="seq_mean_k3"` | Decoupled PPO, Bypass PG | +| **K3-RS-Token-TIS** | `rollout_is="token"` + `rollout_rs="seq_mean_k3"` | Decoupled PPO, Bypass PG | + +**Note:** In bypass mode, `loss_type` controls the loss function. Use "ppo_clip" (default) or "reinforce". + +#### Available Preset Methods + +| Preset Method | Estimator | Mode | Properties | +|---------------|-----------|------|------------| +| **Decoupled PPO Mode** (3 policies: π_rollout, π_old, π_θ) | +| `decoupled_token_is()` | Token-TIS | Decoupled PPO | Per-token IS weights | +| `decoupled_seq_is()` | Seq-TIS | Decoupled PPO | Sequence-level IS weights | +| `decoupled_seq_is_rs()` | Seq-MIS | Decoupled PPO | Sequence IS + sequence RS | +| `decoupled_geo_rs()` | Geo-RS | Decoupled PPO | Geometric RS | +| `decoupled_geo_rs_token_tis()` | Geo-RS-Token-TIS | Decoupled PPO | Geometric filter + token IS | +| **K3 KL Estimator** (more stable for small KL values) | +| `decoupled_k3_rs()` | K3-RS | Decoupled PPO | K3 rejection, no IS weights | +| `decoupled_k3_rs_token_tis()` | K3-RS-Token-TIS | Decoupled PPO | K3 filter + token clipped weight | +| **Bypass Mode (PPO-clip)** (ratio handles IS, RS masks outliers) | +| `bypass_ppo_clip()` | - | Bypass (PPO-clip) | PPO-clip only | +| `bypass_ppo_clip_geo_rs()` | Geo-RS | Bypass (PPO-clip) | PPO-clip + Geo-RS (ratio) | +| `bypass_ppo_clip_k3_rs()` | K3-RS | Bypass (PPO-clip) | PPO-clip + K3-RS | +| **Bypass Mode (REINFORCE)** (explicit IS weights, no PPO clipping) | +| `bypass_pg_is()` | Seq-TIS | Bypass (REINFORCE) | REINFORCE + Seq IS | +| `bypass_pg_geo_rs()` | Geo-RS | Bypass (REINFORCE) | REINFORCE + Geo-RS (ratio) | +| `bypass_pg_geo_rs_token_tis()` | Geo-RS-Token-TIS | Bypass (REINFORCE) | REINFORCE + Geo filter + token IS | +| **Other** | +| `disabled()` | - | - | Metrics only | + +**Note:** Bypass mode sets π_old = π_rollout and uses `loss_type` to select the loss function. + +#### Additional Supported Combinations (Manual Configuration) + +These combinations are **fully supported** but require manual configuration: + +**1. Token IS + Token RS** +```python +config = RolloutCorrectionConfig( + rollout_is="token", + rollout_is_threshold=2.0, + rollout_rs="token_k1", + rollout_rs_threshold="0.5_2.0", +) +``` +**Properties:** Token-level IS weights + token-level RS mask. + +**2. Pure Token RS** +```python +config = RolloutCorrectionConfig( + rollout_is=None, + rollout_rs="token_k1", + rollout_rs_threshold="0.5_2.0", +) +``` +**Properties:** Token-level RS mask only, no IS weights. + +**3. Pure Sequence RS** +```python +config = RolloutCorrectionConfig( + rollout_is=None, + rollout_rs="seq_sum_k1", + rollout_rs_threshold="0.5_2.0", +) +``` +**Properties:** Sequence-level RS mask only, no IS weights. + +**Key properties:** +- Any IS aggregation level (token/sequence) can be used in either decoupled or bypass mode +- Rejection sampling can be added to any combination +- Geometric aggregation is typically used for RS only (not IS weighting) +- Pure RS (`bypass_pg_rs`) uses bypass + geometric RS with `loss_type="reinforce"` for REINFORCE (no IS weights) +- All combinations in the table above are valid and supported by the implementation + +--- + +### 3.7 Common Implementation Mistake + +#### Incorrect LLM-RL Implementation (PPO Without Rollout Correction) + +**Theory:** Naive LLM-RL implementation that incorrectly applies PPO by **ignoring the actual rollout policy** and assuming $\pi_{\text{old}} = \pi_{\text{rollout}}$. + +**Note:** This incorrect implementation pattern was identified in [Liu, Li, et al. (2025)](https://richardli.xyz/rl-collapse) as a key cause of training instability in LLM-RL systems, motivating the development of this rollout correction framework. + +**Loss Function:** + +$$ +L_{\text{PPO}}(\theta) = -\mathbb{E}_t \left[ \min\left( r_t(\theta) A_t, \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon) A_t \right) \right] +$$ + +where $r_t(\theta) = \frac{\pi_{\theta}(a_t|s_t)}{\pi_{\text{old}}(a_t|s_t)}$ (ignores $\pi_{\text{rollout}}$). + +**Why it's wrong:** +- **Ignores $\pi_{\text{rollout}}$**: Uses $\pi_{\text{old}}$ as behavior policy instead of actual $\pi_{\text{rollout}}$ +- **Policy mismatch**: In LLM-RL, rollout typically uses different precision/backend/checkpoint than training, causing $\pi_{\text{rollout}} \neq \pi_{\text{old}}$ even with same model weights +- **Not PPO's fault**: PPO itself is correct; the issue is the incorrect assumption + +**Correct alternatives:** +1. **Decoupled mode**: Three policies with IS correction from $\pi_{\text{rollout}}$ to $\pi_{\text{old}}$ +2. **Bypass mode**: Two policies using $\pi_{\text{rollout}}$ as both behavior policy and proximal policy +3. **Bypass + Policy Gradient mode**: Two policies with IS/RS correction and no PPO clipping + +**Implementation:** `compute_policy_loss()` in [core_algos.py](../../verl/trainer/ppo/core_algos.py#L812-L884) + +--- + +## 4. Off-Policy Diagnostic Metrics + +These metrics quantify the severity of off-policy drift. + +**Note on notation:** Metrics use $\rho_t = \frac{\pi_{\text{old}}(a_t|s_t)}{\pi_{\text{rollout}}(a_t|s_t)}$. In bypass mode, $\pi_{\text{old}} = \pi_{\text{rollout}}$, so metrics measure rollout→current drift using $\rho_t = \frac{\pi_{\theta}}{\pi_{\text{rollout}}}$ instead. + +### 4.1 KL Divergence + +**Direct KL estimator:** + +$$ +\text{KL}(\pi_{\text{rollout}} \| \pi_{\text{old}}) = \mathbb{E}_{t \sim \pi_{\text{rollout}}} \left[ \log \pi_{\text{rollout}}(a_t|s_t) - \log \pi_{\text{old}}(a_t|s_t) \right] +$$ + +**K3 KL estimator** (alternative formulation): + +$$ +\text{KL}_{\text{K3}} = \mathbb{E}_{t \sim \pi_{\text{rollout}}} \left[ \rho_t - \log \rho_t - 1 \right] +$$ + +where $\rho_t = \frac{\pi_{\text{old}}(a_t|s_t)}{\pi_{\text{rollout}}(a_t|s_t)}$. + +### 4.2 Perplexity + +**Old policy perplexity:** + +$$ +\text{PPL}_{\text{old}} = \exp\left( -\frac{1}{|T|} \sum_{t \in T} \log \pi_{\text{old}}(a_t|s_t) \right) +$$ + +**Rollout policy perplexity:** + +$$ +\text{PPL}_{\text{rollout}} = \exp\left( -\frac{1}{|T|} \sum_{t \in T} \log \pi_{\text{rollout}}(a_t|s_t) \right) +$$ + +**PPL ratio** (inverse of geometric mean IS weight): + +$$ +\text{PPL}_{\text{ratio}} = \frac{\text{PPL}_{\text{old}}}{\text{PPL}_{\text{rollout}}} = \exp\left( -\frac{1}{|T|} \sum_{t \in T} \log \rho_t \right) = \left(\prod_{t \in T} \rho_t\right)^{-1/|T|} +$$ + +**Interpretation:** Values > 1 mean $\pi_{\text{old}}$ assigns lower probability than $\pi_{\text{rollout}}$ to the observed actions (distribution shift). + +### 4.3 Chi-squared Divergence + +Measures the second moment of the IS weight distribution. + +**Token-level:** + +$$ +\chi^2_{\text{token}} = \mathbb{E}_{t \sim \pi_{\text{rollout}}} \left[ \rho_t^2 \right] - 1 +$$ + +**Sequence-level:** + +$$ +\chi^2_{\text{seq}} = \mathbb{E}_{\text{seq} \sim \pi_{\text{rollout}}} \left[ \left(\prod_{t \in T} \rho_t\right)^2 \right] - 1 +$$ + +**Interpretation:** +- $\chi^2 = 0$: Policies are identical +- $\chi^2 > 0$: Higher values indicate more severe off-policy distribution shift + +**Implementation:** `compute_offpolicy_metrics()` in [rollout_corr_helper.py](../../verl/trainer/ppo/rollout_corr_helper.py#L670-L776) + +--- + +## 5. Summary and Decision Guide + +### 5.1 Method Summary Table + +| Method | Theory | Policies | PPO Clip | IS Correction | Correctness | Speed | +|--------|--------|----------|----------|---------------|-------------|-------| +| **Bypass Mode** (π_old = π_rollout, `loss_type` selects algorithm) | +| `loss_type="ppo_clip"` (default) | PPO (ratio = π_θ/π_rollout) | 2 (rollout, θ) | ✅ | RS mask only (ratio handles IS) | ✅ Correct | **Fast** | +| `loss_type="reinforce"` | Off-policy REINFORCE | 2 (rollout, θ) | ❌ | ✅ (explicit IS weights) | ✅ Correct | **Fast** | +| **Bypass Mode Presets (PPO-clip)** | +| `bypass_ppo_clip` | PPO only | 2 (rollout, θ) | ✅ | - | ✅ Correct | **Fast** | +| `bypass_ppo_clip_geo_rs` | PPO + Geo-RS | 2 (rollout, θ) | ✅ | Geo-RS mask (ratio) | ✅ Correct | **Fast** | +| **Bypass Mode Presets (REINFORCE)** | +| `bypass_pg_is` | REINFORCE + Seq-TIS | 2 (rollout, θ) | ❌ | ✅ Seq-TIS | ✅ Correct | **Fast** | +| `bypass_pg_geo_rs` | REINFORCE + Geo-RS | 2 (rollout, θ) | ❌ | Geo-RS only (ratio) | ✅ Correct | **Fast** | +| `bypass_pg_geo_rs_token_tis` | REINFORCE + Geo RS + Token IS | 2 (rollout, θ) | ❌ | ✅ Geo-RS-Token-TIS | ✅ Correct | **Fast** | +| **Decoupled PPO Mode** (IS weights = π_old / π_rollout) | +| `decoupled_token_is` | Decoupled PPO | 3 (rollout, old, θ) | ✅ | ✅ Token-TIS | ✅ Correct | Standard | +| `decoupled_seq_is` | Decoupled PPO | 3 (rollout, old, θ) | ✅ | ✅ Seq-TIS | ✅ Correct | Standard | +| `decoupled_seq_is_rs` | Decoupled PPO + RS | 3 (rollout, old, θ) | ✅ | ✅ Seq-MIS | ✅ Correct | Standard | +| `decoupled_geo_rs` | Decoupled PPO + Geo-RS | 3 (rollout, old, θ) | ✅ | Geo-RS only (ratio) | ✅ Correct | Standard | +| `decoupled_geo_rs_token_tis` | Decoupled PPO + Geo RS + Token IS | 3 (rollout, old, θ) | ✅ | ✅ Geo-RS-Token-TIS | ✅ Correct | Standard | +| **Incorrect (for reference)** | +| Naive LLM-RL | Incorrect PPO usage | 2 (old, θ) | ✅ | ❌ | ⚠️ Incorrect | Standard | + +**Notes:** +- **Bypass mode** sets π_old = π_rollout and uses `loss_type` to select the loss function: + - `"ppo_clip"` (default): PPO clipped ratio (IS handled by ratio = π_θ/π_rollout, no explicit IS weights to avoid double-counting) + - `"reinforce"`: Explicit IS weights applied as $w \cdot \log \pi \cdot A$ +- Both loss types benefit from rejection sampling (RS) which masks out-of-distribution samples + +### 5.2 Estimator Hierarchy + +These estimators define **how IS weights and rejection masks are computed**. They are orthogonal to the operating mode (decoupled PPO vs bypass policy gradient) and can be combined with either. + +| Estimator | Configuration | Mechanism | Best For | +|-----------|---------------|-----------|----------| +| **Token-TIS** | `rollout_is="token"` | Clips per-token ratios | Lower variance IS with acceptable bias | +| **Seq-TIS** | `rollout_is="sequence"` | Clips sequence ratio $\rho(\tau) \to \min(\rho(\tau), C)$ | Clean data with moderate mismatch; unbiased | +| **Seq-MIS** | `rollout_is="sequence"` + `rollout_rs="seq_sum_k1"` | Rejects sequences with $\rho(\tau) > C$ | Severe mismatch; filters "toxic tail" (garbage data) | +| **Geo-RS** | `rollout_rs="seq_mean_k1"` | Rejects on geometric mean ratio exp(E[log(r)]) | Length-invariant trust region | +| **Geo-RS-Token-TIS** | `rollout_is="token"` + `rollout_rs="seq_mean_k1"` | Geometric filter + token IS weights | Ratio-based length normalization + lower variance IS | +| **K3-RS** | `rollout_rs="seq_mean_k3"` | Rejects on K3 KL divergence | Small KL values; smooth detector | +| **K3-RS-Token-TIS** | `rollout_is="token"` + `rollout_rs="seq_mean_k3"` | K3 filter + token IS weights | Small KL + lower variance IS | + +**Note:** Each estimator can be used with either: +- **Decoupled PPO** (`bypass_mode=false`): Three policies with PPO clipping +- **Bypass Mode** (`bypass_mode=true`): Two policies with configurable loss type + - `loss_type="ppo_clip"` (default): PPO clipped objective (IS via ratio, RS mask applied) + - `loss_type="reinforce"`: REINFORCE with explicit IS weights + +### 5.3 Method Characteristics by Scenario + +**Choosing estimator by off-policy severity:** +- **Negligible** (same checkpoint, minor differences): No IS correction needed; use bypass mode for efficiency +- **Moderate** (async workers, slight staleness): Token-TIS provides per-token IS correction with lower variance +- **Severe** (replay buffers, old data): Seq-TIS or Seq-MIS provides sequence-level IS correction; use Seq-MIS when high-weight samples are likely garbage + +**Choosing estimator by sequence length:** +- **Short sequences** (standard chat): Seq-TIS is optimal +- **Long sequences** (CoT, agents): K1-RS or K1-RS-Token-TIS to avoid Length Trap + +**Choosing operating mode:** +- **Batch size invariance needed**: Use decoupled mode (`bypass_mode=false`) +- **Computational efficiency needed**: Use bypass mode (`bypass_mode=true`) to skip `old_log_prob` computation +- **No PPO clipping**: Use bypass mode with `loss_type="reinforce"` + +### 5.4 Decoupled Mode vs Bypass Mode + +**Decoupled mode** (computes `old_log_prob` separately): +- Implements full decoupled PPO with three policies (mathematically correct) +- Separately measures and corrects Drift 1 (rollout→old) and Drift 2 (old→current) +- Achieves batch size invariance and efficient stale data utilization +- Enables accurate off-policy metrics monitoring + +**Bypass mode** (sets $\pi_{\text{old}} = \pi_{\text{rollout}}$): +- Uses $\pi_{\text{rollout}}$ as both behavior policy and proximal policy (mathematically correct) +- Computational efficiency: Skips separate `old_log_prob` computation +- Does not achieve batch size invariance (proximal policy depends on data collection) + +--- + +## 6. Implementation References + +- **[Rollout Correction Usage Guide](rollout_corr.md)** - Practical configuration and troubleshooting +- **Config:** [verl/trainer/config/algorithm.py](../../verl/trainer/config/algorithm.py) +- **IS/RS Helper:** [verl/trainer/ppo/rollout_corr_helper.py](../../verl/trainer/ppo/rollout_corr_helper.py) +- **PPO Loss:** [verl/trainer/ppo/core_algos.py](../../verl/trainer/ppo/core_algos.py) +- **Tests:** [tests/trainer/ppo/test_rollout_corr.py](../../tests/trainer/ppo/test_rollout_corr.py) + +--- + +## References + +- **Williams, R. J. (1992).** "Simple statistical gradient-following algorithms for connectionist reinforcement learning." *Machine Learning*, 8(3-4), 229-256. https://doi.org/10.1007/BF00992696 +- **Schulman, J., Wolski, F., Dhariwal, P., Radford, A., & Klimov, O. (2017).** "Proximal policy optimization algorithms." *arXiv preprint arXiv:1707.06347.* https://arxiv.org/abs/1707.06347 +- **Hilton, J., Cobbe, K., & Schulman, J. (2021).** "Batch size-invariance for policy optimization." *arXiv preprint arXiv:2110.00641.* https://arxiv.org/abs/2110.00641 + - Introduced decoupled PPO: separating proximal policy (for controlling policy update size) from behavior policy (for off-policy correction) to achieve batch size invariance diff --git a/verl/docs/algo/spin.md b/verl/docs/algo/spin.md new file mode 100644 index 0000000000000000000000000000000000000000..682ce43954f5cefef358d45b4ff097941bd077c3 --- /dev/null +++ b/verl/docs/algo/spin.md @@ -0,0 +1,179 @@ +# Recipe: Self-Play Fine-Tuning (SPIN) + +Last updated: 05/31/2025. + +`verl` provides a recipe inspired by the paper **"Self-Play Fine-Tuning Converts Weak Language Models to Strong Language Models"** (SPIN). SPIN is a language model finetuning algorithm that enables iterative self-improvement through a self-play mechanism inspired by game theory. + +**Core Idea:** Models learn by playing against themselves, reducing reliance on external preference datasets or stronger teacher models: + +1. **Synthetic Data Generation:** The current model generates responses, creating its own training data from previous iterations. +2. **Two-Player Game Setup:** A game involving two players acted by a single LLM. +3. **Iterative Training:** The model progressively improves by refining its policy, with each iteration's model becoming the opponent for the next iteration. + +Paper Authors: [Zixiang Chen](https://github.com/uclaml/SPIN)\*, [Yihe Deng](https://github.com/uclaml/SPIN)\*, [Huizhuo Yuan](https://scholar.google.com/citations?user=8foZzX4AAAAJ)\*, [Kaixuan Ji](https://scholar.google.com/citations?user=FOoKDukAAAAJ), [Quanquan Gu](https://web.cs.ucla.edu/~qgu/) + +[[Webpage](https://uclaml.github.io/SPIN/)] [[Huggingface](https://huggingface.co/papers/2401.01335)] [[Paper](https://arxiv.org/abs/2401.01335)] [[Original Implementation](https://github.com/uclaml/SPIN)] + +verl Implementation Authors: [Chendong Wang](https://cdwang96.github.io/), [Chenyang Zhao](https://github.com/zhaochenyang20) + +--- + +## Key Function (compute_online_dpo_loss) and Related works +SPIN (Chen et al., 2024) proposes an iterative self-play mechanism to fine-tune language models. In each iteration, SPIN's training objective, when using a logistic loss function, is equivalent to Direct Preference Optimization (DPO) loss (Rafailov et al., 2023). + +This `verl` recipe realizes SPIN's core concept by using DPO loss iteratively (Xu et al., 2023; Xiong et al., 2023; Snorkel AI, 2024). This means that in each iteration, we fine-tune the LLM using DPO loss for preference optimization. Notably, Xu et al. (2023) explored iterative preference optimization with pairwise cringe loss, while Xiong et al. (2023) discussed how to bridge theory and practice for RLHF under KL constraints using iterative training. The concept of iterative preference learning was also explored in online DPO (Guo et al., 2024), which focuses on direct alignment from online AI feedback. In online DPO, preference data is dynamically updated during training, allowing the model to learn from its own generated data. + +Specifically, we developed the **`compute_online_dpo_loss`** function and built this SPIN recipe on top of it. By incorporating online preference generation, this approach enables continuously refining language models without relying on fixed external preference datasets. + +**Reference Papers:** +* [Self-Play Fine-Tuning Converts Weak Language Models to Strong Language Models](https://arxiv.org/abs/2401.01335) (Chen et al., 2024) +* [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://arxiv.org/abs/2305.18290) (Rafailov et al., 2023) +* [Somethings are more cringe than others: Preference optimization with the pairwise cringe loss](https://arxiv.org/abs/2312.16682) (Xu et al., 2023) +* [Iterative preference learning from human feedback: Bridging theory and practice for rlhf under kl-constraint](https://arxiv.org/abs/2312.11456) (Xiong et al., 2023) +* [Snorkel-Mistral-PairRM-DPO](https://huggingface.co/snorkelai/Snorkel-Mistral-PairRM-DPO) (Snorkel AI, 2024) +* [Direct language model alignment from online ai feedback](https://arxiv.org/abs/2402.04792) (Guo et al., 2024) + + +## Our Online DPO Implementation + +Our `compute_online_dpo_loss` function adapts `verl`'s existing PPO infrastructure (based on `verl` v0.3.0.post1) for this iterative online DPO. Key aspects of our implementation include: + +* **No Critic:** Unlike PPO, we omit the value function critic. +* **Dynamic Reference Model:** An explicit reference policy (`ref_policy_wg`) is used for DPO loss. This reference model's weights can be periodically updated from the actor (`ref_update_freq`), providing a dynamic baseline. +* **Online Preference Generation:** The `compute_onlineDPO_pref` function (in `core_algos.py`) dynamically creates chosen/rejected pairs based on a reward source (e.g., rule-based ranking for math problems). +* **DPO Loss Integration:** We replace PPO's policy loss with our `compute_online_dpo_loss` (in `core_algos.py`) within the actor update (`dp_actor.py`), directly optimizing the policy using the generated preferences. +* **Iterative Training Orchestration:** The `SpinTrainer` (in `spin_trainer.py`) manages the entire self-play loop: generation, preference labeling, optional reference model updates, and policy updates, enabling continuous self-improvement aligned with SPIN's principles. + +--- +## Algorithm + +This recipe implements an Online algorithm adapted to the `verl` Reinforcement Learning framework, which provides an alternative to PPO for fine-tuning language models. + +**Online Loop:** Instead of maximizing a scalar reward signal in PPO, this approach directly optimizes the policy model to align with preference data generated *online* during training: + +1. **Generation:** The current model generates multiple responses for each prompt in a batch. +2. **Preference Labeling:** A function evaluates these generated responses to determine which one is preferred (chosen) and which is dispreferred (rejected). This can be done using a reward function or implicit ranking based on specific rules. (In this recipe, we use rule-based ranking on the math problem). +3. **Update:** This preference tuple (`prompt`, `chosen_response`, `rejected_response`) is used to update the actor model using `compute_online_dpo_loss`, comparing against a reference model. + +**Connection with SPIN:** +Instead of only using a fixed target data distribution, the online generation loop in step 2 will dynamically change the target data distribution by using a certain Preference Labeling method (rule-based ranking on the math problem by selecting the better one in this recipe). This explores the direction mentioned in SPIN's paper Section 7 about "dynamically changing target data distribution" to potentially elevate LLM performance beyond the fixed human-annotated data ceiling. + +--- + +## Reproduce the Experiment (Example Setup) + +The following steps outline how to set up the environment and run the SPIN recipe, based on the provided test log using GSM8K and Qwen2.5-3B-Instruct. + +1. **Setup Environment (Example using Docker):** + ```bash + # Start a container with GPU access and shared memory + docker run -it --name spin_test --gpus all \ + --shm-size=32g \ + --ipc=host \ + -v /path/to/host/.cache:/root/.cache \ + -e HF_TOKEN= \ + lmsysorg/sglang:latest \ + /bin/bash + + # Inside the container or on your host machine: + # Ensure /tmp is writable + mkdir -p /tmp + chmod 1777 /tmp + + # Install Python 3.10 (if not present) and venv + sudo apt update + sudo apt install -y python3.10 python3.10-venv tmux + python3 -m ensurepip --upgrade + + # Create and activate a virtual environment + python3 -m venv ~/.python/spin_env + source ~/.python/spin_env/bin/activate + + # Install uv (fast package installer) + python3 -m pip install uv + ``` + +2. **Install verl and Dependencies:** + ```bash + # Clone the verl repository and checkout the spin branch + cd ~ + git clone git@github.com:verl-project/verl.git && cd verl + + # Install flash-attn (handle potential build issues) + python3 -m uv pip install wheel packaging + python3 -m uv pip install flash-attn --no-build-isolation --no-deps + + # Install verl with sglang extras + python3 -m uv pip install -e ".[sglang]" + ``` + *Note: If `flash-attn` installation fails, try the manual steps again or consult its documentation.* + +3. **Login & Download Data/Model:** + ```bash + # Login to Weights & Biases (optional, for logging) + export WANDB_API_KEY= + # wandb login + + # Download the GSM8K dataset + python3 examples/data_preprocess/gsm8k.py --local_save_dir ~/data/gsm8k # Adjusted path + + # Download the base model (Example: Qwen2.5-3B-Instruct) + hf download Qwen/Qwen2.5-3B-Instruct --local-dir $HOME/models/Qwen2.5-3B-Instruct + ``` + +4. **Configure:** + * Modify the configuration file (e.g., `config/spin_trainer.yaml` or the one specified in the run script) with correct paths to your downloaded model, data, desired hyperparameters (`dpo_beta`, learning rate, etc.), and distributed training settings (nodes, GPUs per node). + * Pay attention to `actor_rollout_ref.model`, `data` paths, `reward_model` config (if using one), and `trainer.ref_update_freq`. + +5. **Run Training:** + ```bash + # Set CUDA visible devices (adjust based on your hardware and config) + export CUDA_VISIBLE_DEVICES=0,1,2,3 + + # Launch the training script (e.g., test.sh or a custom script) + # Ensure test.sh points to the correct config and main script + bash recipe/spin/run_spin.sh + ``` + +--- + +## Configuration + +* The primary configuration is typically managed through a YAML file specified in the launch script (e.g., `config/spin_trainer.yaml`). +* Key configuration sections: + * `data`: Paths to training/validation prompt files, batch sizes, sequence lengths. + * `actor_rollout_ref`: Paths to the base model (used for actor and initial reference), FSDP settings, optimization parameters (learning rate, scheduler). + * `reward_model`: Configuration for the reward model used for online preference labeling (path, batch size, etc.). Can be omitted if using a simpler reward function. + * `algorithm`: DPO-specific hyperparameters like `dpo_beta`, `dpo_loss_type`. + * `trainer`: Distributed training settings (nodes, GPUs per node), logging (WandB), checkpointing frequency, and `ref_update_freq` (set > 0 to enable periodic reference model updates from the actor). + +--- + +## Key Files + +* `main_spin.py`: Main entry point using Hydra to load the config and launch the `SpinTrainer`. +* `spin_trainer.py`: Defines the `SpinTrainer` class, orchestrating the Online DPO training loop. +* `fsdp_workers.py`: Implements Ray workers (Actor, Reference) potentially using FSDP. +* `dp_actor.py`: Contains the actor class, including the DPO policy update logic. +* `core_algos.py`: Includes helper functions for `compute_online_dpo_loss` and `compute_onlineDPO_pref`. +* `config/spin_trainer.yaml` (or similar): Main Hydra configuration file for the recipe. +* `run_spin.sh` (or similar): Example bash script for launching a training run. +* `README.md`: This file. + +--- + +## Acknowledgement + +We sincerely thank the contribution and guidance from the `verl` community and advisors, including (adapted from SPPO): + +* [Zixiang Chen](https://sites.google.com/view/zxchen) +* [Yuhao Yang](https://github.com/yhyang201) +* [Yifan Zhang](https://github.com/yifanzhang-pro) +* [Yongan Xiang](https://github.com/BearBiscuit05) +* [Junrong Lin](https://github.com/ocss884) +* [Yuxuan Tong](https://github.com/tongyx361) +* [Guangming Shen](https://github.com/PeterSH6) +* [Biao He](https://www.linkedin.com/in/biao-he/) +* [Qingquan Song](https://qingquansong.github.io/) +* [Chenyang Zhao](https://zhaochenyang20.github.io/Chayenne/) +* [Quanquan Gu](https://web.cs.ucla.edu/~qgu/) diff --git a/verl/docs/algo/sppo.md b/verl/docs/algo/sppo.md new file mode 100644 index 0000000000000000000000000000000000000000..77dbbacdd15b01a13cf8614738fca91b3174234c --- /dev/null +++ b/verl/docs/algo/sppo.md @@ -0,0 +1,52 @@ +# Recipe: Self-Play Preference Optimization (SPPO) + +Last updated: 05/28/2025. + +verl provides a community recipe implementation for the paper [Self-Play Preference Optimization for Language Model Alignment](https://arxiv.org/abs/2405.00675). SPPO can significantly enhance the performance of an LLM without strong external signals such as responses or preferences from GPT-4. It can outperform the model trained with iterative direct preference optimization (DPO), among other methods. SPPO is theoretically grounded, ensuring that the LLM can converge to the von Neumann winner (i.e., Nash equilibrium) under general, potentially intransitive preference, and empirically validated through extensive evaluations on multiple datasets. + +Paper Authors: [Yue Wu](https://yuewu.us/)\*, [Zhiqing Sun](https://www.cs.cmu.edu/~zhiqings/)\*, [Huizhuo Yuan](https://scholar.google.com/citations?user=8foZzX4AAAAJ)\*, [Kaixuan Ji](https://scholar.google.com/citations?user=FOoKDukAAAAJ), [Yiming Yang](https://www.cs.cmu.edu/~yiming/), [Quanquan Gu](https://web.cs.ucla.edu/~qgu/) + +verl Implementation Authors: [Yuhao Yang](https://github.com/yhyang201), [Chenyang Zhao](https://github.com/zhaochenyang20) + +[[Webpage](https://uclaml.github.io/SPPO/)] [[Huggingface](https://huggingface.co/papers/2405.00675)] [[Paper](https://arxiv.org/abs/2405.00675)][[Original Implementation](https://github.com/uclaml/SPPO)] + +## Reproduce the Experiment + +We evaluate the performance of SPPO on the MATH dataset. Starting from an initial score of 46.6 with Qwen2.5-7B-Instruct, we achieve a score of 65.6 after 20 epochs of training, placing our model approximately in the top 20 on the [MATH leaderboard](https://paperswithcode.com/sota/math-word-problem-solving-on-math). It's important to note that verl's internal evaluation metrics may not perfectly align with the official evaluation methodology for Qwen2.5-7B-Instruct. Therefore, for consistency and fair comparison, we report only the results based on verl's evaluation framework. + +``` +git clone git@github.com:verl-project/verl.git +cd verl +python3 -m uv pip install -e ".[sglang]" + +export WANDB_API_KEY= + +python3 examples/data_preprocess/math_dataset.py --local_dir ~/data/math +hf download Qwen/Qwen2.5-7B-Instruct --local-dir $HOME/models/Qwen2.5-7B-Instruct + +export CUDA_VISIBLE_DEVICES=0,1,2,3 +bash recipe/sppo/run_qwen2.5-7b_rm.sh +``` + +Note that the installation would occasionally fail to install flash-attn. If this happens, you can install it manually by running: + +```bash +python3 -m uv pip install wheel +python3 -m uv pip install packaging +python3 -m uv pip install flash-attn --no-build-isolation --no-deps +``` + +## Acknowledgement + +We sincerely thank the contribution and guidance from: + +- [Yue Wu](https://yuewu.us/) +- [Chendong Wang](https://cdwang96.github.io/) +- [Yifan Zhang](https://github.com/yifanzhang-pro) +- [Yongan Xiang](https://github.com/BearBiscuit05) +- [Junrong Lin](https://github.com/ocss884) +- [Yuxuan Tong](https://github.com/tongyx361) +- [Guangming Shen](https://github.com/PeterSH6) +- [Biao He](https://www.linkedin.com/in/biao-he/) +- [Qingquan Song](https://qingquansong.github.io/) +- [Quanquan Gu](https://web.cs.ucla.edu/~qgu/) diff --git a/verl/docs/amd_tutorial/amd_build_dockerfile_page.rst b/verl/docs/amd_tutorial/amd_build_dockerfile_page.rst new file mode 100644 index 0000000000000000000000000000000000000000..096dad26e604e00e263b6a3a6364f38246accf49 --- /dev/null +++ b/verl/docs/amd_tutorial/amd_build_dockerfile_page.rst @@ -0,0 +1,796 @@ +Getting started with AMD (ROCM Kernel) +===================================================== + +Last updated: 07/06/2025. + +Author: `Yusheng Su `_ + +Setup +----- + +If you run on AMD GPUs (MI300) with ROCM platform, you cannot use the previous quickstart to run verl. You should follow the following steps to build a docker and set ``RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES`` or ``RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES`` when starting ray in verl's RLHF training. + + +docker/Dockerfile.rocm +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. code-block:: bash + + FROM "rlsys/rocm-6.3.4-patch:rocm6.3.4-numa-patch_ubuntu-22.04" + + SHELL ["/bin/bash", "-ceuxo", "pipefail"] + + ENV MAX_JOBS=512 + + ENV PATH="/usr/local/python3.12/bin:$PATH" + RUN ln -sf /usr/bin/python3.12 /usr/bin/python && \ + ln -sf /usr/bin/pip3.12 /usr/bin/pip + + ############################################ + RUN apt-get update + RUN apt-get install -y pkg-config liblzma-dev + ############################################ + + ########################################### + ##########Install TransformerEngine######## + ########################################### + WORKDIR /workspace/ + # transformer-engine install + # https://github.com/ROCm/TransformerEngine + RUN rm -rf TransformerEngine + RUN git clone --recursive https://github.com/ROCm/TransformerEngine.git + WORKDIR /workspace/TransformerEngine + git checkout 236178e5 + # git checkout bb061ade + # git checkout 864405c + ENV NVTE_FRAMEWORK=pytorch + ENV NVTE_ROCM_ARCH=gfx942 + ENV NVTE_USE_HIPBLASLT=1 + ENV NVTE_USE_ROCM=1 + # export CMAKE_PREFIX_PATH="/opt/rocm:/opt/rocm/hip:/usr/local:/usr:${CMAKE_PREFIX_PATH:-}" + ENV CMAKE_PREFIX_PATH="/opt/rocm:/opt/rocm/hip:/usr/local:/usr" + RUN MAX_JOBS=$(MAX_JOBS) pip install . -vvv + WORKDIR /workspace/ + ########################################### + ########################################### + ########################################### + + + + + + #################################################################################### + ################Install vllm - sglang require vllm 0.6.7 dependency################# + #################################################################################### + #### Require vllm 0.6.7 - checkout 113274a0 + WORKDIR /workspace/ + RUN rm -rf vllm + RUN pip uninstall -y vllm + # Refer to here (down-grade vllm to 0.6.3): https://docs.vllm.ai/en/v0.6.3/getting_started/amd-installation.html + RUN git clone https://github.com/ROCm/vllm.git + # git clone https://github.com/vllm-project/vllm.git + WORKDIR /workspace/vllm + RUN git checkout 113274a0 + ENV PYTORCH_ROCM_ARCH="gfx90a;gfx942" + #ENV MAX_JOBS=512 + ENV MAX_JOBS=${MAX_JOBS} + RUN pip install "boto3>=1.26.0" + RUN pip install setuptools_scm + # will add src into py. You can delete the repo + RUN python3 setup.py install + WORKDIR /workspace/ + #################################################################################### + #################################################################################### + #################################################################################### + + + + ########################################### + ############For hack docker################ + ########################################### + RUN pip install setuptools==75.8.0 + ########################################### + ########################################### + ########################################### + + + + ########################################### + ############build sgalng################### + ########################################### + # Set environment variables + ENV BASE_DIR=/sgl-workspace + ENV BUILD_TYPE=all + ENV SGL_REPO=https://github.com/sgl-project/sglang + ENV SGL_BRANCH=v0.4.6.post5 + ENV TRITON_REPO=https://github.com/ROCm/triton.git + ENV TRITON_COMMIT=improve_fa_decode_3.0.0 + ENV AITER_REPO=https://github.com/ROCm/aiter.git + ENV AITER_COMMIT=v0.1.2 + # v0.1.2 version - commit id: 9d11f47 + # ENV AITER_COMMIT=9d11f47 + ENV HIP_FORCE_DEV_KERNARG=1 + ENV HSA_NO_SCRATCH_RECLAIM=1 + ENV SGLANG_SET_CPU_AFFINITY=1 + ENV SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1 + ENV NCCL_MIN_NCHANNELS=112 + ENV MOE_PADDING=1 + ENV VLLM_FP8_PADDING=1 + ENV VLLM_FP8_ACT_PADDING=1 + ENV VLLM_FP8_WEIGHT_PADDING=1 + ENV VLLM_FP8_REDUCE_CONV=1 + ENV TORCHINDUCTOR_MAX_AUTOTUNE=1 + ENV TORCHINDUCTOR_MAX_AUTOTUNE_POINTWISE=1 + ENV HIPCC_COMPILE_FLAGS_APPEND="--offload-arch=gfx942" + ENV AMDGPU_TARGETS=gfx942 + ENV ROCM_ARCH=gfx942 + ENV PYTORCH_ROCM_ARCH="gfx90a;gfx942" + # Switch to working directory + WORKDIR /sgl-workspace + # Clean and create directory + RUN rm -rf /sgl-workspace && mkdir -p /sgl-workspace + + # Clone and build sglang + RUN git clone ${SGL_REPO} \ + && cd sglang \ + && git checkout ${SGL_BRANCH} || echo "Using default branch" \ + && cd sgl-kernel \ + && rm -f pyproject.toml \ + && mv pyproject_rocm.toml pyproject.toml \ + && python setup_rocm.py install \ + && cd .. \ + && if [ "$BUILD_TYPE" = "srt" ]; then \ + python -m pip --no-cache-dir install -e "python[srt_hip]"; \ + else \ + python -m pip --no-cache-dir install -e "python[all_hip]"; \ + fi \ + && cd /sgl-workspace \ + && cp -r /sgl-workspace/sglang /sglang \ + && python -m pip cache purge + + # Install common Python packages + RUN pip install IPython orjson python-multipart torchao pybind11 + # Rebuild Triton + RUN pip uninstall -y triton || true \ + && git clone ${TRITON_REPO} \ + && cd triton \ + && git checkout ${TRITON_COMMIT} \ + && cd python \ + && python3 setup.py install \ + && cd /sgl-workspace + # ENV HIPCC_COMPILE_FLAGS_APPEND="--offload-arch=gfx942 --amdgpu-lower-module-lds-strategy=1" + # ENV HIPCC_COMPILE_FLAGS_APPEND="--offload-arch=gfx942" + + # Build aiter + #version: Commit 9d11f47 + # && git checkout ${AITER_COMMIT} \ + RUN pip uninstall -y aiter || true + RUN git clone ${AITER_REPO} \ + && cd aiter \ + && git checkout ${AITER_COMMIT} \ + && git submodule sync \ + && git submodule update --init --recursive \ + && PREBUILD_KERNELS=1 GPU_ARCHS=gfx942 python3 setup.py install \ + && cd /sgl-workspace + + # Copy MI300X config + RUN find /sgl-workspace/sglang/python/sglang/srt/layers/quantization/configs/ \ + /sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs/ \ + -type f -name '*MI300X*' | \ + xargs -I {} sh -c 'vf_config=$(echo "$1" | sed "s/MI300X/MI300X_VF/"); cp "$1" "$vf_config"' -- {} + + # Environment setup complete. + RUN echo "Environment setup complete." + + WORKDIR /workspace/ + ########################################### + ########################################### + ########################################### + + + + + + + ########################################### + ###############vllm v0.8.5################# + ########################################### + WORKDIR /workspace/ + + ENV VLLM_TARGET_DEVICE=rocm + ENV ROCM_PATH=/opt/rocm + ENV SETUPTOOLS_SCM_PRETEND_VERSION=0.8.5.dev + # Find the repo path in: DockerFile/Dockerfile.rocm_yang + # RUN git clone https://github.com/RLFoundation/vllm-patch.git + RUN pip uninstall -y vllm || true + RUN rm -rf vllm-patch + RUN git clone https://github.com/RLFoundation/vllm-patch.git \ + && cd vllm-patch \ + && git checkout v0.8.5-sleep-numa \ + && rm -rf build/ dist/ *.egg-info \ + && ln -sf /opt/rocm/lib/libamdhip64.so /usr/lib/libamdhip64.so \ + && SETUPTOOLS_SCM_PRETEND_VERSION=0.8.5.dev PYTORCH_ROCM_ARCH="gfx90a;gfx942" MAX_JOBS=${MAX_JOBS} python3 setup.py install + # RUN SETUPTOOLS_SCM_PRETEND_VERSION=0.8.5.dev PYTORCH_ROCM_ARCH="gfx90a;gfx942" MAX_JOBS=${MAX_JOBS} python3 setup.py develop + WORKDIR /workspace/ + ########################################### + ########################################### + ########################################### + + + + + ######################################### + #### Install megatron-core############### + ######################################### + RUN pip uninstall -y megatron-core && \ + git clone https://github.com/yushengsu-thu/Megatron-LM-amd_version.git && \ + cd Megatron-LM-amd_version && \ + pip install -vvv -e . && \ + cd /workspace/ + ######################################### + ######################################### + ######################################### + + + + + ####################################### + ################apex################### + ####################################### + WORKDIR /workspace/ + RUN pip uninstall -y apex && \ + git clone git@github.com:ROCm/apex.git && \ + cd apex && \ + python setup.py install && \ + cd /workspace/ + ####################################### + ####################################### + ####################################### + + + ################################################################################ + ###########################Add torch_memory_saver############################### + ################################################################################ + # Set environment variables + ENV HIPCC_COMPILE_FLAGS_APPEND="--amdgpu-target=gfx90a;gfx942 -D__HIP_PLATFORM_AMD__" + ENV CFLAGS="-D__HIP_PLATFORM_AMD__" + ENV CXXFLAGS="-D__HIP_PLATFORM_AMD__" + RUN pip install "git+https://github.com/YangWang92/torch_memory_saver_numa.git@numa" + ################################################################################ + ################################################################################ + ################################################################################ + + + + ######################################## + ######Install ray####################### + ######################################## + # need to add this patch: https://github.com/ray-project/ray/pull/53531/files + RUN pip uninstall ray -y + RUN pip install "ray[data,train,tune,serve]>=2.47.0" + ######################################## + ######################################## + ######################################## + + + ########################################## + #######Install other dependencies######### + ########################################## + RUN pip install "tensordict==0.6.2" --no-deps && \ + pip install accelerate \ + codetiming \ + datasets \ + dill \ + hydra-core \ + liger-kernel \ + numpy \ + pandas \ + peft \ + "pyarrow>=15.0.0" \ + pylatexenc \ + torchdata \ + wandb \ + orjson \ + pybind11 + + WORKDIR /workspace/ + RUN git clone https://github.com/verl-project/verl.git && \ + cd verl && \ + pip install -e . + ########################################## + ########################################## + ########################################## + + WORKDIR /workspace/ + CMD ["/usr/bin/bash"] + + +Build the image: +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. code-block:: bash + + docker docker/build -t verl-rocm . + +Run the container +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Note: You can pull the docker from this DockerHub: [RLSys Foundation](https://hub.docker.com/u/yushengsuthu) +Pull the image: +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. code-block:: bash + + docker pull rlsys/verl:verl-0.4.1_ubuntu-22.04_rocm6.3.4-numa-patch_vllm0.8.5_sglang0.4.6.post4 + + docker tag rlsys/verl:verl-0.4.1_ubuntu-22.04_rocm6.3.4-numa-patch_vllm0.8.5_sglang0.4.6.post4 verl-rocm:latest + +Run the container +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + +Optional: Running without root and with user permissions +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. code-block:: bash + + docker run --rm -it \ + --device /dev/dri \ + --device /dev/kfd \ + -p 8265:8265 \ + --group-add video \ + --cap-add SYS_PTRACE \ + --security-opt seccomp=unconfined \ + --privileged \ + -v $HOME/.ssh:/root/.ssh \ + -v $HOME:$HOME \ + --shm-size 128G \ + -w $PWD \ + verl-rocm \ + /bin/bash + +(Optional): If you do not want to root mode and require assign yourself as the user +Please add ``-e HOST_UID=$(id -u)`` and ``-e HOST_GID=$(id -g)`` into the above docker launch script. + +Example +------- + +Due to to special setting in AMD (ROCM) torch, +1. If your ``ray>=2.45.0`` (default), you need to set ``RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES`` when starting ray in verl's RLHF training and add this [patch](https://github.com/ray-project/ray/pull/53531/files). +2. If your ``ray<2.45.0``, you need to set ``RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES`` when starting ray in verl's RLHF training. +Inference ``$ENGINE`` can be ``vllm`` or ``sglang``. We choose ``vllm`` as default in the following examples. + + + +PPO +~~~ + +.. code-block:: bash + + YOUR_PROJECT_NAME=r1-verl-ppo-upstream + YOUR_RUN_NAME=r1-training_ppo-upstream + # export HYDRA_FULL_ERROR=1 + + export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 + + # [ray] < 2.45.0 + #export RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES=1 + + # [ray] >= 2.45.0 + export RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES=1 # Patch with https://github.com/ray-project/ray/pull/52794 + + GPUS_PER_NODE=8 + MODEL_PATH=Qwen/Qwen2.5-0.5B-Instruct + python3 examples/data_preprocess/gsm8k.py --local_save_dir data/gsm8k + python3 -c "import transformers; transformers.pipeline('text-generation', model='$MODEL_PATH')" + ENGINE=vllm #sglang + + PYTHONUNBUFFERED=1 python3 -m verl.trainer.main_ppo \ + data.train_files=data/gsm8k/train.parquet \ + data.val_files=data/gsm8k/test.parquet \ + data.train_batch_size=256 \ + data.val_batch_size=1312 \ + data.max_prompt_length=512 \ + data.max_response_length=256 \ + actor_rollout_ref.model.path=$MODEL_PATH \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.actor.ppo_mini_batch_size=64 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=1 \ + actor_rollout_ref.rollout.name=$ENGINE \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \ + critic.optim.lr=1e-5 \ + critic.model.path=$MODEL_PATH \ + critic.ppo_micro_batch_size_per_gpu=4 \ + algorithm.kl_ctrl.kl_coef=0.001 \ + trainer.logger=console \ + trainer.project_name=$YOUR_PROJECT_NAME \ + trainer.experiment_name=$YOUR_RUN_NAME \ + trainer.val_before_train=False \ + trainer.n_gpus_per_node=$GPUS_PER_NODE \ + trainer.nnodes=1 \ + trainer.save_freq=10 \ + trainer.test_freq=10 \ + trainer.total_epochs=15 #2>&1 | tee verl_demo.log + +GRPO +~~~~ + +.. code-block:: bash + + YOUR_PROJECT_NAME=r1-verl-grpo-upstream + YOUR_RUN_NAME=r1-training_grpo-upstream + # export HYDRA_FULL_ERROR=1 + # export FSDP_VERBOSE=1 + + #export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 + + # [ray] < 2.45.0 + #export RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES=1 + + # [ray] >= 2.45.0 + export RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES=1 # Patch with https://github.com/ray-project/ray/pull/52794 + + GPUS_PER_NODE=8 + MODEL_PATH=Qwen/Qwen2.5-0.5B-Instruct + # MODEL_PATH=Qwen/Qwen2-7B-Instruct + python3 examples/data_preprocess/gsm8k.py --local_save_dir data/gsm8k + python3 -c "import transformers; transformers.pipeline('text-generation', model='$MODEL_PATH')" + ENGINE=vllm #sglang + + python3 -m verl.trainer.main_ppo \ + algorithm.adv_estimator=grpo \ + data.train_files=data/gsm8k/train.parquet \ + data.val_files=data/gsm8k/test.parquet \ + data.train_batch_size=1024 \ + data.val_batch_size=1312 \ + data.max_prompt_length=512 \ + data.max_response_length=1024 \ + actor_rollout_ref.model.path=$MODEL_PATH \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.model.use_remove_padding=True \ + actor_rollout_ref.actor.ppo_mini_batch_size=256 \ + actor_rollout_ref.actor.use_dynamic_bsz=True \ + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=24000 \ + actor_rollout_ref.actor.use_kl_loss=True \ + actor_rollout_ref.actor.kl_loss_coef=0.001 \ + actor_rollout_ref.actor.kl_loss_type=low_var_kl \ + actor_rollout_ref.model.enable_gradient_checkpointing=Flase \ + actor_rollout_ref.actor.fsdp_config.param_offload=False \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \ + actor_rollout_ref.rollout.tensor_model_parallel_size=2 \ + actor_rollout_ref.rollout.name=$ENGINE \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \ + actor_rollout_ref.rollout.n=5 \ + actor_rollout_ref.ref.fsdp_config.param_offload=False \ + algorithm.kl_ctrl.kl_coef=0.001 \ + trainer.critic_warmup=0 \ + trainer.logger=console \ + trainer.project_name=$YOUR_PROJECT_NAME \ + trainer.experiment_name=$YOUR_RUN_NAME \ + trainer.n_gpus_per_node=$GPUS_PER_NODE \ + trainer.val_before_train=False \ + trainer.nnodes=1 \ + trainer.save_freq=-1 \ + trainer.test_freq=10 \ + trainer.total_epochs=15 + + + +Multi-node training: slurm with Docker/Podman container +--------------------------------------------------------------------------------------- + +If you want to run multi-node training with slurm, you can use the following script. + +.. note:: + 1. You need to use ``podman`` or ``docker`` in the following script. We will release the apptainer script later. + 2. If you want to use ``podman``, you just replace ``docker`` with ``podman`` in the following script. + +The script includes the following steps: + +1. SLURM Configuration +2. Environment Setup +3. Docker/Podman Container Setup +4. Ray Cluster Initialization +5. Data Preprocessing +6. Model Setup +7. Training Launch + + +slurm_script.sh +~~~~~~~~~~~~~~~~~~~~ + +.. code-block:: bash + + #!/bin/bash + + #SBATCH --job-name=verl-ray-on-slurm + #SBATCH --nodes=2 + #SBATCH --ntasks-per-node=2 + #SBATCH --mem=200G + #SBATCH --time=30-00:00:00 + #SBATCH --gpus-per-node=8 + #SBATCH --cpus-per-task=28 + #SBATCH --output=../verl_log/slurm-%j.out + #SBATCH --error=../verl_log/slurm-%j.err + #SBATCH --nodelist=gpu-[0,1] + + + # load necessary modules + ### Run this setup + # [Cluster]: Use docker + # docker pull docker.io/rocm/vllm:rocm6.2_mi300_ubuntu20.04_py3.9_vllm_0.6.4 + + + ########################################################################## + ###The following setting should be set in different project and cluster### + ########################################################################## + + ### Project + CONTAINER_NAME="multinode_verl_training" + IMG="verl.rocm" + DOCKERFILE="docker/Dockerfile.rocm" + # echo $PWD + verl_workdir="${HOME}/projects/verl_upstream" + export TRANSFORMERS_CACHE="${HOME}/.cache/huggingface" + export HF_HOME=$TRANSFORMERS_CACHE + + ### Cluster Network Setting + export NCCL_DEBUG=TRACE + export GPU_MAX_HW_QUEUES=2 + export TORCH_NCCL_HIGH_PRIORITY=1 + export NCCL_CHECKS_DISABLE=1 + # export NCCL_IB_HCA=rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7 + export NCCL_IB_HCA=mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_8,mlx5_9 + export NCCL_IB_GID_INDEX=3 + export NCCL_CROSS_NIC=0 + export CUDA_DEVICE_MAX_CONNECTIONS=1 + export NCCL_PROTO=Simple + export RCCL_MSCCL_ENABLE=0 + export TOKENIZERS_PARALLELISM=false + export HSA_NO_SCRATCH_RECLAIM=1 + ########################################################################## + + ## Assign using GPUs + export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 + + ### For rocm and training script + # [ray] < 2.45.0 + #export RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES=1 + + # [ray] >= 2.45.0 + export RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES=1 # Patch with https://github.com/ray-project/ray/pull/52794 + + + # Build and launch the Docker container + srun bash -c " + # Exit on any error + set -e + + # Clean up dangling images (images with tag) + docker image prune -f + + # Need to pull the docker first + docker pull rlsys/verl:verl-0.4.1_ubuntu-22.04_rocm6.3.4-numa-patch_vllm0.8.5_sglang0.4.6.post4 + + if ! docker images --format "{{.Repository}}:{{.Tag}}" | grep -q "${IMG}"; then + echo \"Building ${IMG} image...\" + docker build -f \"${DOCKERFILE}\" -t \"${IMG}\" . + else + echo \"${IMG} image already exists, skipping build\" + fi + + # Removing old container if exists + docker rm \"${CONTAINER_NAME}\" 2>/dev/null || true + + # Checking network devices + ibdev2netdev + + # Launch the docker + docker run --rm -d \ + -e HYDRA_FULL_ERROR=1 \ + -e RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES=1 \ + -e RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES=1 \ + -e NCCL_DEBUG=${NCCL_DEBUG} \ + -e GPU_MAX_HW_QUEUES=${GPU_MAX_HW_QUEUES} \ + -e TORCH_NCCL_HIGH_PRIORITY=${TORCH_NCCL_HIGH_PRIORITY} \ + -e NCCL_CHECKS_DISABLE=${NCCL_CHECKS_DISABLE} \ + -e NCCL_IB_HCA=${NCCL_IB_HCA} \ + -e NCCL_IB_GID_INDEX=${NCCL_IB_GID_INDEX} \ + -e NCCL_CROSS_NIC=${NCCL_CROSS_NIC} \ + -e CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS} \ + -e NCCL_PROTO=${NCCL_PROTO} \ + -e RCCL_MSCCL_ENABLE=${RCCL_MSCCL_ENABLE} \ + -e TOKENIZERS_PARALLELISM=${TOKENIZERS_PARALLELISM} \ + -e HSA_NO_SCRATCH_RECLAIM=${HSA_NO_SCRATCH_RECLAIM} \ + -e TRANSFORMERS_CACHE=${TRANSFORMERS_CACHE} \ + -e HF_HOME=${HF_HOME} \ + --network host \ + --device /dev/dri \ + --device /dev/kfd \ + --device /dev/infiniband \ + --group-add video \ + --cap-add SYS_PTRACE \ + --security-opt seccomp=unconfined \ + --privileged \ + -v \${HOME}:\${HOME} \ + -v \${HOME}/.ssh:/root/.ssh \ + -w "${verl_workdir}" \ + --shm-size 128G \ + --name \"${CONTAINER_NAME}\" \ + \"${IMG}\" \ + tail -f /dev/null + + echo \"Container setup completed\" + " + # (Optional): If you do not want to root mode and require assign yuorself as the user + # Please add `-e HOST_UID=$(id -u)` and `-e HOST_GID=$(id -g)` into the above docker launch script. + + + + + + ### Ray launch the nodes before training + + # Getting the node names + nodes_array=($(scontrol show hostnames "$SLURM_JOB_NODELIST" | tr '\n' ' ')) + + head_node=${nodes_array[0]} + head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address) + + # if we detect a space character in the head node IP, we'll + # convert it to an ipv4 address. This step is optional. + if [[ "$head_node_ip" == *" "* ]]; then + IFS=' ' read -ra ADDR <<<"$head_node_ip" + if [[ ${#ADDR[0]} -gt 16 ]]; then + head_node_ip=${ADDR[1]} + else + head_node_ip=${ADDR[0]} + fi + echo "IPV6 address detected. We split the IPV4 address as $head_node_ip" + fi + + port=6379 + ip_head=$head_node_ip:$port + export ip_head + echo "IP Head: $ip_head" + + # make sure we set environment variables before Ray initialization + + # Print out all env variables + printenv + + echo "Starting HEAD at $head_node" + srun --nodes=1 --ntasks=1 -w "$head_node" \ + docker exec "${CONTAINER_NAME}" \ + ray start --head --node-ip-address="$head_node_ip" --port=$port \ + --dashboard-port=8266 \ + --num-cpus "${SLURM_CPUS_PER_TASK}" --num-gpus "${SLURM_GPUS_PER_NODE}" --block & + # optional, though may be useful in certain versions of Ray < 1.0. + sleep 10 + + # number of nodes other than the head node + worker_num=$((SLURM_JOB_NUM_NODES - 1)) + + for ((i = 1; i <= worker_num; i++)); do + node_i=${nodes_array[$i]} + echo "Debug: Starting worker on node_i = ${node_i}" + if [ -z "$node_i" ]; then + echo "Error: Empty node name for worker $i" + continue + fi + echo "Starting WORKER $i at $node_i" + srun --nodes=1 --ntasks=1 -w "$node_i" \ + docker exec "${CONTAINER_NAME}" \ + ray start --address "$ip_head" --num-cpus "${SLURM_CPUS_PER_TASK}" --num-gpus "${SLURM_GPUS_PER_NODE}" --block & + sleep 5 + done + + + + + # Ray initlization test (See whether any error in the above execution) + echo "Testing Ray initialization in the slurm nodes..." + docker exec "${CONTAINER_NAME}" python3 -c ' + import ray + try: + ray.init(address="auto") + print("\n=== Ray Cluster Status ===") + print(f"Number of nodes: {len(ray.nodes())}") + for node in ray.nodes(): + print("Node: {}, Status: {}".format(node["NodeManagerHostname"], node["Alive"])) + # print(f"Node: {node}") + ray.shutdown() + print("Ray initialization successful!") + except Exception as e: + print(f"Ray initialization failed: {str(e)}") + ' + echo "=== Ray test completed ===" + ###### + + + + # Run data preprocessing + + echo "Starting data preprocessing..." + docker exec "${CONTAINER_NAME}" \ + python3 "examples/data_preprocess/gsm8k.py" "--local_save_dir" "../data/gsm8k" + + echo "Starting data preprocessing..." + docker exec "${CONTAINER_NAME}" \ + python3 "examples/data_preprocess/math_dataset.py" "--local_dir" "../data/math" + + train_files="../data/gsm8k/train.parquet" + val_files="../data/gsm8k/test.parquet" + + # Download and test model + echo "Loading model..." + docker exec "${CONTAINER_NAME}" \ + python3 -c "import transformers; transformers.pipeline('text-generation', model='Qwen/Qwen2-7B-Instruct')" + MODEL_PATH="Qwen/Qwen2-7B-Instruct" + + # Set model path after pipeline test + MODEL_PATH="Qwen/Qwen2.5-0.5B-Instruct" + + echo "== Data and model loading Done ==" + + echo "Start to train..." + + docker exec "${CONTAINER_NAME}" \ + python3 -c "import transformers; transformers.pipeline('text-generation', model='Qwen/Qwen2-7B-Instruct')" + MODEL_PATH="Qwen/Qwen2-7B-Instruct" + + + PYTHONUNBUFFERED=1 srun --overlap --nodes=${SLURM_NNODES} --ntasks=1 -w "$head_node" \ + docker exec "${CONTAINER_NAME}" \ + python3 -m verl.trainer.main_ppo \ + data.train_files=$train_files \ + data.val_files=$val_files \ + data.train_batch_size=1024 \ + data.max_prompt_length=1024 \ + data.max_response_length=1024 \ + actor_rollout_ref.model.path=$MODEL_PATH \ + actor_rollout_ref.model.enable_gradient_checkpointing=False \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.model.use_remove_padding=True \ + actor_rollout_ref.actor.ppo_mini_batch_size=256 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=8 \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.fsdp_config.param_offload=False \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=2 \ + actor_rollout_ref.rollout.name=vllm \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.9 \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \ + actor_rollout_ref.ref.fsdp_config.param_offload=True \ + critic.optim.lr=1e-5 \ + critic.model.use_remove_padding=True \ + critic.model.path=$MODEL_PATH \ + critic.model.enable_gradient_checkpointing=False \ + critic.ppo_micro_batch_size_per_gpu=8 \ + critic.fsdp.param_offload=False \ + critic.fsdp.optimizer_offload=False \ + algorithm.kl_ctrl.kl_coef=0.0001 \ + trainer.critic_warmup=0 \ + trainer.logger='["console","wandb"]' \ + trainer.project_name='verl_example' \ + trainer.experiment_name='Qwen2.5-32B-Instruct_function_rm' \ + trainer.n_gpus_per_node=${SLURM_GPUS_PER_NODE} \ + trainer.val_before_train=False \ + trainer.nnodes=${SLURM_NNODES} \ + trainer.save_freq=-1 \ + trainer.test_freq=10 \ + trainer.total_epochs=15 + + +Run slurm_script.sh +~~~~~~~~~~~~~~~~~~~~ +Just sbatch your slurm_script.sh + +.. code-block:: bash + + sbatch slurm_script.sh + diff --git a/verl/docs/amd_tutorial/amd_quick_start.rst b/verl/docs/amd_tutorial/amd_quick_start.rst new file mode 100644 index 0000000000000000000000000000000000000000..48dcd1d5b3c9ec271c5df5384f8e55d0c9c2a91c --- /dev/null +++ b/verl/docs/amd_tutorial/amd_quick_start.rst @@ -0,0 +1,136 @@ +Getting started with AMD ROCm +========================================= + +Last updated: 05/17/2026. + +Author: `Mingjie Lu `_, `Xiaohong Kou `_, `Fuwei Yang `_ + +Overview +-------- + +This document is a quick-start tutorial for running VeRL on AMD ROCm. +It provides a production-style bring-up flow for container startup, environment +verification, and training examples. + +Current software and hardware scope: + +- Runtime modes: fully supports **Fully Async** and **Colocate**. +- Inference engine: **vLLM** validated; **SGLang** support is ongoing. +- Trainer backends: **FSDP**, **FSDP2** and **Megatron**. +- GPU targets: + + - MI300X / MI325X (``gfx942``) + - MI355X (``gfx950``) + +Software Baseline +----------------- + +Use the following prebuilt image for tutorial and validation: + +- ``amdagi/verl-dev:rocm7.0.2_56_te2.10_vllm0.20_py312`` + +The Docker build recipe remains unchanged: + +- `docker/rocm/Dockerfile.rocm `_ + +Host Prerequisites +------------------ + +Before launching the container, ensure: + +1. AMD ROCm 7.0.2 host driver stack is installed and healthy. +2. Docker has access to ``/dev/kfd`` and ``/dev/dri``. +3. Dataset and model storage paths are ready. + +Launch Container +---------------- + +.. code-block:: bash + + NAME=verl_release + DOCKER=amdagi/verl-dev:rocm7.0.2_56_te2.10_vllm0.20_py312 + + docker pull $DOCKER + + docker run -it --name $NAME --device /dev/kfd --device /dev/dri \ + --privileged --network=host \ + --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \ + --shm-size=2048g \ + --ulimit memlock=-1 --ulimit stack=67108864 \ + -w /workspace \ + $DOCKER \ + /bin/bash + +Environment Check (Inside Container) +------------------------------------ + +.. code-block:: bash + + # ROCm and visible GPU targets + rocminfo | grep -E "gfx942|gfx950" || true + + # PyTorch + ROCm sanity check + python - <<'PY' + import torch + print("torch:", torch.__version__) + print("rocm :", torch.version.hip) + print("cuda_available:", torch.cuda.is_available()) + if torch.cuda.is_available(): + print("gpu_count:", torch.cuda.device_count()) + print("device_0:", torch.cuda.get_device_name(0)) + PY + +Feature Support Matrix +---------------------- + +.. list-table:: Current support status + :header-rows: 1 + + * - Category + - Status + - Notes + * - Runtime mode + - Fully supported + - Fully Async and Colocate are production-ready + * - Inference engine + - vLLM validated + - SGLang integration is ongoing + * - Trainer backend + - Fully supported + - FSDP, Megatron + * - Hardware + - Fully supported + - MI300X / MI325X (gfx942), MI355X (gfx950) + +Example Workflow +---------------- + +1) Colocate mode + FSDP (GRPO, Qwen3-8B) +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +For Qwen3-8B FSDP training, enable both parameter and optimizer offload to avoid OOM. + +.. code-block:: bash + + # Configure these in your launch script or Hydra overrides: + # actor_rollout_ref.actor.fsdp_config.param_offload=True + # actor_rollout_ref.actor.fsdp_config.optimizer_offload=True + bash examples/grpo_trainer/run_qwen3_8b_fsdp.sh + +2) Colocate mode + Megatron (GRPO, Qwen3.5-35B) +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. code-block:: bash + + bash examples/grpo_trainer/run_qwen3_5-35b-megatron.sh + +3) Fully Async mode +~~~~~~~~~~~~~~~~~~~ + +``RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES`` and +``RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES`` are no longer required in this release. + +.. code-block:: bash + + # For qwen2.5-math-7b, update max_position_embeddings to 32768 in config.json after model download. + bash verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_4_4.sh diff --git a/verl/docs/amd_tutorial/amd_vllm_page.rst b/verl/docs/amd_tutorial/amd_vllm_page.rst new file mode 100644 index 0000000000000000000000000000000000000000..7c230acab8792406e0ecb82d1a4fb417ba027a2e --- /dev/null +++ b/verl/docs/amd_tutorial/amd_vllm_page.rst @@ -0,0 +1,41 @@ +verl performance tuning for AMD (ROCm Kernel) +===================================================== + +Last updated: 11/13/2025. + +Author: `Yang Wang `_, `Songlin Jiang `_ + +Use vLLM Sleep Mode for AMD MI3xx series GPUs +-------------------------------------------------------------- + +By default, verl requires vLLM to enable sleep mode, which allows vLLM to offload GPU memory to CPU memory after rollout. This feature has been merged into the main branch of vLLM for version later than 0.11.0. + +For now, you can use the vLLM main branch and build it from the source code, or you can directly install vLLM from the pre-built ROCm wheels for vLLM version later than 0.11.0 when it's available. + +1. Clone the vLLM repository and build it with the following commands: + +.. code-block:: bash + + git clone https://github.com/vllm-project/vllm.git + cd vllm + git reset --hard 4ca5cd5740c0cd7788cdfa8b7ec6a27335607a48 # You can also use a later commit as you wish + python -m pip install -r requirements/rocm.txt + VLLM_TARGET_DEVICE=rocm ROCM_PATH=/opt/rocm/ python3 setup.py develop + +2. Additionally, we recommend you to use the ROCm version later than or equal to ROCm 7.0. + +After the upgrade, you can verify whether sleep mode is working by trying out `these scripts `_. + +If sleep mode is working, you should see the memory usage reduce after sleep. + +After applying the vLLM patch and completing the installation, you can enable sleep mode in verl to reduce memory overhead. This allows verl to offload unused GPU memory during rollout, significantly lowering the memory footprint during long-context training or multi-node reinforcement learning. + + +Enable CUDA Graph and Bypass ROCm-related issues +-------------------------------------------------------------- + +Due to potential issues with CUDA graph capture in ROCm, we've found that vLLM's CUDA graph feature cannot be enabled on multiple nodes in verl on AMD platforms with vLLM V1 mode. This leads to significantly slower rollout performance. + +Our investigation shows that ROCm may trigger an unexpected crash when attempting to capture large batches with CUDA graph. One workaround is to set ``actor_rollout_ref.rollout.cudagraph_capture_sizes`` to values such as ``[1, 2, 4, 8, 16, 32, 64]`` (change depending on your GPU memory size). + +Then, you can choose to enable CUDA graph by setting ``actor_rollout_ref.rollout.enforce_eager`` to ``False`` in your verl configuration file. diff --git a/verl/docs/api/data.rst b/verl/docs/api/data.rst new file mode 100644 index 0000000000000000000000000000000000000000..5baa5b51bfdb79f6ead72f1f46141720248bd813 --- /dev/null +++ b/verl/docs/api/data.rst @@ -0,0 +1,61 @@ +Data interface +========================= + +Last updated: 05/19/2025 (API docstrings are auto-generated). + +DataProto is the interface for data exchange. + +The :class:`verl.DataProto` class contains two key members: + +- batch: a :class:`tensordict.TensorDict` object for the actual data +- meta_info: a :class:`Dict` with additional meta information + +TensorDict +~~~~~~~~~~~~ + +:attr:`DataProto.batch` is built on top of :class:`tensordict`, a project in the PyTorch ecosystem. +A TensorDict is a dict-like container for tensors. To instantiate a TensorDict, you must specify key-value pairs as well as the batch size. + +.. code-block:: python + + >>> import torch + >>> from tensordict import TensorDict + >>> tensordict = TensorDict({"zeros": torch.zeros(2, 3, 4), "ones": torch.ones(2, 3, 5)}, batch_size=[2,]) + >>> tensordict["twos"] = 2 * torch.ones(2, 5, 6) + >>> zeros = tensordict["zeros"] + >>> tensordict + TensorDict( + fields={ + ones: Tensor(shape=torch.Size([2, 3, 5]), device=cpu, dtype=torch.float32, is_shared=False), + twos: Tensor(shape=torch.Size([2, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False), + zeros: Tensor(shape=torch.Size([2, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([2]), + device=None, + is_shared=False) + +One can also index a tensordict along its batch_size. The contents of the TensorDict can be manipulated collectively as well. + +.. code-block:: python + + >>> tensordict[..., :1] + TensorDict( + fields={ + ones: Tensor(shape=torch.Size([1, 3, 5]), device=cpu, dtype=torch.float32, is_shared=False), + twos: Tensor(shape=torch.Size([1, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False), + zeros: Tensor(shape=torch.Size([1, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([1]), + device=None, + is_shared=False) + >>> tensordict = tensordict.to("cuda:0") + >>> tensordict = tensordict.reshape(6) + +For more about :class:`tensordict.TensorDict` usage, see the official tensordict_ documentation. + +.. _tensordict: https://pytorch.org/tensordict/stable/overview.html + + +Core APIs +~~~~~~~~~~~~~~~~~ + +.. autoclass:: verl.DataProto + :members: to, select, union, make_iterator, concat diff --git a/verl/docs/api/single_controller.rst b/verl/docs/api/single_controller.rst new file mode 100644 index 0000000000000000000000000000000000000000..44ea366ffe4b12ce5293821877ce70a0073f2152 --- /dev/null +++ b/verl/docs/api/single_controller.rst @@ -0,0 +1,30 @@ +Single Controller interface +============================ + +Last updated: 05/27/2025 (API docstrings are auto-generated). + +The Single Controller provides a unified interface for managing distributed workers +using Ray or other backends and executing functions across them. +It simplifies the process of dispatching tasks and collecting results, particularly +when dealing with data parallelism or model parallelism. + + +Core APIs +~~~~~~~~~~~~~~~~~ + +.. autoclass:: verl.single_controller.Worker + :members: __init__, __new__, get_master_addr_port, get_cuda_visible_devices, world_size, rank + +.. autoclass:: verl.single_controller.WorkerGroup + :members: __init__, world_size + +.. autoclass:: verl.single_controller.ClassWithInitArgs + :members: __init__, __call__ + +.. autoclass:: verl.single_controller.ResourcePool + :members: __init__, world_size, local_world_size_list, local_rank_list + +.. autoclass:: verl.single_controller.ray.RayWorkerGroup + :members: __init__ + +.. autofunction:: verl.single_controller.ray.create_colocated_worker_cls \ No newline at end of file diff --git a/verl/docs/api/trainer.rst b/verl/docs/api/trainer.rst new file mode 100644 index 0000000000000000000000000000000000000000..abfa51f01a31606f436a95fde13770577b9ab540 --- /dev/null +++ b/verl/docs/api/trainer.rst @@ -0,0 +1,31 @@ +Trainer Interface +================================ + +Last updated: 06/08/2025 (API docstrings are auto-generated). + +Trainers drive the training loop. Introducing new trainer classes in case of new training paradiam is encouraged. + +.. autosummary:: + :nosignatures: + + verl.trainer.ppo.ray_trainer.RayPPOTrainer + + +Core APIs +~~~~~~~~~~~~~~~~~ + +.. autoclass:: verl.trainer.ppo.ray_trainer.RayPPOTrainer + :members: __init__, init_workers, fit + +.. automodule:: verl.utils.tokenizer + :members: hf_tokenizer + +.. automodule:: verl.trainer.ppo.core_algos + :members: agg_loss, kl_penalty, compute_policy_loss, kl_penalty + +.. automodule:: verl.trainer.ppo.reward + :members: load_reward_manager, compute_reward, compute_reward_async + +.. autoclass:: verl.workers.reward_manager.NaiveRewardManager + +.. autoclass:: verl.workers.reward_manager.DAPORewardManager diff --git a/verl/docs/api/utils.rst b/verl/docs/api/utils.rst new file mode 100644 index 0000000000000000000000000000000000000000..e15e3a5a32bdbb129a25d93b12e751385caa30b5 --- /dev/null +++ b/verl/docs/api/utils.rst @@ -0,0 +1,76 @@ +Utilities +============ + +Last updated: 05/19/2025 (API docstrings are auto-generated). + +This section documents the utility functions and classes in the VERL library. + +Python Functional Utilities +------------------------------ + +.. automodule:: verl.utils.py_functional + :members: append_to_dict + +File System Utilities +------------------------ + +.. automodule:: verl.utils.fs + :members: copy_to_local + +Tracking Utilities +--------------------- + +.. automodule:: verl.utils.tracking + :members: Tracking + +Metrics Utilities +--------------------- + +.. automodule:: verl.utils.metric + :members: reduce_metrics + +Checkpoint Management +------------------------ + +.. automodule:: verl.utils.checkpoint.checkpoint_manager + :members: find_latest_ckpt_path + +.. automodule:: verl.utils.checkpoint.fsdp_checkpoint_manager + :members: FSDPCheckpointManager + +Dataset Utilities +--------------------- + +.. automodule:: verl.utils.dataset.rl_dataset + :members: RLHFDataset, collate_fn + +Torch Functional Utilities +----------------------------- + +.. automodule:: verl.utils.torch_functional + :members: get_constant_schedule_with_warmup, masked_whiten, masked_mean, logprobs_from_logits + +Sequence Length Balancing +---------------------------- + +.. automodule:: verl.utils.seqlen_balancing + :members: get_reverse_idx, rearrange_micro_batches + +Ulysses Utilities +-------------------- + +.. automodule:: verl.utils.ulysses + :members: gather_outputs_and_unpad, ulysses_pad_and_slice_inputs + +FSDP Utilities +------------------ + +.. automodule:: verl.utils.fsdp_utils + :members: get_fsdp_wrap_policy, get_init_weight_context_manager, init_fn, load_fsdp_model_to_gpu, load_fsdp_optimizer, offload_fsdp_model_to_cpu, offload_fsdp_optimizer, + +Debug Utilities +------------------- + +.. automodule:: verl.utils.profiler + :members: log_gpu_memory_usage, GPUMemoryLogger + diff --git a/verl/docs/ascend_tutorial/README.md b/verl/docs/ascend_tutorial/README.md new file mode 100644 index 0000000000000000000000000000000000000000..763698c300209d7e13af3fde5a661b81b4611eb9 --- /dev/null +++ b/verl/docs/ascend_tutorial/README.md @@ -0,0 +1,69 @@ + +## 简介 + +昇腾全面支持 Verl 使用与开发，本文档全面介绍了如何在华为昇腾芯片 NPU 上使用 Verl。 + +Last updated: 05/14/2026. + +## 目录结构 + +``` +ascend_tutorial/ +├── get_start/ # 快速入门指南 +├── feature_support/ # 特性支持说明 +├── model_support/ # 模型支持说明 +├── dev_guide/ # 开发指南 +├── faq/ # 常见问题解答 +└── contribution_guide/ # 社区贡献指南 +``` +## 最新消息 +- [verl-ascend-recipe 建仓](https://github.com/verl-project/verl-ascend-recipe) - 新增昇腾recipe +- [verl on ascend 2026Q2 roadmap](https://github.com/verl-project/verl/issues/5526) - 2026Q2 RoadMap 已发布 + +## 快速开始 +- [Docker 构建使用指南](./get_start/dockerfile_build_guidance.rst) - 构建并使用昇腾环境的 Docker 镜像 +- [自定义环境安装](./get_start/install_guidance.rst) - 在昇腾 NPU 上自定义安装 Verl +- [快速上手](./get_start/quick_start.rst) - 快速上手在昇腾 NPU 上运行 Verl + +## 特性支持说明 + +- [verl特性支持](./dev_guide/model_dev/parameter_and_metrics.md) - 支持的verl框架特性/参数列表 +- [NPU特性支持](./feature_support/npu_advance_features.md) - NPU相关常用特性/环境变量说明 + +## 模型支持说明 + +- [模型与算法支持说明](./model_support/model_and_algorithm_support.md) - 支持的模型/算法列表 +- [最佳实践示例](./model_support/examples) - 最佳实践与模型部署示例 + + +## 开发指南 + +- [模型开发](./dev_guide/model_dev) + - [模型迁移](./dev_guide/model_dev/transfer_to_npu_guide.md) - 模型迁移指南 + - [训练参数与指标](./dev_guide/model_dev/parameter_and_metrics.md) - 训练参数与指标 + - [模型评测](./dev_guide/model_dev/evaluation.md) - 模型评测指南 +- [精度调试](./dev_guide/precision_analysis) + - [精度分析](./dev_guide/precision_analysis/precision_alignment_zh.md) - 精度对齐指南 + - [精度调试器](./dev_guide/precision_analysis/precision_debugger_zh.md) - 精度问题排查工具 +- [性能调优](./dev_guide/performance) + - [性能分析](./dev_guide/performance/ascend_performance_analysis_guide.md) - 性能分析指南 + - [性能调优](./dev_guide/performance/perf_tuning_on_ascend.rst) - 性能调优指南 + - [profiling采集](./dev_guide/performance/ascend_profiling_zh.rst) - profiling 工具使用指南 + + +## 支持与反馈 + +如果您在使用过程中遇到问题，欢迎通过以下方式获取帮助： + +1. 查看 [FAQ](./faq/faq.rst) +2. 在 GitHub Issues 中提交问题 +3. 联系昇腾技术支持 + +## 贡献指南 +- [verl 社区贡献](../contributing) - Verl 社区贡献指南 +- [昇腾 CI 指南](./contribution_guide/ascend_ci_guide_zh.rst) - 昇腾环境 CI 配置与测试 + +## 相关资源 + +- [Verl 官方文档](https://verl.readthedocs.io/) +- [昇腾开发者社区](https://www.hiascend.com/) \ No newline at end of file diff --git a/verl/docs/ascend_tutorial/contribution_guide/ascend_ci_guide_zh.rst b/verl/docs/ascend_tutorial/contribution_guide/ascend_ci_guide_zh.rst new file mode 100644 index 0000000000000000000000000000000000000000..87171b0f1b1d1971b33d3cd504ac2183947a060f --- /dev/null +++ b/verl/docs/ascend_tutorial/contribution_guide/ascend_ci_guide_zh.rst @@ -0,0 +1,171 @@ +NPU-CI 添加指导 +=========== + +Last updated: 02/02/2026. + +我们在 verl 上增加基于华为昇腾设备的CI用例添加指导。 + +verl 仓库使用 GitHub Actions 作为 CI 平台，通过分层测试架构保障代码质量与系统稳定性。 +NPU 相关的工作流主要包括： + +* ``npu_unit_test.yml``：运行单元测试。 +* 以 ``_ascend.yml`` 结尾的文件：运行针对 Ascend NPU 的端到端测试或专项测试。 + +添加新用例指南 +----------------------------------- + +1. 数据集与权重 +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +流水机器上的权重与绝对路径： + ++---------------------------------------+-------------------------------------------------------------------+ +| 模型名称 | 绝对路径 | ++=======================================+===================================================================+ +| Qwen3-30B-A3B-Instruct-2507 | ``${HOME}/.cache/models/Qwen/Qwen3-30B-A3B-Instruct-2507`` | ++---------------------------------------+-------------------------------------------------------------------+ +| Qwen2.5-VL-3B-Instruct | ``${HOME}/.cache/models/Qwen/Qwen2.5-VL-3B-Instruct`` | ++---------------------------------------+-------------------------------------------------------------------+ +| Qwen2.5-0.5B | ``${HOME}/.cache/models/Qwen/Qwen2.5-0.5B`` | ++---------------------------------------+-------------------------------------------------------------------+ +| Qwen2.5-0.5B-Instruct | ``${HOME}/.cache/models/Qwen/Qwen2.5-0.5B-Instruct`` | ++---------------------------------------+-------------------------------------------------------------------+ +| Qwen2.5-1.5B-Instruct | ``${HOME}/.cache/models/Qwen/Qwen2.5-1.5B-Instruct`` | ++---------------------------------------+-------------------------------------------------------------------+ +| Skywork-Reward-V2-Llama-3.2-1B | ``${HOME}/.cache/models/Skywork/Skywork-Reward-V2-Llama-3.2-1B`` | ++---------------------------------------+-------------------------------------------------------------------+ + +流水机器上的数据集与绝对路径： + ++--------------+---------------------------------------------------+ +| 数据集名称 | 绝对路径 | ++==============+===================================================+ +| gsm8k | ``${HOME}/.cache/datasets/openai/gsm8k`` | ++--------------+---------------------------------------------------+ +| geo3k | ``${HOME}/.cache/datasets/hiyouga/geometry3k`` | ++--------------+---------------------------------------------------+ + +**Note** + + ${HOME}是root + + gpu用例中权重在~/models/路径下，如需适配可以用软链接，``ln -s /root/.cache/models ~/models`` + + 此处为原始数据集，按需进行数据处理，如下。 + + ``python examples/data_preprocess/gsm8k_multiturn_sft.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k`` + + +2. 工作流 YAML 模板 +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +如需新增一个工作流，可参考以下模板创建 ``.github/workflows/your_yml_ascend.yml`` 文件。 + +主要修改部分包括： + +* 工作流名称（``name``） +* 触发条件（``on``） +* 运行环境（``runs-on``） +* 容器镜像（``container.image``） +* 具体执行步骤（``jobs..steps``） + +.. code-block:: yaml + :linenos: + + name: your_yml_ascend # 工作流唯一标识 + # 触发条件配置 + on: + push: + branches: + - main + - v0.* + pull_request: + branches: + - main + paths: + - ".github/workflows/your_yml_ascend.yml" # 必须包含此工作流文件路径 + - "path/to/affected_files" # 需监控的相关代码路径 + + # 并发控制策略 + concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} # 仅非main分支取消进行中的任务 + + permissions: + contents: read # 最小权限原则 + + jobs: + your_job_name: # 任务唯一标识 + if: github.repository_owner == 'verl-project' # 仅在主仓库运行 + runs-on: linux-aarch64-a2-4 # 硬件规格：a2实例，4卡NPU + timeout-minutes: 60 # 任务超时阈值（分钟） + container: + #运行镜像 该示例为vllm的镜像 + image: swr.ap-southeast-1.myhuaweicloud.com/base_image/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest + options: >- + --shm-size 16g # 共享内存配置 + env: + HF_ENDPOINT: "https://hf-mirror.com" + HF_HUB_ENABLE_HF_TRANSFER: "0" + steps: + - name: Check npu and CANN info + run: | + cat /usr/local/Ascend/ascend-toolkit/latest/"$(uname -i)"-linux/ascend_toolkit_install.info + npu-smi info + - name: Check initial pip list from image + run: pip list + - name: Checkout repository + uses: actions/checkout@v4 + with: + fetch-depth: 0 + clean: true + - name: Install dependencies + run: | + pip install --no-deps -e . + - name: Verify environment + run: pip list + # 以下为具体测试步骤（根据需求定制） + - name: Preprocess dataset + run: python examples/data_preprocess/your_script.py --local_dataset_path ${HOME}/.cache/datasets/your_dataset + - name: Execute NPU test + run: | + ray stop --force + bash tests/special_npu/your_test_script.sh + +**Note** + + + ${HOME}/.cache/文件夹内一旦添加新内容，不会因CI跑完容器销毁而删除，请避免往该文件夹添加内容。 + + +3. 添加单元测试 +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +步骤： + +(1) 在 ``tests/`` 目录下创建或修改单元测试文件（例如 ``test_xxx.py``）。 +(2) 若测试文件路径未被 ``npu_unit_test.yml`` 中的 ``--ignore-glob`` 规则排除，则会在以下命令中自动执行： + + .. code-block:: yaml + + pytest -s -x --ignore-glob="xxx" --ignore-glob="xxx" tests/ + +(3) 若测试路径在 ``--ignore-glob`` 排除范围内，需在 ``npu_unit_test.yml`` 中新增一个 step 来显式运行该测试。 +(4) 如新增一批相关用例，建议单独创建专门的工作流文件以保持清晰。 + +4. 添加端到端测试脚本 +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +步骤： + +(1) 在 ``tests/special_npu/`` 目录下创建端到端测试脚本。 +(2) 在 ``.github/workflows/`` 目录中找到功能最接近的以 ``_ascend.yml`` 结尾的工作流文件，在其中添加一个 step 调用该脚本。 +(3) 若测试场景独立或较复杂，可考虑单独创建新的工作流文件。 + +5. 测试策略建议 +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +* **单元测试**：覆盖核心函数、类与方法，确保逻辑正确。 +* **集成/端到端测试**：覆盖典型训练、推理 pipeline，验证多模块协同与硬件适配。 +* **资源管理**：一个workflow里的多个job为并行运行，请合理设置超时时间，避免任务长时间挂起，请控制单个 job 的运行时间在 40min 以内。 + +通过以上步骤，可系统化地为 verl 仓库添加 NPU 相关的自动化测试，确保代码变更在合并前经过充分验证。 diff --git a/verl/docs/ascend_tutorial/dev_guide/model_dev/evaluation.md b/verl/docs/ascend_tutorial/dev_guide/model_dev/evaluation.md new file mode 100644 index 0000000000000000000000000000000000000000..6e021a02a8f45df30effc2c84f01e885d4bb2bd7 --- /dev/null +++ b/verl/docs/ascend_tutorial/dev_guide/model_dev/evaluation.md @@ -0,0 +1,105 @@ +# 模型评测 + +Last updated: 05/14/2026. + +不同模型步骤一致,仅以Qwen3-30B为例列举 + +我们通过 AISBenchmark 评估模型,该工具支持vllm/sglang多种推理后端的评估 + +## 1.安装方法 + +~~~bash +git clone https://gitee.com/aisbench/benchmark.git +cd benchmark +pip install -e . +~~~ + + +## 2.下载评估数据集 + +~~~bash +cd path/to/benchmark/ais_bench/datasets +wget http://opencompass.oss-cn-shanghai.aliyuncs.com/datasets/data/math.zip +unzip math.zip +rm math.zip +~~~ + +## 3.权重转换 + +当前verl已经支持mbridge直接保存hf格式模型权重,无需转换即可使用. + +如果模型权重不是hf格式,需要先转换为hf格式,再进行评估 + +此处参照verl原生[转换方法](verl\docs\advance\checkpoint.rst) + +## 4.vllm推理评测 + +**启动vllm_server服务** + +通过以下命令拉起NPU服务端，需要修改的参数：model和tensor-parallel-size。 + +model：保存训练后权重转换完的huggingface模型地址； + +tensor-parallel-size：张量并行副本数，TP建议和训练时infer的配置保持一致； + +data-parallel-size：数据并行副本数，DP建议和训练时infer的配置保持一致，默认为1； + +port：可任意设置空闲端口； + +~~~bash +vllm serve /path/to/Qwen3-30B/ \ + --served-model-name auto \ + --gpu-memory-utilization 0.9 \ + --max-num-seqs 24 \ + --max-model-len 22528 \ + --max-num-batched-tokens 22528 \ + --enforce-eager \ + --trust-remote-code \ + --distributed_executor_backend=mp \ + --tensor-parallel-size 8 \ + --data-parallel-size 1 \ + --generation-config vllm \ + --port 6380 +~~~ + +**修改aisbench推理配置启动vllm_client评测** + +打开推理配置文件 benchmark/ais_bench/benchmark/configs/models/vllm_api/vllm_api_stream_chat.py + +host_port需与服务端的port一致，根据模型配置修改max_seq_len和max_out_len +~~~bash +from ais_bench.benchmark.models import VLLMCustomAPIChatStream +from ais_bench.benchmark.utils.model_postprocessors import extract_non_reasoning_content + +models = [ + dict( + attr="service", + type=VLLMCustomAPIChatStream, + abbr='vllm-api-stream-chat', + path="", + model="", + request_rate = 0, + retry = 2, + host_ip = "localhost", + host_port = 8080, + max_out_len = 512, + batch_size=1, + trust_remote_code=False, + generation_kwargs = dict( + temperature = 0.5, + top_k = 10, + top_p = 0.95, + seed = None, + repetition_penalty = 1.03, + ), + pred_postprocessor=dict(type=extract_non_reasoning_content) + ) +] +~~~ + +另起一个窗口进行评测，开启评测命令： +~~~bash + ais_bench --models vllm_api_stream_chat --datasets math500_gen_0_shot_cot_chat_prompt +~~~ +## 5.sglang推理评测 +参照 [sglang最佳实践](../../model_support/examples/ascend_sglang_best_practices.rst)中评测进行 \ No newline at end of file diff --git a/verl/docs/ascend_tutorial/dev_guide/model_dev/parameter_and_metrics.md b/verl/docs/ascend_tutorial/dev_guide/model_dev/parameter_and_metrics.md new file mode 100644 index 0000000000000000000000000000000000000000..4258076203a1c50a7eee48538944227592ce9e78 --- /dev/null +++ b/verl/docs/ascend_tutorial/dev_guide/model_dev/parameter_and_metrics.md @@ -0,0 +1,880 @@ +# 训练配置参数与指标说明 + +Last updated: 05/12/2026. + +verl 通过层级化的 YAML 配置文件管理所有参数，涉及到的所有配置文件均在 `verl\trainer\config` 目录下。 + +--- + +## 1. 配置参数说明 + +### 1.1 公共配置参数 + +以下参数在 FSDP 方案和 Megatron 方案中均存在且含义一致。 + +#### 1.1.1 Actor 优化器配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `actor_rollout_ref.actor.optim.lr` | `1.0e-06` | Actor 学习率 | +| `actor_rollout_ref.actor.optim.lr_warmup_steps_ratio` | `0.0` | 学习率预热步数占总训练步数的比例 | +| `actor_rollout_ref.actor.optim.total_training_steps` | `-1` | 总训练步数，-1 表示自动计算 | +| `actor_rollout_ref.actor.optim.weight_decay` | `0.01` | 权重衰减，用于防止模型过拟合 | +| `actor_rollout_ref.actor.optim.lr_warmup_steps` | `-1` | 学习率预热步数，-1 表示由 ratio 自动计算 | +| `actor_rollout_ref.actor.optim.betas` | `[0.9, 0.999]` | Adam 优化器的一阶和二阶动量系数 | +| `actor_rollout_ref.actor.optim.clip_grad` | `1.0` | 梯度裁剪阈值 | +| `actor_rollout_ref.actor.optim.override_optimizer_config` | `null` / `{}` | 覆盖优化器配置（FSDP 为 null，Megatron 为 {}） | + +#### 1.1.2 Actor 策略配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `actor_rollout_ref.actor.strategy` | `fsdp` / `megatron` | 训练策略，FSDP 方案为 fsdp，Megatron 方案为 megatron | +| `actor_rollout_ref.actor.ppo_mini_batch_size` | `256` | PPO 训练的 mini batch 大小 | +| `actor_rollout_ref.actor.ppo_micro_batch_size` | `null` | PPO 训练的 micro batch 大小 | +| `actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu` | `null` | 每 GPU 的 PPO micro batch 大小 | +| `actor_rollout_ref.actor.use_dynamic_bsz` | `false` | 是否使用动态 batch size | +| `actor_rollout_ref.actor.ppo_max_token_len_per_gpu` | `16384` | 每 GPU 的 PPO 最大 token 长度 | +| `actor_rollout_ref.actor.clip_ratio` | `0.2` | PPO 裁剪比例，控制策略更新幅度，一般取值范围 [0.1, 0.3] | +| `actor_rollout_ref.actor.clip_ratio_low` | `0.2` | PPO 下界裁剪比例 | +| `actor_rollout_ref.actor.clip_ratio_high` | `0.2` | PPO 上界裁剪比例 | +| `actor_rollout_ref.actor.tau_pos` | `1.0` | 正优势裁剪的 tau 参数 | +| `actor_rollout_ref.actor.tau_neg` | `1.05` | 负优势裁剪的 tau 参数 | +| `actor_rollout_ref.actor.freeze_vision_tower` | `false` | 是否冻结视觉塔（多模态模型） | +| `actor_rollout_ref.actor.clip_ratio_c` | `3.0` | 裁剪比例的上限常数 | +| `actor_rollout_ref.actor.loss_agg_mode` | `token-mean` | 损失聚合模式，可选 token-mean 等 | +| `actor_rollout_ref.actor.loss_scale_factor` | `null` | 损失缩放因子 | +| `actor_rollout_ref.actor.entropy_coeff` | `0` | 熵正则化系数，控制策略探索程度 | +| `actor_rollout_ref.actor.calculate_entropy` | `false` | 是否计算策略熵 | +| `actor_rollout_ref.actor.use_kl_loss` | `false` | 是否使用 KL 散度损失 | +| `actor_rollout_ref.actor.use_prefix_grouper` | `false` | 是否使用前缀分组器 | +| `actor_rollout_ref.actor.use_torch_compile` | `true` | 是否使用 torch.compile 加速 | +| `actor_rollout_ref.actor.kl_loss_coef` | `0.001` | KL 损失系数 | +| `actor_rollout_ref.actor.kl_loss_type` | `low_var_kl` | KL 损失类型，可选 low_var_kl 等 | +| `actor_rollout_ref.actor.ppo_epochs` | `1` | PPO 更新轮数 | +| `actor_rollout_ref.actor.shuffle` | `false` | 训练时是否对 mini batch 进行 shuffle | +| `actor_rollout_ref.actor.data_loader_seed` | `42` | 数据加载器随机种子 | +| `actor_rollout_ref.actor.grad_clip` | `1.0` | 梯度裁剪值 | +| `actor_rollout_ref.actor.ulysses_sequence_parallel_size` | `1` | Ulysses 序列并行大小 | +| `actor_rollout_ref.actor.entropy_from_logits_with_chunking` | `false` | 是否使用分块方式从 logits 计算熵 | +| `actor_rollout_ref.actor.entropy_checkpointing` | `false` | 是否对熵计算使用梯度检查点 | +| `actor_rollout_ref.actor.use_remove_padding` | 引用自 `model.use_remove_padding` | 是否移除 padding | +| `actor_rollout_ref.actor.calculate_sum_pi_squared` | `false` | 是否计算策略概率平方和 | +| `actor_rollout_ref.actor.sum_pi_squared_checkpointing` | `false` | 是否对策略概率平方和计算使用梯度检查点 | +| `actor_rollout_ref.actor.use_fused_kernels` | 引用自 `model.use_fused_kernels` | 是否使用融合内核 | + +#### 1.1.3 Policy Loss 配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `actor_rollout_ref.actor.policy_loss.loss_mode` | `vanilla` | 策略损失模式，可选 vanilla、clip_cov、kl_cov、dppo_tv、dppo_kl、gspo、sapo、geo_mean、cispo、gpg、bypass_mode、reinforce_is 等 | +| `actor_rollout_ref.actor.policy_loss.clip_cov_ratio` | `0.0002` | clip_cov 模式的协方差比率 | +| `actor_rollout_ref.actor.policy_loss.clip_cov_lb` | `1.0` | clip_cov 模式的协方差下界 | +| `actor_rollout_ref.actor.policy_loss.clip_cov_ub` | `5.0` | clip_cov 模式的协方差上界 | +| `actor_rollout_ref.actor.policy_loss.kl_cov_ratio` | `0.0002` | kl_cov 模式的协方差比率 | +| `actor_rollout_ref.actor.policy_loss.ppo_kl_coef` | `0.1` | PPO KL 散度系数 | + +#### 1.1.4 Rollout 配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `actor_rollout_ref.rollout.name` | `???` | Rollout 引擎名称，需用户指定 | +| `actor_rollout_ref.rollout.mode` | `async` | Rollout 模式，可选 async、sync 等 | +| `actor_rollout_ref.rollout.nnodes` | `0` | Rollout 使用的节点数 | +| `actor_rollout_ref.rollout.n_gpus_per_node` | 引用自 `trainer.n_gpus_per_node` | 每节点 GPU 数 | +| `actor_rollout_ref.rollout.temperature` | `1.0` | 采样温度，控制生成随机性 | +| `actor_rollout_ref.rollout.top_k` | `-1` | Top-K 采样参数，-1 表示不启用 | +| `actor_rollout_ref.rollout.top_p` | `1` | Top-P（nucleus）采样参数 | +| `actor_rollout_ref.rollout.prompt_length` | 引用自 `data.max_prompt_length` | Prompt 最大长度 | +| `actor_rollout_ref.rollout.response_length` | 引用自 `data.max_response_length` | Response 最大长度 | +| `actor_rollout_ref.rollout.dtype` | `bfloat16` | Rollout 推理数据类型 | +| `actor_rollout_ref.rollout.gpu_memory_utilization` | `0.5` | GPU 内存利用率，推理时使用 GPU 内存的比例 | +| `actor_rollout_ref.rollout.ignore_eos` | `false` | 是否忽略 EOS token | +| `actor_rollout_ref.rollout.enforce_eager` | `false` | 是否强制使用 PyTorch eager 模式 | +| `actor_rollout_ref.rollout.cudagraph_capture_sizes` | `null` | CUDA Graph 捕获大小列表 | +| `actor_rollout_ref.rollout.free_cache_engine` | `true` | 是否在每次推理后释放缓存引擎 | +| `actor_rollout_ref.rollout.tensor_model_parallel_size` | `2` | 推理时 TP 并行大小 | +| `actor_rollout_ref.rollout.data_parallel_size` | `1` | 推理时数据并行大小 | +| `actor_rollout_ref.rollout.expert_parallel_size` | `1` | 推理时专家并行大小 | +| `actor_rollout_ref.rollout.pipeline_model_parallel_size` | `1` | 推理时 PP 并行大小 | +| `actor_rollout_ref.rollout.max_num_batched_tokens` | `8192` | 单步最大批处理 token 数 | +| `actor_rollout_ref.rollout.max_model_len` | `null` | 模型最大序列长度，null 表示自动推断 | +| `actor_rollout_ref.rollout.max_num_seqs` | `1024` | 推理并发最大样本数 | +| `actor_rollout_ref.rollout.enable_chunked_prefill` | `true` | 是否启用分块预填充 | +| `actor_rollout_ref.rollout.enable_prefix_caching` | `true` | 是否启用前缀缓存（KV Cache 复用） | +| `actor_rollout_ref.rollout.logprobs_mode` | `processed_logprobs` | logprobs 计算模式 | +| `actor_rollout_ref.rollout.scheduling_policy` | `fcfs` | 调度策略，可选 fcfs 等 | +| `actor_rollout_ref.rollout.load_format` | `dummy` | 模型加载格式 | +| `actor_rollout_ref.rollout.log_prob_micro_batch_size` | `null` | log prob 计算的 micro batch 大小 | +| `actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu` | `null` | 每 GPU 的 log prob micro batch 大小 | +| `actor_rollout_ref.rollout.log_prob_use_dynamic_bsz` | 引用自 `actor.use_dynamic_bsz` | log prob 是否使用动态 batch size | +| `actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu` | 引用自 `actor.ppo_max_token_len_per_gpu` | 每 GPU 的 log prob 最大 token 长度 | +| `actor_rollout_ref.rollout.disable_log_stats` | `true` | 是否禁用推理日志统计 | +| `actor_rollout_ref.rollout.do_sample` | `true` | 是否进行采样（false 则为贪心解码） | +| `actor_rollout_ref.rollout.n` | `1` | 每个 prompt 生成的 response 数量 | +| `actor_rollout_ref.rollout.over_sample_rate` | `0` | 过采样率 | +| `actor_rollout_ref.rollout.multi_stage_wake_up` | `false` | 是否启用多阶段唤醒 | +| `actor_rollout_ref.rollout.calculate_log_probs` | `false` | 是否在 rollout 阶段计算 log probs | +| `actor_rollout_ref.rollout.skip_rollout` | `false` | 是否跳过 rollout | +| `actor_rollout_ref.rollout.skip_dump_dir` | `/tmp/rollout_dump` | 跳过 rollout 时的 dump 目录 | +| `actor_rollout_ref.rollout.skip_tokenizer_init` | `true` | 是否跳过分词器初始化 | +| `actor_rollout_ref.rollout.enable_rollout_routing_replay` | `false` | 是否启用 rollout 路由重放 | +| `actor_rollout_ref.rollout.quantization` | `null` | 量化方式 | +| `actor_rollout_ref.rollout.quantization_config_file` | `null` | 量化配置文件路径 | +| `actor_rollout_ref.rollout.layered_summon` | `false` | 是否启用分层召唤（仅 FSDP 方案） | + +#### 1.1.5 Rollout 验证采样配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `actor_rollout_ref.rollout.val_kwargs.top_k` | `-1` | 验证时 Top-K 采样参数 | +| `actor_rollout_ref.rollout.val_kwargs.top_p` | `1.0` | 验证时 Top-P 采样参数 | +| `actor_rollout_ref.rollout.val_kwargs.temperature` | `0` | 验证时采样温度，0 表示贪心解码 | +| `actor_rollout_ref.rollout.val_kwargs.n` | `1` | 验证时每个 prompt 生成的 response 数 | +| `actor_rollout_ref.rollout.val_kwargs.do_sample` | `false` | 验证时是否采样 | + +#### 1.1.6 Multi-Turn 配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `actor_rollout_ref.rollout.multi_turn.enable` | `false` | 是否启用多轮对话 | +| `actor_rollout_ref.rollout.multi_turn.max_assistant_turns` | `null` | 最大助手轮数 | +| `actor_rollout_ref.rollout.multi_turn.tool_config_path` | `null` | 工具配置文件路径 | +| `actor_rollout_ref.rollout.multi_turn.max_user_turns` | `null` | 最大用户轮数 | +| `actor_rollout_ref.rollout.multi_turn.max_parallel_calls` | `1` | 最大并行工具调用数 | +| `actor_rollout_ref.rollout.multi_turn.max_tool_response_length` | `256` | 工具响应最大长度 | +| `actor_rollout_ref.rollout.multi_turn.tool_response_truncate_side` | `middle` | 工具响应截断方向 | +| `actor_rollout_ref.rollout.multi_turn.interaction_config_path` | `null` | 交互配置文件路径 | +| `actor_rollout_ref.rollout.multi_turn.use_inference_chat_template` | `false` | 是否使用推理聊天模板 | +| `actor_rollout_ref.rollout.multi_turn.tokenization_sanity_check_mode` | `strict` | 分词完整性检查模式 | +| `actor_rollout_ref.rollout.multi_turn.format` | `hermes` | 多轮对话格式 | +| `actor_rollout_ref.rollout.multi_turn.num_repeat_rollouts` | `null` | 重复 rollout 次数 | + +#### 1.1.7 Agent 配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `actor_rollout_ref.rollout.agent.num_workers` | `8` | Agent 工作进程数 | +| `actor_rollout_ref.rollout.agent.default_agent_loop` | `single_turn_agent` | 默认 Agent 循环类型 | +| `actor_rollout_ref.rollout.agent.agent_loop_config_path` | `null` | Agent 循环配置文件路径 | +| `actor_rollout_ref.rollout.agent.custom_async_server.path` | `null` | 自定义异步服务路径 | +| `actor_rollout_ref.rollout.agent.custom_async_server.name` | `null` | 自定义异步服务名称 | + +#### 1.1.8 Checkpoint Engine 配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `actor_rollout_ref.rollout.checkpoint_engine.backend` | `naive` | Checkpoint 引擎后端 | +| `actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes` | `2048` | 权重更新桶大小（MB） | + +#### 1.1.9 Trace 配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `actor_rollout_ref.rollout.trace.project_name` | 引用自 `trainer.project_name` | 追踪项目名称 | +| `actor_rollout_ref.rollout.trace.experiment_name` | 引用自 `trainer.experiment_name` | 追踪实验名称 | +| `actor_rollout_ref.rollout.trace.backend` | `null` | 追踪后端 | +| `actor_rollout_ref.rollout.trace.token2text` | `false` | 是否将 token 转为文本 | +| `actor_rollout_ref.rollout.trace.max_samples_per_step_per_worker` | `null` | 每步每 worker 最大样本数 | + +#### 1.1.10 Prometheus 配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `actor_rollout_ref.rollout.prometheus.enable` | `false` | 是否启用 Prometheus 监控 | +| `actor_rollout_ref.rollout.prometheus.port` | `9090` | Prometheus 端口 | +| `actor_rollout_ref.rollout.prometheus.file` | `/tmp/ray/session_latest/metrics/prometheus/prometheus.yml` | Prometheus 配置文件路径 | +| `actor_rollout_ref.rollout.prometheus.served_model_name` | 引用自 `model.path` | 服务模型名称 | + +#### 1.1.11 Reference 模型配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `actor_rollout_ref.ref.rollout_n` | 引用自 `rollout.n` | Rollout 次数 | +| `actor_rollout_ref.ref.strategy` | 引用自 `actor.strategy` | 训练策略 | +| `actor_rollout_ref.ref.use_torch_compile` | 引用自 `actor.use_torch_compile` | 是否使用 torch.compile | +| `actor_rollout_ref.ref.log_prob_micro_batch_size` | `null` | log prob 计算的 micro batch 大小 | +| `actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu` | `null` | 每 GPU 的 log prob micro batch 大小 | +| `actor_rollout_ref.ref.log_prob_use_dynamic_bsz` | 引用自 `actor.use_dynamic_bsz` | log prob 是否使用动态 batch size | +| `actor_rollout_ref.ref.log_prob_max_token_len_per_gpu` | 引用自 `actor.ppo_max_token_len_per_gpu` | 每 GPU 的 log prob 最大 token 长度 | +| `actor_rollout_ref.ref.ulysses_sequence_parallel_size` | 引用自 `actor.ulysses_sequence_parallel_size` | Ulysses 序列并行大小 | +| `actor_rollout_ref.ref.entropy_from_logits_with_chunking` | `false` | 是否使用分块方式从 logits 计算熵 | +| `actor_rollout_ref.ref.entropy_checkpointing` | `false` | 是否对熵计算使用梯度检查点 | + +#### 1.1.12 Critic 优化器配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `critic.optim.lr` | `1.0e-05` | Critic 学习率 | +| `critic.optim.lr_warmup_steps_ratio` | `0.0` | 学习率预热步数比例 | +| `critic.optim.total_training_steps` | `-1` | 总训练步数 | +| `critic.optim.weight_decay` | `0.01` | 权重衰减 | +| `critic.optim.lr_warmup_steps` | `-1` | 学习率预热步数 | +| `critic.optim.betas` | `[0.9, 0.999]` | Adam 优化器动量系数 | +| `critic.optim.clip_grad` | `1.0` | 梯度裁剪阈值 | +| `critic.optim.override_optimizer_config` | `null` / `{}` | 覆盖优化器配置 | + +#### 1.1.13 Critic 策略配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `critic.strategy` | `fsdp` / `megatron` | 训练策略 | +| `critic.enable` | `null` | 是否启用 Critic，null 表示自动决定 | +| `critic.ppo_mini_batch_size` | 引用自 `actor.ppo_mini_batch_size` | PPO mini batch 大小 | +| `critic.ppo_micro_batch_size` | `null` | PPO micro batch 大小 | +| `critic.ppo_micro_batch_size_per_gpu` | `null` | 每 GPU 的 PPO micro batch 大小 | +| `critic.use_dynamic_bsz` | 引用自 `actor.use_dynamic_bsz` | 是否使用动态 batch size | +| `critic.ppo_max_token_len_per_gpu` | `32768` | 每 GPU 的 PPO 最大 token 长度 | +| `critic.forward_max_token_len_per_gpu` | 引用自 `.ppo_max_token_len_per_gpu` | 前向计算每 GPU 最大 token 长度 | +| `critic.ppo_epochs` | 引用自 `actor.ppo_epochs` | PPO 更新轮数 | +| `critic.shuffle` | 引用自 `actor.shuffle` | 是否 shuffle | +| `critic.data_loader_seed` | `42` / 引用自 `actor.data_loader_seed` | 数据加载器随机种子 | +| `critic.cliprange_value` | `0.5` | Critic 值函数裁剪范围，一般取值范围 [0.1, 0.3] | +| `critic.loss_agg_mode` | 引用自 `actor.loss_agg_mode` | 损失聚合模式 | +| `critic.grad_clip` | `1.0` | 梯度裁剪值 | +| `critic.ulysses_sequence_parallel_size` | `1` | Ulysses 序列并行大小 | +| `critic.forward_micro_batch_size` | 引用自 `.ppo_micro_batch_size` | 前向计算 micro batch 大小 | +| `critic.forward_micro_batch_size_per_gpu` | 引用自 `.ppo_micro_batch_size_per_gpu` | 前向计算每 GPU micro batch 大小 | + +#### 1.1.14 Critic 模型配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `critic.model.path` | `~/models/deepseek-llm-7b-chat` | Critic 模型路径 | +| `critic.model.tokenizer_path` | 引用自 `model.path` | 分词器路径 | +| `critic.model.override_config` | `{}` | 覆盖模型配置 | +| `critic.model.external_lib` | 引用自 `model.external_lib` | 外部库路径 | +| `critic.model.trust_remote_code` | 引用自 `model.trust_remote_code` | 是否信任远程代码 | +| `critic.model.use_shm` | `false` | 是否使用共享内存 | +| `critic.model.enable_gradient_checkpointing` | `true` | 是否启用梯度检查点 | +| `critic.model.enable_activation_offload` | `false` | 是否启用激活卸载 | +| `critic.model.use_remove_padding` | `false` / `true` | 是否移除 padding | +| `critic.model.lora_rank` | `0` | LoRA 秩 | +| `critic.model.lora_alpha` | `16` | LoRA alpha | +| `critic.model.target_modules` | `all-linear` | LoRA 目标模块 | +| `critic.model.tiled_mlp.enabled` | `false` | 是否启用分片 MLP | +| `critic.model.tiled_mlp.num_shards` | `4` | MLP 分片数 | + +#### 1.1.15 数据配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `data.tokenizer` | `null` | 分词器路径 | +| `data.use_shm` | `false` | 是否使用共享内存 | +| `data.train_files` | `~/data/rlhf/gsm8k/train.parquet` | 训练数据文件路径 | +| `data.val_files` | `~/data/rlhf/gsm8k/test.parquet` | 验证数据文件路径 | +| `data.train_max_samples` | `-1` | 训练最大样本数，-1 表示不限制 | +| `data.val_max_samples` | `-1` | 验证最大样本数 | +| `data.prompt_key` | `prompt` | 数据中 prompt 的键名 | +| `data.reward_fn_key` | `data_source` | 奖励函数的键名 | +| `data.max_prompt_length` | `512` | 最大 prompt 长度 | +| `data.max_response_length` | `512` | 最大 response 长度 | +| `data.train_batch_size` | `1024` | 训练 batch 大小 | +| `data.val_batch_size` | `null` | 验证 batch 大小 | +| `data.tool_config_path` | 引用自 `rollout.multi_turn.tool_config_path` | 工具配置文件路径 | +| `data.return_raw_input_ids` | `false` | 是否返回原始 input ids | +| `data.return_raw_chat` | `true` | 是否返回原始聊天内容 | +| `data.return_full_prompt` | `false` | 是否返回完整 prompt | +| `data.shuffle` | `true` | 是否 shuffle 训练数据 | +| `data.seed` | `null` | 数据 shuffle 随机种子 | +| `data.dataloader_num_workers` | `8` | 数据加载器工作进程数 | +| `data.image_patch_size` | `14` | 图像 patch 大小 | +| `data.validation_shuffle` | `false` | 验证时是否 shuffle | +| `data.filter_overlong_prompts` | `false` | 是否过滤超长 prompt | +| `data.filter_overlong_prompts_workers` | `1` | 过滤超长 prompt 的工作进程数 | +| `data.truncation` | `error` | 截断策略 | +| `data.image_key` | `images` | 图像数据的键名 | +| `data.video_key` | `videos` | 视频数据的键名 | +| `data.trust_remote_code` | `false` | 是否信任远程代码 | +| `data.return_multi_modal_inputs` | `true` | 是否返回多模态输入 | + +#### 1.1.16 奖励配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `reward.num_workers` | `8` | 奖励计算工作进程数 | +| `reward.custom_reward_function.path` | `null` | 自定义奖励函数路径 | +| `reward.custom_reward_function.name` | `compute_score` | 自定义奖励函数名称 | +| `reward.reward_manager.source` | `register` | 奖励管理器来源 | +| `reward.reward_manager.name` | `naive` | 奖励管理器名称 | +| `reward.reward_model.enable` | `false` | 是否启用奖励模型 | +| `reward.reward_model.enable_resource_pool` | `false` | 是否启用奖励模型资源池 | +| `reward.reward_model.n_gpus_per_node` | `8` | 奖励模型每节点 GPU 数 | +| `reward.reward_model.nnodes` | `0` | 奖励模型节点数 | +| `reward.reward_model.model_path` | `null` | 奖励模型路径 | +| `reward.sandbox_fusion.url` | `null` | Sandbox Fusion URL | +| `reward.sandbox_fusion.max_concurrent` | `64` | Sandbox Fusion 最大并发数 | +| `reward.sandbox_fusion.memory_limit_mb` | `1024` | Sandbox Fusion 内存限制（MB） | + +#### 1.1.17 算法配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `algorithm.gamma` | `1.0` | 折扣因子 | +| `algorithm.lam` | `1.0` | GAE lambda 参数 | +| `algorithm.adv_estimator` | `gae` | 优势估计方法，可选 gae 等 | +| `algorithm.norm_adv_by_std_in_grpo` | `true` | GRPO 中是否按标准差归一化优势 | +| `algorithm.use_kl_in_reward` | `false` | 是否在奖励中使用 KL 惩罚 | +| `algorithm.kl_penalty` | `kl` | KL 惩罚类型 | +| `algorithm.kl_ctrl.type` | `fixed` | KL 控制器类型，可选 fixed、kl_adapter 等 | +| `algorithm.kl_ctrl.kl_coef` | `0.001` | KL 惩罚系数 | +| `algorithm.kl_ctrl.horizon` | `10000` | KL 适配器的 horizon | +| `algorithm.kl_ctrl.target_kl` | `0.1` | 目标 KL 散度 | +| `algorithm.use_pf_ppo` | `false` | 是否使用 PF-PPO | +| `algorithm.pf_ppo.reweight_method` | `pow` | PF-PPO 重加权方法 | +| `algorithm.pf_ppo.weight_pow` | `2.0` | PF-PPO 加权幂次 | + +#### 1.1.18 Rollout Correction 配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `algorithm.rollout_correction.rollout_is` | `null` | 是否启用 IS 重要性采样校正 | +| `algorithm.rollout_correction.rollout_is_threshold` | `2.0` | IS 权重阈值 | +| `algorithm.rollout_correction.rollout_rs` | `null` | 是否启用拒绝采样校正 | +| `algorithm.rollout_correction.rollout_rs_threshold` | `null` | RS 阈值 | +| `algorithm.rollout_correction.bypass_mode` | `false` | 是否启用旁路模式 | +| `algorithm.rollout_correction.loss_type` | `ppo_clip` | 校正损失类型 | +| `algorithm.rollout_correction.rollout_is_batch_normalize` | `false` | IS 权重是否批量归一化 | + +#### 1.1.19 训练器配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `trainer.balance_batch` | `true` | 是否平衡 batch | +| `trainer.total_epochs` | `30` | 总训练 epoch 数 | +| `trainer.total_training_steps` | `null` | 总训练步数，null 表示由 epoch 自动计算 | +| `trainer.project_name` | `verl_examples` | 项目名称 | +| `trainer.experiment_name` | `gsm8k` | 实验名称 | +| `trainer.logger` | `[console, wandb]` | 日志后端列表 | +| `trainer.log_val_generations` | `0` | 验证生成日志数量 | +| `trainer.nnodes` | `1` | 训练节点数 | +| `trainer.n_gpus_per_node` | `8` | 每节点 GPU 数 | +| `trainer.save_freq` | `-1` | 保存频率，-1 表示不保存 | +| `trainer.esi_redundant_time` | `0` | ESI 冗余时间 | +| `trainer.resume_mode` | `auto` | 恢复模式，可选 auto 等 | +| `trainer.resume_from_path` | `null` | 恢复路径 | +| `trainer.val_before_train` | `true` | 训练前是否先验证 | +| `trainer.val_only` | `false` | 是否仅验证 | +| `trainer.test_freq` | `-1` | 测试频率 | +| `trainer.critic_warmup` | `0` | Critic 预热步数 | +| `trainer.default_hdfs_dir` | `null` | 默认 HDFS 目录 | +| `trainer.del_local_ckpt_after_load` | `false` | 加载后是否删除本地 checkpoint | +| `trainer.default_local_dir` | `checkpoints/${trainer.project_name}/${trainer.experiment_name}` | 默认本地 checkpoint 目录 | +| `trainer.max_actor_ckpt_to_keep` | `null` | 最多保留的 Actor checkpoint 数 | +| `trainer.max_critic_ckpt_to_keep` | `null` | 最多保留的 Critic checkpoint 数 | +| `trainer.ray_wait_register_center_timeout` | `300` | Ray 注册中心等待超时（秒） | +| `trainer.device` | `cuda` | 训练设备 | +| `trainer.use_legacy_worker_impl` | `auto` | 是否使用旧版 worker 实现 | +| `trainer.rollout_data_dir` | `null` | 保存每轮 rollout 结果的地址配置 | + +#### 1.1.20 模型配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `actor_rollout_ref.model.path` | `~/models/deepseek-llm-7b-chat` | 模型路径 | +| `actor_rollout_ref.model.hf_config_path` | `null` | HuggingFace 配置路径 | +| `actor_rollout_ref.model.tokenizer_path` | `null` | 分词器路径 | +| `actor_rollout_ref.model.use_shm` | `false` | 是否使用共享内存 | +| `actor_rollout_ref.model.trust_remote_code` | `false` | 是否信任远程代码 | +| `actor_rollout_ref.model.custom_chat_template` | `null` | 自定义聊天模板 | +| `actor_rollout_ref.model.external_lib` | `null` | 外部库路径 | +| `actor_rollout_ref.model.override_config` | `{}` | 覆盖模型配置 | +| `actor_rollout_ref.model.enable_gradient_checkpointing` | `true` | 是否启用梯度检查点 | +| `actor_rollout_ref.model.enable_activation_offload` | `false` | 是否启用激活卸载 | +| `actor_rollout_ref.model.use_remove_padding` | `true` / `false` | 是否移除 padding | +| `actor_rollout_ref.model.lora_rank` | `0` | LoRA 秩，0 表示不使用 LoRA | +| `actor_rollout_ref.model.lora_alpha` | `16` | LoRA alpha | +| `actor_rollout_ref.model.target_modules` | `all-linear` | LoRA 目标模块 | +| `actor_rollout_ref.model.exclude_modules` | `null` | LoRA 排除模块 | +| `actor_rollout_ref.model.lora_adapter_path` | `null` | LoRA 适配器路径 | +| `actor_rollout_ref.model.use_liger` | `false` | 是否使用 Liger 内核 | +| `actor_rollout_ref.model.use_fused_kernels` | `false` | 是否使用融合内核 | +| `actor_rollout_ref.model.fused_kernel_options.impl_backend` | `torch` | 融合内核实现后端 | +| `actor_rollout_ref.model.tiled_mlp.enabled` | `false` | 是否启用分片 MLP | +| `actor_rollout_ref.model.tiled_mlp.num_shards` | `4` | MLP 分片数 | + +#### 1.1.21 公共引擎配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `actor_rollout_ref.hybrid_engine` | `true` | 是否使用混合引擎（训练推理共享权重） | +| `actor_rollout_ref.nccl_timeout` | `600` | NCCL 通信超时（秒） | +| `transfer_queue.enable` | `false` | 是否启用传输队列 | + +--- + +### 1.2 FSDP 专属配置参数 + +以下参数仅在 FSDP 方案（`_generated_ppo_trainer.yaml`）中存在。 + +#### 1.2.1 FSDP 优化器配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `actor_rollout_ref.actor.optim.optimizer` | `AdamW` | 优化器类型 | +| `actor_rollout_ref.actor.optim.optimizer_impl` | `torch.optim` | 优化器实现 | +| `actor_rollout_ref.actor.optim.min_lr_ratio` | `0.0` | 最小学习率比例 | +| `actor_rollout_ref.actor.optim.num_cycles` | `0.5` | 余弦调度周期数 | +| `actor_rollout_ref.actor.optim.lr_scheduler_type` | `constant` | 学习率调度器类型 | +| `actor_rollout_ref.actor.optim.zero_indexed_step` | `true` | 步数是否从 0 开始计数 | +| `actor_rollout_ref.actor.optim.warmup_style` | `null` | 预热风格 | + +#### 1.2.2 Actor FSDP 引擎配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `actor_rollout_ref.actor.fsdp_config.wrap_policy.min_num_params` | `0` | FSDP 包装的最小参数数 | +| `actor_rollout_ref.actor.fsdp_config.param_offload` | `false` | 是否将参数卸载到 CPU | +| `actor_rollout_ref.actor.fsdp_config.optimizer_offload` | `false` | 是否将优化器状态卸载到 CPU | +| `actor_rollout_ref.actor.fsdp_config.offload_policy` | `false` | 卸载策略 | +| `actor_rollout_ref.actor.fsdp_config.reshard_after_forward` | `true` | 前向计算后是否重新分片 | +| `actor_rollout_ref.actor.fsdp_config.fsdp_size` | `-1` | FSDP 组大小，-1 表示全局 | +| `actor_rollout_ref.actor.fsdp_config.forward_prefetch` | `false` | 是否预取前向参数 | +| `actor_rollout_ref.actor.fsdp_config.model_dtype` | `fp32` | 模型计算数据类型 | +| `actor_rollout_ref.actor.fsdp_config.use_orig_params` | `false` | 是否使用原始参数 | +| `actor_rollout_ref.actor.fsdp_config.seed` | `42` | 随机种子 | +| `actor_rollout_ref.actor.fsdp_config.full_determinism` | `false` | 是否启用完全确定性 | +| `actor_rollout_ref.actor.fsdp_config.forward_only` | `false` | 是否仅前向计算（Actor 为 false） | +| `actor_rollout_ref.actor.fsdp_config.strategy` | `fsdp` | 策略类型 | +| `actor_rollout_ref.actor.fsdp_config.dtype` | `bfloat16` | 模型存储数据类型 | + +#### 1.2.3 Reference FSDP 引擎配置 + +与 Actor FSDP 引擎配置结构相同，主要区别： + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `actor_rollout_ref.ref.fsdp_config.forward_only` | `true` | Reference 模型仅前向计算 | + +其余参数（`wrap_policy`、`param_offload`、`optimizer_offload`、`reshard_after_forward`、`fsdp_size`、`dtype` 等）默认值与 Actor FSDP 引擎配置一致。 + +#### 1.2.4 Critic FSDP 引擎配置 + +与 Actor FSDP 引擎配置结构相同，主要区别： + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `critic.model.fsdp_config.forward_only` | `false` | Critic 模型需要训练 | +| `critic.model.fsdp_config.use_remove_padding` | `false` | Critic 不移除 padding | + +其余参数默认值与 Actor FSDP 引擎配置一致。 + +--- + +### 1.3 Megatron 专属配置参数 + +以下参数仅在 Megatron 方案（`_generated_ppo_megatron_trainer.yaml`）中存在。 + +#### 1.3.1 Megatron 优化器配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `actor_rollout_ref.actor.optim.optimizer` | `adam` | 优化器类型 | +| `actor_rollout_ref.actor.optim.lr_warmup_init` | `0.0` | 学习率预热初始值 | +| `actor_rollout_ref.actor.optim.lr_decay_steps` | `null` | 学习率衰减步数 | +| `actor_rollout_ref.actor.optim.lr_decay_style` | `constant` | 学习率衰减风格，可选 constant、cosine、exponential 等 | +| `actor_rollout_ref.actor.optim.min_lr` | `0.0` | 最小学习率 | +| `actor_rollout_ref.actor.optim.weight_decay_incr_style` | `constant` | 权重衰减增长风格 | +| `actor_rollout_ref.actor.optim.lr_wsd_decay_style` | `exponential` | WSD 学习率衰减风格 | +| `actor_rollout_ref.actor.optim.lr_wsd_decay_steps` | `null` | WSD 学习率衰减步数 | +| `actor_rollout_ref.actor.optim.use_checkpoint_opt_param_scheduler` | `false` | 是否使用 checkpoint 优化器参数调度器 | + +#### 1.3.2 Actor Megatron 引擎配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `actor_rollout_ref.actor.megatron.param_offload` | `false` | 是否将参数卸载到 CPU | +| `actor_rollout_ref.actor.megatron.grad_offload` | `false` | 是否将梯度卸载到 CPU | +| `actor_rollout_ref.actor.megatron.optimizer_offload` | `false` | 是否将优化器状态卸载到 CPU | +| `actor_rollout_ref.actor.megatron.tensor_model_parallel_size` | `1` | TP 并行大小 | +| `actor_rollout_ref.actor.megatron.expert_model_parallel_size` | `1` | 专家并行大小 | +| `actor_rollout_ref.actor.megatron.expert_tensor_parallel_size` | `null` | 专家 TP 并行大小 | +| `actor_rollout_ref.actor.megatron.pipeline_model_parallel_size` | `1` | PP 并行大小 | +| `actor_rollout_ref.actor.megatron.virtual_pipeline_model_parallel_size` | `null` | 虚拟 PP 并行大小 | +| `actor_rollout_ref.actor.megatron.context_parallel_size` | `1` | 上下文并行大小 | +| `actor_rollout_ref.actor.megatron.sequence_parallel` | `true` | 是否启用序列并行 | +| `actor_rollout_ref.actor.megatron.use_distributed_optimizer` | `true` | 是否使用分布式优化器 | +| `actor_rollout_ref.actor.megatron.use_dist_checkpointing` | `false` | 是否使用分布式 checkpoint | +| `actor_rollout_ref.actor.megatron.dist_checkpointing_path` | `null` | 分布式 checkpoint 路径 | +| `actor_rollout_ref.actor.megatron.dist_checkpointing_prefix` | `''` | 分布式 checkpoint 前缀 | +| `actor_rollout_ref.actor.megatron.dist_ckpt_optim_fully_reshardable` | `false` | 分布式 checkpoint 优化器是否完全可重分片 | +| `actor_rollout_ref.actor.megatron.distrib_optim_fully_reshardable_mem_efficient` | `false` | 分布式优化器重分片是否内存高效 | +| `actor_rollout_ref.actor.megatron.seed` | `42` | 随机种子 | +| `actor_rollout_ref.actor.megatron.use_mbridge` | `true` | 是否使用 mBridge | +| `actor_rollout_ref.actor.megatron.vanilla_mbridge` | `true` | 是否使用原始 mBridge | +| `actor_rollout_ref.actor.megatron.use_remove_padding` | `true` | 是否移除 padding | +| `actor_rollout_ref.actor.megatron.forward_only` | `false` | 是否仅前向计算 | +| `actor_rollout_ref.actor.megatron.dtype` | `bfloat16` | 模型数据类型 | +| `actor_rollout_ref.actor.megatron.load_weight` | `true` | 是否加载权重 | + +#### 1.3.3 Megatron Transformer 覆盖配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `override_transformer_config.recompute_granularity` | `null` | 重计算粒度 | +| `override_transformer_config.recompute_modules` | `[core_attn]` | 重计算模块列表 | +| `override_transformer_config.recompute_method` | `null` | 重计算方法 | +| `override_transformer_config.recompute_num_layers` | `null` | 重计算层数 | +| `override_transformer_config.attention_backend` | `flash` | 注意力后端 | + +#### 1.3.4 Reference Megatron 引擎配置 + +与 Actor Megatron 引擎配置结构相同，主要区别： + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `actor_rollout_ref.ref.megatron.forward_only` | `true` | Reference 模型仅前向计算 | + +其余参数默认值引用自 Actor Megatron 引擎配置（如 `param_offload`、`tensor_model_parallel_size` 等）。 + +#### 1.3.5 Critic Megatron 引擎配置 + +与 Actor Megatron 引擎配置结构相同，主要区别： + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `critic.megatron.forward_only` | `false` | Critic 模型需要训练 | + +#### 1.3.6 Megatron LoRA 配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `model.lora.type` | `lora` | LoRA 类型 | +| `model.lora.merge` | `false` | 是否合并 LoRA 权重 | +| `model.lora.rank` | `0` | LoRA 秩，0 表示不使用 | +| `model.lora.alpha` | `32` | LoRA alpha | +| `model.lora.dropout` | `0.0` | LoRA dropout | +| `model.lora.target_modules` | `[linear_qkv, linear_proj, linear_fc1, linear_fc2]` | LoRA 目标模块 | +| `model.lora.exclude_modules` | `[]` | LoRA 排除模块 | +| `model.lora.dropout_position` | `pre` | LoRA dropout 位置 | +| `model.lora.lora_A_init_method` | `xavier` | LoRA A 矩阵初始化方法 | +| `model.lora.lora_B_init_method` | `zero` | LoRA B 矩阵初始化方法 | +| `model.lora.a2a_experimental` | `false` | 是否启用 a2a 实验性功能 | +| `model.lora.dtype` | `null` | LoRA 数据类型 | +| `model.lora.adapter_path` | `null` | LoRA 适配器路径 | +| `model.lora.freeze_vision_model` | `true` | 是否冻结视觉模型 | +| `model.lora.freeze_vision_projection` | `true` | 是否冻结视觉投影 | +| `model.lora.freeze_language_model` | `true` | 是否冻结语言模型 | + +#### 1.3.7 模型 override_config（Megatron 方案） + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `model.override_config.model_config` | `{}` | 模型配置覆盖 | +| `model.override_config.moe_config.freeze_moe_router` | `false` | 是否冻结 MoE 路由 | + +#### 1.3.8 Rollout layer_name_map（Megatron 方案） + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `rollout.layer_name_map.qkv_layer_name` | `qkv` | QKV 层名称映射 | +| `rollout.layer_name_map.gate_proj_layer_name` | `gate_up` | Gate 投影层名称映射 | + +--- + +### 1.4 高级配置参数 + +#### 1.4.1 Profiler 配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `profiler.enable` | `false` | 是否启用 Profiler | +| `profiler.tool` | 引用自 `global_profiler.tool` | Profiler 工具，可选 nsys、npu、torch、torch_memory | +| `profiler.all_ranks` | `false` | 是否在所有 rank 上启用 | +| `profiler.ranks` | `[]` | 指定启用的 rank 列表 | +| `profiler.save_path` | 引用自 `global_profiler.save_path` | Profiler 结果保存路径 | + +#### 1.4.2 Global Profiler 配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `global_profiler.tool` | `null` | 全局 Profiler 工具 | +| `global_profiler.steps` | `null` | Profiler 采集步数 | +| `global_profiler.profile_continuous_steps` | `false` | 是否连续步采集 | +| `global_profiler.save_path` | `outputs/profile` | 全局保存路径 | + +#### 1.4.3 Router Replay 配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `router_replay.mode` | `disabled` | 路由重放模式，可选 disabled、record、replay | +| `router_replay.record_file` | `null` | 路由记录文件路径 | +| `router_replay.replay_file` | `null` | 路由重放文件路径 | + +#### 1.4.4 Checkpoint 配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `checkpoint.save_contents` | `[model, optimizer, extra]` | Checkpoint 保存内容 | +| `checkpoint.load_contents` | 引用自 `.save_contents` | Checkpoint 加载内容 | +| `checkpoint.async_save` | `false` | 是否异步保存 Checkpoint | +| `checkpoint.mbridge_config` | `{}` | mBridge 配置 | + +#### 1.4.5 QAT 配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `qat.enable` | `false` | 是否启用量化感知训练（QAT） | +| `qat.mode` | `w4a16` | 量化模式 | +| `qat.group_size` | `16` | 量化分组大小 | +| `qat.ignore_patterns` | `[lm_head, embed_tokens, re:.*mlp.gate$]` | 量化忽略的模式列表 | +| `qat.activation_observer` | `static_minmax` | 激活值观察器类型 | +| `qat.quantization_config_path` | `null` | 量化配置文件路径 | + +#### 1.4.6 MTP 配置 + +| 参数名 | 默认值 | 说明 | +|--------|--------|------| +| `mtp.enable` | `false` | 是否启用多 token 预测（MTP） | +| `mtp.enable_train` | `false` | 是否在训练中启用 MTP | +| `mtp.enable_rollout` | `false` | 是否在推理中启用 MTP | +| `mtp.detach_encoder` | `false` | 是否分离编码器 | +| `mtp.mtp_loss_scaling_factor` | `0.1` | MTP 损失缩放因子 | +| `mtp.speculative_algorithm` | `EAGLE` | 推测解码算法 | +| `mtp.speculative_num_steps` | `3` | 推测步数 | +| `mtp.speculative_eagle_topk` | `1` | EAGLE Top-K | +| `mtp.speculative_num_draft_tokens` | `4` | 推测 draft token 数 | +| `mtp.method` | `mtp` | MTP 方法 | +| `mtp.num_speculative_tokens` | `1` | 推测 token 数 | + +--- + +## 2. 训练指标说明 + +强化学习算法每个 iteration 打印的日志指标说明如下： + +### 2.1 训练基础指标 + +| 指标 | 说明 | +|------|------| +| `training/global_step` | 当前全局训练步数 | +| `training/epoch` | 当前训练 epoch | + +### 2.2 Actor 模型指标 + +| 指标 | 说明 | +|------|------| +| `actor/pg_loss` | 策略梯度损失（PPO clip loss），基于优势函数的策略梯度目标函数值 | +| `actor/kl_loss` | KL 散度损失，衡量当前策略与参考策略之间的偏离程度（仅 `use_kl_loss=True` 时打印） | +| `actor/entropy` | 策略熵，表示策略的随机性或探索能力（仅 `calculate_entropy=True` 或 `entropy_coeff!=0` 时打印） | +| `actor/grad_norm` | Actor 梯度范数（裁剪后），表示反向传播中参数梯度的整体幅度 | +| `actor/lr` | Actor 当前学习率 | +| `actor/pg_clipfrac` | PPO 裁剪机制生效的比例，反映策略更新幅度的稳定性 | +| `actor/ppo_kl` | PPO 算法的实际 KL 散度（当前策略 vs 旧策略） | +| `actor/pg_clipfrac_lower` | PPO 下界裁剪比例（部分 `loss_mode` 有此指标） | +| `actor/reward_kl_penalty` | KL 惩罚值，当前策略与参考策略的 KL 均值（仅 `use_kl_in_reward=True` 时打印） | +| `actor/reward_kl_penalty_coeff` | KL 惩罚系数 beta（仅 `use_kl_in_reward=True` 时打印） | +| `actor/kl_coef` | KL 损失系数（仅 `use_kl_loss=True` 时打印） | + +### 2.3 Critic 模型指标 + +| 指标 | 说明 | +|------|------| +| `critic/vf_loss` | 值函数损失 | +| `critic/vf_clipfrac` | Critic 裁剪机制生效的比例，反映值函数更新幅度的稳定性 | +| `critic/vpred_mean` | 预测值的均值 | +| `critic/grad_norm` | Critic 梯度范数（裁剪后） | +| `critic/lr` | Critic 当前学习率 | +| `critic/vf_explained_var` | 值函数解释方差 1 - Var(returns-values)/Var(returns)（仅 `use_critic=True` 时打印） | + +### 2.4 数据统计指标 + +| 指标 | 说明 | +|------|------| +| `critic/score/mean` | 非中止样本的序列分数均值 | +| `critic/score/max` | 非中止样本的序列分数最大值 | +| `critic/score/min` | 非中止样本的序列分数最小值 | +| `critic/rewards/mean` | 非中止样本的序列奖励均值 | +| `critic/rewards/max` | 非中止样本的序列奖励最大值 | +| `critic/rewards/min` | 非中止样本的序列奖励最小值 | +| `critic/advantages/mean` | 有效 token 的优势值均值 | +| `critic/advantages/max` | 有效 token 的优势值最大值 | +| `critic/advantages/min` | 有效 token 的优势值最小值 | +| `critic/returns/mean` | 有效 token 的回报均值 | +| `critic/returns/max` | 有效 token 的回报最大值 | +| `critic/returns/min` | 有效 token 的回报最小值 | +| `critic/values/mean` | 有效 token 的 Critic 值均值（仅 `use_critic=True` 时打印） | +| `critic/values/max` | 有效 token 的 Critic 值最大值（仅 `use_critic=True` 时打印） | +| `critic/values/min` | 有效 token 的 Critic 值最小值（仅 `use_critic=True` 时打印） | +| `response_length/mean` | 响应长度均值（含中止样本） | +| `response_length/max` | 响应长度最大值 | +| `response_length/min` | 响应长度最小值 | +| `response_length/clip_ratio` | 响应长度达到最大长度的比例 | +| `response_length_non_aborted/mean` | 非中止样本的响应长度均值 | +| `response_length_non_aborted/max` | 非中止样本的响应长度最大值 | +| `response_length_non_aborted/min` | 非中止样本的响应长度最小值 | +| `response_length_non_aborted/clip_ratio` | 非中止样本的响应长度达到最大长度的比例 | +| `response/aborted_ratio` | 中止样本（响应长度为 0）的比例 | +| `prompt_length/mean` | 提示长度均值 | +| `prompt_length/max` | 提示长度最大值 | +| `prompt_length/min` | 提示长度最小值 | +| `prompt_length/clip_ratio` | 提示长度达到最大长度的比例 | +| `num_turns/mean` | 多轮对话轮数均值（仅多轮对话时打印） | +| `num_turns/max` | 多轮对话轮数最大值（仅多轮对话时打印） | +| `num_turns/min` | 多轮对话轮数最小值（仅多轮对话时打印） | +| `tool_call_counts/mean` | 工具调用次数均值（仅存在 `tool_call_counts` 时打印） | +| `tool_call_counts/max` | 工具调用次数最大值 | +| `tool_call_counts/min` | 工具调用次数最小值 | + +### 2.5 时间指标 + +| 指标 | 说明 | +|------|------| +| `timing_s/gen` | 生成（rollout）耗时（秒） | +| `timing_s/ref` | Reference 模型计算 log_p 耗时（秒） | +| `timing_s/values` | Critic 模型计算 values 耗时（秒） | +| `timing_s/adv` | 计算优势值耗时（秒） | +| `timing_s/update_critic` | Critic 模型更新耗时（秒） | +| `timing_s/update_actor` | Actor 模型更新耗时（秒） | +| `timing_s/step` | 一步总耗时（秒） | +| `timing_s/old_log_prob` | Actor 模型计算旧 log_p 耗时（秒） | +| `timing_s/reward` | 奖励计算耗时（秒） | +| `timing_s/testing` | 验证耗时（秒） | +| `timing_s/save_checkpoint` | 保存 checkpoint 耗时（秒） | +| `timing_s/update_weights` | 权重同步耗时（秒） | +| `timing_per_token_ms/gen` | 生成阶段每 token 耗时（毫秒） | +| `timing_per_token_ms/ref` | Reference 模型每 token 耗时（毫秒） | +| `timing_per_token_ms/values` | Critic 模型每 token 耗时（毫秒） | +| `timing_per_token_ms/adv` | 优势值计算每 token 耗时（毫秒） | +| `timing_per_token_ms/update_critic` | Critic 更新每 token 耗时（毫秒） | +| `timing_per_token_ms/update_actor` | Actor 更新每 token 耗时（毫秒） | + +### 2.6 性能指标 + +| 指标 | 说明 | +|------|------| +| `perf/total_num_tokens` | 本步处理的总 token 数 | +| `perf/time_per_step` | 本步总耗时（秒） | +| `perf/throughput` | 吞吐量：tokens / (time * n_gpus) | +| `perf/max_memory_allocated_gb` | GPU 最大已分配内存（GB） | +| `perf/max_memory_reserved_gb` | GPU 最大预留内存（GB） | +| `perf/cpu_memory_used_gb` | CPU 已使用内存（GB） | +| `perf/mfu/actor` | Actor 训练的 MFU（模型浮点利用率） | +| `perf/mfu/critic` | Critic 训练的 MFU | +| `perf/mfu/actor_infer` | Actor 推理阶段的 MFU | + +### 2.7 方差代理指标 + +| 指标 | 说明 | +|------|------| +| `variance_proxy/proxy1_signal_strength` | 信号强度：梯度均值的平方范数 \|\|g_mean\|\|^2 | +| `variance_proxy/proxy2_total_power` | 总功率：梯度平方范数的期望 E[\|\|g_tau\|\|^2] | +| `variance_proxy/proxy3_pure_noise` | 纯噪声：梯度方差代理 (1/(N-1)) * (Proxy2 - Proxy1) | +| `variance_proxy/expected_a_squared` | 优势平方的期望 E[A^2] | +| `variance_proxy/expected_w` | W-score 代理的期望 E[W] | + +### 2.8 条件性指标 + +以下指标仅在特定条件满足时打印： + +#### 2.8.1 Rollout Correction 指标 + +仅启用 `rollout_correction` 时打印，均带 `rollout_corr/` 前缀。 + +**IS 权重指标**（仅启用 IS 校正时）： + +| 指标 | 说明 | +|------|------| +| `rollout_corr/rollout_is_mean` | IS 权重均值 | +| `rollout_corr/rollout_is_max` | IS 权重最大值 | +| `rollout_corr/rollout_is_min` | IS 权重最小值 | +| `rollout_corr/rollout_is_std` | IS 权重标准差 | +| `rollout_corr/rollout_is_ratio_fraction_high` | 超过上限阈值的 IS 权重比例 | +| `rollout_corr/rollout_is_ratio_fraction_low` | 低于下限阈值的 IS 权重比例 | +| `rollout_corr/rollout_is_eff_sample_size` | 有效样本大小（ESS） | +| `rollout_corr/rollout_is_seq_mean` | 序列级 IS 权重均值 | +| `rollout_corr/rollout_is_seq_std` | 序列级 IS 权重标准差 | +| `rollout_corr/rollout_is_seq_max` | 序列级 IS 权重最大值 | +| `rollout_corr/rollout_is_seq_min` | 序列级 IS 权重最小值 | +| `rollout_corr/rollout_is_seq_max_deviation` | 序列级 IS 权重与理想值 1.0 的最大偏差 | +| `rollout_corr/rollout_is_seq_fraction_high` | 序列级 IS 权重超过上限的比例 | +| `rollout_corr/rollout_is_seq_fraction_low` | 序列级 IS 权重低于下限的比例 | +| `rollout_corr/rollout_is_batch_norm_factor` | IS 权重批量归一化因子（仅 `rollout_is_batch_normalize=True` 时打印） | + +**Rejection Sampling 指标**（仅启用 RS 校正时）： + +| 指标 | 说明 | +|------|------| +| `rollout_corr/rollout_rs_{option}_mean` | RS 统计量均值 | +| `rollout_corr/rollout_rs_{option}_max` | RS 统计量最大值 | +| `rollout_corr/rollout_rs_{option}_min` | RS 统计量最小值 | +| `rollout_corr/rollout_rs_{option}_std` | RS 统计量标准差 | +| `rollout_corr/rollout_rs_{option}_fraction_high` | 超过上限阈值的比例 | +| `rollout_corr/rollout_rs_{option}_fraction_low` | 低于下限阈值的比例 | +| `rollout_corr/rollout_rs_{option}_seq_mean` | 序列级 RS 统计量均值 | +| `rollout_corr/rollout_rs_{option}_seq_std` | 序列级 RS 统计量标准差 | +| `rollout_corr/rollout_rs_{option}_seq_max` | 序列级 RS 统计量最大值 | +| `rollout_corr/rollout_rs_{option}_seq_min` | 序列级 RS 统计量最小值 | +| `rollout_corr/rollout_rs_{option}_seq_max_deviation` | 序列级 RS 统计量与 0 的最大偏差 | +| `rollout_corr/rollout_rs_{option}_seq_fraction_high` | 序列级超过上限的比例 | +| `rollout_corr/rollout_rs_{option}_seq_fraction_low` | 序列级低于下限的比例 | +| `rollout_corr/rollout_rs_{option}_masked_fraction` | token 级被 mask 掉的比例 | +| `rollout_corr/rollout_rs_{option}_seq_masked_fraction` | 序列级被 mask 掉的比例 | +| `rollout_corr/rollout_rs_masked_fraction` | 总体 token 级被 mask 掉的比例 | +| `rollout_corr/rollout_rs_seq_masked_fraction` | 总体序列级被 mask 掉的比例 | + +**Off-policy 诊断指标**（仅启用 off-policy 诊断时）： + +| 指标 | 说明 | +|------|------| +| `rollout_corr/training_ppl` | 训练策略的困惑度 | +| `rollout_corr/training_log_ppl` | 训练策略的 log 困惑度 | +| `rollout_corr/kl` | KL(π_rollout \|\| π_training) 直接估计 | +| `rollout_corr/k3_kl` | K3 KL 估计（更稳定） | +| `rollout_corr/rollout_ppl` | Rollout 策略的困惑度 | +| `rollout_corr/rollout_log_ppl` | Rollout 策略的 log 困惑度 | +| `rollout_corr/log_ppl_diff` | log PPL 差值（rollout - training） | +| `rollout_corr/log_ppl_abs_diff` | log PPL 绝对差值的均值 | +| `rollout_corr/log_ppl_diff_max` | log PPL 差值最大值 | +| `rollout_corr/log_ppl_diff_min` | log PPL 差值最小值 | +| `rollout_corr/ppl_ratio` | PPL 比率（training_ppl / rollout_ppl） | +| `rollout_corr/chi2_token` | token 级卡方散度 | +| `rollout_corr/chi2_seq` | 序列级卡方散度 | + +#### 2.8.2 序列长度平衡指标 + +仅启用 `balance_batch` 时打印： + +| 指标 | 说明 | +|------|------| +| `global_seqlen/min` | 平衡前各 DP 分区的最小序列长度和 | +| `global_seqlen/max` | 平衡前各 DP 分区的最大序列长度和 | +| `global_seqlen/minmax_diff` | 平衡前 max - min 差值 | +| `global_seqlen/balanced_min` | 平衡后各 DP 分区的最小序列长度和 | +| `global_seqlen/balanced_max` | 平衡后各 DP 分区的最大序列长度和 | +| `global_seqlen/mean` | 各分区的平均序列长度和 | + +#### 2.8.3 GDPO 奖励指标 + +仅使用 GDPO 估计器时打印： + +| 指标 | 说明 | +|------|------| +| `gdpo/{key}/mean` | GDPO 各奖励分量的均值 | +| `gdpo/{key}/std` | GDPO 各奖励分量的标准差 | +| `gdpo/{key}/max` | GDPO 各奖励分量的最大值 | +| `gdpo/{key}/min` | GDPO 各奖励分量的最小值 | + +#### 2.8.4 训推一致性指标 + +仅存在 `actor_rollout_ref.rollout.calculate_log_probs=True` 时打印： + +| 指标 | 说明 | +|------|------| +| `training/rollout_probs_diff_valid` | 标记为 1（有效） | +| `training/rollout_probs_diff_max` | rollout 与 actor 概率差异的最大值 | +| `training/rollout_probs_diff_mean` | rollout 与 actor 概率差异的均值 | +| `training/rollout_probs_diff_std` | rollout 与 actor 概率差异的标准差 | +| `training/rollout_actor_probs_pearson_corr` | rollout 与 actor 概率的 Pearson 相关系数 | + +#### 2.8.5 验证指标 + +验证阶段打印： + +| 指标 | 说明 | +|------|------| +| `val-core/{data_source}/{var_name}/{metric_name}` | 核心验证指标（mean@N, maj@N, best@N 等） | +| `val-aux/{data_source}/{var_name}/{metric_name}` | 辅助验证指标（std@N, worst@N 等） | +| `val-aux/num_turns/mean` | 验证集多轮对话轮数均值 | +| `val-aux/num_turns/max` | 验证集多轮对话轮数最大值 | +| `val-aux/num_turns/min` | 验证集多轮对话轮数最小值 | diff --git a/verl/docs/ascend_tutorial/dev_guide/model_dev/transfer_to_npu_guide.md b/verl/docs/ascend_tutorial/dev_guide/model_dev/transfer_to_npu_guide.md new file mode 100644 index 0000000000000000000000000000000000000000..235b0cef3137e9913813e5ae1bcf5acae5b583d9 --- /dev/null +++ b/verl/docs/ascend_tutorial/dev_guide/model_dev/transfer_to_npu_guide.md @@ -0,0 +1,296 @@ +# Transfer to NPU guide + +Last updated: 05/14/2026 + +本文为开发者提供从 GPU 迁移至 NPU或在 NPU 上独立适配模型的完整实践经验，涵盖前期准备、各组件打通、精度对齐、性能优化及长跑评测全流程。 + +## 一、前期准备 + +搭建可支持 NPU 运行的基础运行环境，保证模型正常加载、数据集可顺利读取，作为后续迁移调试、业务跑通的基础。 + + +### 1.1 软硬件环境与依赖配置 + +参照官方文档[install\_guidance.rst](https://github.com/verl-project/verl/blob/main/docs/ascend_tutorial/get_start/install_guidance.rst)；若模型依赖的推理引擎 vllm、vllm_ascend 和训练引擎Megatron、MindSpeed、transformers 版本与教程存在差异，**以模型实际适配版本为准**。 + +### 1.2 模型权重 + +BF16 为 VeRL 框架中 FSDP 与 Megatron 等训练后端**默认混合精度训练数据类型**。昇腾 NPU 环境统一采用 **BF16** 作为基准精度格式，权重需对齐反量化为 BF16。目前 A2、A3 机型**暂不支持 FP8 精度训练**，仅支持 BF16 精度；A5 机型后续版本将开放 FP8 低精度训练能力。 + +### 1.3 数据准备 + +数据需参照[Prepare Data for Post-Training](https://verl.readthedocs.io/en/latest/preparation/prepare_data.html)将数据集预处理为 parquet 格式：(1) 确保它包含计算强化学习奖励所需的必要字段；(2) 读取速度更快。 + +## 二、各组件联调打通 + +VeRL 框架采用推理引擎、训练引擎与权重同步桥接（Checkpoint Engine）相解耦的架构设计，可实现计算与数据的深度分离，为模型向昇腾 NPU 迁移适配提供了灵活的扩展基础。在开展模型在NPU上的迁移与适配工作时，建议优先完成推理引擎、训练引擎、Megatron-Bridge 各组件的单独适配与验证，待各组件运行稳定后，再推进 VeRL 整网链路的打通与调试。关于 VeRL 不同推理、训练后端在昇腾 NPU 上的具体特性支持，可参考[昇腾特性文档](https://github.com/verl-project/verl/blob/main/docs/ascend_tutorial/feature_support/ascend_backend_features.md)。 + +### 2.1 推理引擎适配 + +VeRL 推理引擎采用分层架构设计，通过抽象接口与工厂模式，实现 vllm、sglang 等多种主流推理后端的灵活支持。在完成 GPU 向 NPU 的迁移适配过程中，推理引擎适配推荐按以下流程操作： + +在 NPU 上跑通 VeRL 整网链路前，建议参考 [vllm-ascend](https://github.com/vllm-project/vllm-ascend/tree/main/docs/source/tutorials/models)、[sglang](https://github.com/sgl-project/sglang/blob/main/docs_new/docs/basic_usage) 官方模型部署教程，优先调通**单实例推理链路**，完整验证模型加载与初始化、Tokenizer 加载正常、单轮 / 批量生成、停止词终止、长上下文推理等**基础推理功能**，前置底层推理引擎稳定可用后，再接入 VeRL 训练流程。 + +### 2.2 训练引擎选择与适配 + +VeRL 主线代码将训练引擎抽象为 `Engine`类，通过标准化接口层实现调度逻辑与底层训练实现的解耦。该架构设计支持 FSDP、Megatron、MindSpeed-LLM 等多种训练后端灵活接入、即插即用，无需修改 VeRL 核心算法与调度逻辑，大幅降低迁移适配成本。 + +当前 NPU 已通过 `is_npu_available` 接口完成设备自动检测，并自动应用对应的 NPU 设备适配补丁。目前只需通过配置 model_engine=fsdp/megatron，即可一键切换训练后端至 FSDP、Megatron，系统会自动加载对应后端的 NPU 适配逻辑，无需额外修改代码。VeRL中昇腾对Megatron做了适配与优化，具体特性配置参考[verl-MindSpeed特性文档](https://gitcode.com/Ascend/MindSpeed/blob/master/docs/zh/user-guide/verl.md)设置。 + +### 2.3 Megatron-Bridge适配 + +Megatron-Bridge 主要用于在 VeRL 框架下，完成推理引擎依赖的 HuggingFace 权重与 Megatron-Core 所需 mcore 权重的双向转换，可通过以下配置启用该功能： + +``` +actor_rollout_ref.actor.megatron.use_mbridge=True +actor_rollout_ref.actor.megatron.vanilla_mbridge=False \ +``` + +Megatron-Bridge已在社区原生适配大量主流模型结构，支持列表可参考：[supported model](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/docs/models/README.md)，在昇腾 NPU 环境开展模型迁移适配时，可基于社区现有能力完成基础配置，但仍有部分模型特殊结构与场景需要补充定制化适配。 + +以​DSA （DeepSeek Sparse Attention）稀疏注意力结构为示例，介绍定制化适配的方法。昇腾 MindSpeed 支持基于吸收矩阵的 DSA能力，该特性要求将 Megatron 中原有的 `linear_kv_up_proj` 算子拆分为 `linear_k_up_proj` 与 `linear_v_up_proj` 两个独立算子。拆分所需权重需从 HuggingFace 格式的 `self_attn.kv_b_proj.weight` 转换生成，而上述原生 PR 并未适配该算子拆分逻辑。 + +因此需手动改造适配相关权重转换逻辑，保障吸收矩阵可正常加载与生效。只有在吸收矩阵可用的基础上，才能正常使能 [sparse\_flash\_attention](https://gitcode.com/cann/ops-transformer/tree/master/attention/sparse_flash_attention) 与 [lightning\_indexer](https://gitcode.com/cann/ops-transformer/tree/master/attention/lightning_indexer) 融合算子；通过引入两个融合算子，可大幅减少内存访问频次、优化内存占用率，同时提升计算性能，最终实现大模型训练与推理链路的运行效率提升及资源开销降低。 + +### 2.4 整网功能打通 + +完成推理引擎适配验证、训练引擎适配开发，参照[参数配置说明](https://github.com/verl-project/verl/blob/main/docs/ascend_tutorial/dev_guide/model_dev/parameter_and_metrics.md)根据实际业务需求，配置推理引擎、训练引擎的相关参数，完成 VeRL 整网功能打通，确保全流程稳定运行。 + +## 三、精度对齐 + +大模型强化学习的精度问题定位链路复杂、影响因素繁多，各类精度问题通常由**训练阶段、推理阶段、训推一致性**问题引入。**精度对齐**是保障训练流程可复现、问题可调试的核心关键。 + +训练与推理阶段的精度对齐，可参考官方文档：[精度对齐文档](https://github.com/verl-project/verl/blob/main/docs/ascend_tutorial/dev_guide/precision_analysis/precision_alignment_zh.md)。因此本节不再赘述基础阶段对齐流程，将**重点围绕训推一致性场景**，基于 msprobe 精度工具，开展精度对齐落地实践与问题定位排查工作。 + +### 3.1 精度监控配置 + +整网跑通后，需启用精度监控参数，设置 `actor_rollout_ref.rollout.calculate_log_probs=True`。在训练过程中需重点观察以下关键指标，以此判断训推一致性及模型训练稳定性： + +* **训推一致性参考指标**： + * `training/rollout_probs_diff_mean`（rollout概率差异均值），模型正常收敛状态下，该指标建议维持在 0.01 以内；若数值持续高于 0.01 或与 GPU 基准存在明显偏离，可判定存在训推精度异常。 + * `training/rollout_probs_diff_max`（rollout概率差异最大值） + * `training/rollout_actor_probs_pearson_corr`（rollout与actor概率的皮尔逊相关系数） +* **模型训练稳定性指标**： + * `actor/grad_norm`：需关注其是否呈整体下降趋势，以此判断模型训练是否正常收敛。 + +此外，配置参数 `trainer.rollout_data_dir=./rollout_dump/` 用于保存训练过程中的 Rollout 中间结果。通过人工核查导出的 Rollout 数据，校验模型回复是否符合预期、输出有无乱码与重复回答现象，可进一步从表象上确认推理引擎适配无误。 + +### 3.2 采集精度数据 + +当 training/rollout_probs_diff_mean 超出 0.01 合理阈值、或与 GPU 基准标杆出现明显偏离时，需进一步通过[msprobe](https://github.com/verl-project/verl/blob/main/docs/ascend_tutorial/dev_guide/precision_analysis/precision_debugger_zh.md)精度工具采集数据做根因定位。 + +### 3.3 训推差异点排查与对齐实践 + +完成数据采集后，优先读取 `construction.json` 文件进行模块级数据比对。先保证 `layer.0.input_layernorm` 输入数据完全一致，再逐模块逐层校验，定位训练与推理输出首次出现不一致的位置。 + +而对于大尺寸模型，微小数值差异会随着逐层累积、放大，导致训练（training）和推理（rollout）的结果差异明显，甚至会出现同一个token训推输出概率分别为0和1的现象，因此需尽可能将每一处差异点对齐至完全相等。 + +定位到差异节点后，适配修改方案同样是关键难点。由于业内各开源社区对相关模块存在多套不同实现，为保障模型实现逻辑的正确性，需多方参考权威源码与技术报告，综合确定最终对齐方案。 + +#### 3.3.1 常见训推不一致 + +在大模型强化学习实践中，可将训推不一致的典型根因归纳为以下五类： + +1. **框架实现不一致**：由于训练、推理框架实现逻辑不同导致。有时是“语义正确”的（如算子拆分方式不同但数学等价），有时是“语义错误”的（如遗漏了某个缩放因子，或多了某个操作），需结合源码与技术报告严格鉴定。 +2. **精度类型差异**：如训练侧全程 BF16，而推理侧在某些归一化等敏感算子中隐式升精到 FP32 计算再降精，导致截断误差。 +3. **超参数不一致**：如LayerNorm模块中硬编码的 `eps` 值未做统一。 +4. **并行策略**：训练时的张量并行 vs 推理时的连续批处理，导致浮点累加顺序差异。 +5. **随机性控制**：Dropout、采样策略在训推阶段的实现偏差。 + +下面列举 GLM-5 模型迁移适配过程中遇到的典型训推不一致实际案例。 + +#### 3.3.2 案例一：FFN激活函数的框架实现不一致 + +从上往下依次比对，排查到第一层的 MLP 激活函数处输出不一致。 + +推理侧已正常使用 NPU 优化的 `npu_swiglu` 融合算子，但训练侧仍执行原生 GLU 小算子实现。 + +* **根因**：尽管已在 Verl 参数中添加了 `swiglu` 使能配置，但 Megatron-Bridge 在 NPU 适配 PR 中，未显式配置 `provider.bias_activation_fusion=True`，导致代码未进入 NPU 融合算子分支。 + ``` + +actor_rollout_ref.actor.megatron.override_transformer_config.swiglu=True \ + +actor_rollout_ref.actor.megatron.override_transformer_config.use_fused_swiglu=True \ + ``` +* **修复方案**：在 Megatron-Bridge 中添加配置项使训练侧正确调用融合算子： + +#### 3.3.3 案例二：indexer_k_norm 的精度与超参数不一致 + +在严格对齐过程中，发现 `indexer_k_norm` 处存在精度类型与超参不一致： + +* **精度差异**：推理侧在 LayerNorm 中存在升精度到 fp32 操作 `F.layer_norm( x.float(), (self.dim,), self.weight, self.bias, self.eps).type_as(x)`，而训练侧 Megatron 实现为 BF16。微小差异经多层累积不可忽视。 +* **修复方案**：统一训练侧代码增加升精降精操作。 +* **超参差异**：GLM5 推理侧vllm继承 DeepSeekV32 逻辑，`k_norm` 的 EPS 值被硬编码为 `1e-6`；而训练引擎及官方技术报告统一采用 `1e-5`。 +* **修复方案**：将推理侧 EPS 修改为 `1e-5` 与训练侧对齐。 + +``` +self.k_norm=LayerNorm(self.head_dim,eps=1e-6 -> 1e-5) +``` + +#### 3.3.4 案例三：lightning_indexer 模块逻辑缺失与冗余 + +排查发现lightning_indexer模块，训练与推理侧该模块具体实现存在不一致，具体表现为： + +* **缺失（推理侧遗漏）**：推理侧缺失了 `weights` 的缩放逻辑。参考 Megatron 训练侧、slime 及 transformers 的标准实现，均包含该缩放，故在推理侧补齐以对齐前向： + +``` +weights, _ = self.weights_proj(x) ++weights = weights * (self.n_head**-0.5) * (self.head_dim**-0.5) +``` + +* **冗余（训练侧多余错误实现）**：训练侧 Megatron 实现中多出了 `rotate_activation`（哈达玛变换）。经查阅大量资料确认，该操作专用于量化场景，在 BF16 格式中属于错误实现。参考 [Transformer PR#45017](https://github.com/huggingface/transformers/pull/45017)，将训练侧该冗余逻辑移除。 + +``` +class DSAIndexer(MegatronModule): + def forward_with_scores( +- q = rotate_activation(q) +- k = rotate_activation(k) +``` + +### 3.4 MoE 大模型通用路由稳定方案 + +在典型的 RL 训练流程中，通常采用 vLLM 等高效推理引擎完成样本采样，再将采样数据送入 Megatron 等训练框架做模型训练优化。 + +对于常规稠密模型，推理与训练框架间的实现、环境差异仅会产生轻微数值偏差；但**大尺寸 MoE 模型**下该问题会被急剧放大。核心根源在于 MoE 动态路由机制：微小的框架实现、运行环境差异，就可能导致同一输入 Token 被分配至完全不同的专家组合，从而走向截然不同的计算路径。 + +这种路由决策的不一致，可能会导致MoE模型RL训练不稳定。它使得从推理阶段获取的“经验”对于训练阶段而言变得完全不同，优化信号因此失真，最终导致灾难性的后果。 + +为了解决这一通用问题，业界引入了 **Routing Replay（路由回放）** 机制。其核心思想是通过锁定特定阶段的专家路由路径，屏蔽微小扰动对路由决策的干扰，从而保证模型训练的稳定性。目前主流包含R2和R3两种变体： + +* **（1）Vanilla Routing Replay (R2)**： (对应`actor_rollout_ref.actor.router_replay.mode="R2"`) + + * **机制**：在梯度更新阶段，复现训练引擎在上一轮采样阶段计算出的专家路径。 + * **作用**：主要缓解**策略陈旧性**对路由的影响。随着策略的更新，当前前向传播计算出的路由可能与生成旧数据时的路由不一致，R2通过回放旧路由来维持优化信号的连贯性。 +* **（2）Rollout Routing Replay (R3)**：(对应`actor_rollout_ref.actor.megatron.router_replay.mode="R3"`) + + * **机制**：在序列生成过程中捕捉推理引擎的路由分布，并将其直接重放到训练引擎中。 + * **作用**：同时解决**训练-推理偏差**和**策略陈旧性**两个问题。它确保了训练阶段计算Loss所依据的专家路径，与实际推理生成结果时的专家路径绝对一致。 + +因此，无论是侧重缓解策略陈旧的 R2，还是实现全链路对齐的 R3，Routing Replay 机制都有效弥合了推理与训练框架间的路由鸿沟。在**大尺寸 MoE 模型**的训推一致性对齐中，该机制已成为保障精度对齐与训练稳定的核心手段。目前，DeepSeek-V3.2、GLM-5、MiMo-V2 等主流大模型均已采用了 R3 模式的 Routing Replay 技术。 + +因此对于大尺寸 MoE 模型，在实际配置中通常推荐使用对齐更彻底的 R3 模式： + +``` +actor_rollout_ref.actor.router_replay.mode="R3" \ +actor_rollout_ref.rollout.enable_rollout_routing_replay=True \ +``` + +## 四、性能优化 + +在昇腾 NPU 上进行大模型 RL（强化学习）训练性能优化时，基础配置调优可优先参考官方文档：[perf_tuning.rst](https://github.com/verl-project/verl/blob/04833f01/docs/perf/perf_tuning.rst)。为实现更高效的优化，建议遵循**数据采集​​→​瓶颈定位​→配置调优→迭代验证**的标准化流程，该流程可显著提升 Rollout、Reward、Update 等核心阶段的吞吐量，同时有效降低资源空泡与负载不均问题。性能分析与调优的具体操作，可严格参照以下官方指引： + +1. Ascend Performance Analysis Guide：[https://github.com/verl-project/verl/blob/main/docs/ascend\_tutorial/examples/ascend\_performance\_analysis\_guide.md](https://github.com/verl-project/verl/blob/main/docs/ascend_tutorial/examples/ascend_performance_analysis_guide.md) +2. Profiling 数据采集与使能配置：[https://github.com/verl-project/verl/blob/main/docs/ascend\_tutorial/profiling/ascend\_profiling\_zh.rst](https://github.com/verl-project/verl/blob/main/docs/ascend_tutorial/profiling/ascend_profiling_zh.rst) + +### 4.1 推理性能优化 + +Rollout 阶段作为大模型 RL 训练的核心推理环节，其推理耗时在整个训练流程中占据绝大部分，以下是该阶段常见的性能优化手段： + +1. **启用图模式功能**：图模式将整个计算图提前编译优化，可以实现算子融合、内存复用、常量折叠等深度优化，显著提升执行效率。 +2. **CPU 绑核加速算子下发**：通过 CPU 绑核可提升算子下发效率；自 vllm-ascend v0.18.0rc1 版本起，ARM 架构昇腾服务器已默认开启该能力。 +3. **HCCL 通信算法配置为 AIV 模式**：将环境变量 `HCCL_OP_EXPANSION_MODE` 设置为 `AIV` 模式，指定通信算法的编排与展开逻辑运行在 Device 侧 Vector Core 计算单元。 +4. **启用异步调度**：能够消除 Worker 连续两次 execute_model 执行间隙，让 Worker 可直接获取已调度完成的 SchedulerOutput 进行模型推理，无需阻塞等待调。 + +对应配置参数如下： + +``` +# 图模式启用 +actor_rollout_ref.rollout.enforce_eager=False +actor_rollout_ref.rollout.engine_kwargs.vllm.compilation_config.cudagraph_mode="FULL_DECODE_ONLY" ++actor_rollout_ref.rollout.engine_k +wargs.vllm.compilation_config.cudagraph_capture_sizes="[2, 4, 8, 16, 24, 32]" +# CPU绑核 +++actor_rollout_ref.rollout.engine_kwargs.vllm.additional_config.enable_cpu_binding=True +# 异步调度启用 +++actor_rollout_ref.rollout.engine_kwargs.vllm.async_scheduling=True +``` + +### 4.2训练性能优化： + +大模型 RL 训练的 Update 阶段具有序列长度差异大、显存消耗高等特点。除了基础的算子融合，还需结合序列并行与显存-计算权衡策略来打破瓶颈。常见训练性能优化特性可参考 [MindSpeed-verl 文档](https://gitcode.com/Ascend/MindSpeed/blob/master/docs/zh/user-guide/verl.md) 完成启用，核心优化手段包括： + +1. ​**算子融合**​：启用 RoPE、SwiGLU、RMSNorm、DSA 等融合算子。通过算子融合减少计算开销与显存，提升训练效率。 +2. ​**Remove padding**​：RL 训练中各 Response 长度参差不齐，传统 Padding 策略会导致大量无效计算。开启 Remove padding后可将多个短序列打包填满 Tensor，极大提升 NPU 计算单元的利用率（MFU）。 + +## 五、评测验证 + +训练完成后，需对目标数据集进行评测验证，确保模型迁移后的业务效果达标。不同模型的评测步骤一致，以下以 GLM-5 为例，详细说明评测流程（采用 AISBenchmark 工具，支持 vllm/sglang 多种推理后端的评估）。 + +评测采用了数学类的数据集aime2025与研究生级专业理科数据集gpqa，验证在目标方向上分数上升，且无关方向不会出现知识灾难遗忘情况。 + +### 5.1 安装aisbench + +```shell +git clone https://gitee.com/aisbench/benchmark.git +cd benchmark +pip install -e . +``` + +### 5.2 下载评估数据集 + +```shell +# linux服务器内，处于工具根路径下 +cd path/to/benchmark/ais_bench/datasets +wget http://opencompass.oss-cn-shanghai.aliyuncs.com/datasets/data/aime2025.zip +unzip aime2025.zip +rm aime2025.zip +``` + +### 5.3 修改AISBench配置代码使能vllm/sglang推理评测 + +打开 benchmark/ais_bench/benchmark/configs/models/vllm_api/vllm_api_stream_chat.py 文件，这是推理评测配置文件，输出长度`max_out_len`建议与训练的`max_response_len`保持一致。 + +```shell +from ais_bench.benchmark.models import VLLMCustomAPIChat +from ais_bench.benchmark.utils.postprocess.model_postprocessors import extract_non_reasoning_content + +models = [ + dict( + attr="service", + type=VLLMCustomAPIChat, + abbr='vllm-api-general-chat', + path="/path/to/GLM-5", # 修改为 GLM-5 模型路径 + model="GLM-5", + stream=True, + request_rate = 0, + use_timestamp=False, + max_seq_len=2048, + retry = 2, + api_key="", + host_ip = "localhost", # 推理服务的IP + host_port = 12890 , # 推理服务的端口 + max_out_len = 8192, # 最大输出tokens长度 + batch_size=48, # 推理的最大并发数 + trust_remote_code=False, + generation_kwargs = dict( + temperature = 0, + seed = 1234, + ), + pred_postprocessor=dict(type=extract_non_reasoning_content) + ) +] +``` + +### 5.4 多机拉起推理服务端 + +参考[vllm_ascend GLM5指南](https://github.com/vllm-project/vllm-ascend/blob/main/docs/source/tutorials/models/GLM5.md#multi-node-deployment)拉起双机A3推理服务，`host_port`与上一小节配置保持一致，`max_model_len`设置为训练时的`max_prompt_length`与`max_response`之和。 + +### 5.5 启动vllm评测任务 + +执行以下命令启动在线推理评测任务，调用已部署的 vLLM 推理后端，加载对应模型配置完成自动化评测： + +``` +ais_bench --models vllm_api_stream_chat --datasets aime2025_gen_0_shot_chat_prompt +``` + +模型经过训练后，核心能力指标实现稳定提升：在AIME2025 数学推理数据集上评测得分稳步上涨，同时在GPQA 研究生级专业理科数据集上也实现了持续的分数增益，无知识退化、无灾难性遗忘问题，训练优化效果符合预期。 + +| 评测数据集 | GLM5-base | 10step | 15step | 40step | 50step | +| ---------- | --------- | ------ | ------ | ------ | ------ | +| aime2025 | 47.5 | 49.17 | 49.17 | 48.33 | 52.5 | +| gpqa | 64.65 | 68.81 | 68.43 | 69.07 | 71.21 | + +## 六、总结 + +本文完整覆盖了大模型从 GPU 迁移至昇腾 NPU 或在 NPU 上独立适配的全流程实践，主要分为环境搭建、组件联调、精度对齐、性能优化、评测验证五大关键环节，为开发者提供可落地、可复用的操作指南与问题解决方案。 + +前期准备阶段需重把控环境依赖版本、模型权重精度与数据集格式，为后续适配奠定基础；组件联调环节需遵循先单组件验证后整网打通的原则，优先确保推理、训练引擎及权重转换工具的稳定适配，针对特殊模型结构需完成定制化改造；精度对齐是迁移适配的核心，需重点监控训推一致性指标，通过逐模块排查解决框架实现、精度类型等常见差异，MoE 模型需启用 Routing Replay 机制保障训练稳定；性能优化需遵循标准化流程，聚焦推理与训练核心阶段，通过图模式、算子融合等手段提升效率、降低资源消耗；最终通过标准化评测验证，确保模型迁移后业务效果达标、无知识退化。 + +整体而言，遵循本文流程可有效降低 NPU 迁移适配成本，规避常见坑点，实现大模型在昇腾 NPU 上的稳定、高效运行。 diff --git a/verl/docs/ascend_tutorial/dev_guide/performance/ascend_performance_analysis_guide.md b/verl/docs/ascend_tutorial/dev_guide/performance/ascend_performance_analysis_guide.md new file mode 100644 index 0000000000000000000000000000000000000000..e3ea1870de1bd6507823053a47b8c751a05ef65d --- /dev/null +++ b/verl/docs/ascend_tutorial/dev_guide/performance/ascend_performance_analysis_guide.md @@ -0,0 +1,153 @@ +# Ascend Performance Analysis Guide + +Last updated: 02/24/2026. + +## 背景介绍 + +随着DeepSeek-R1的发布，大模型强化学习（RL）训练受到广泛关注。在昇腾NPU环境下，verl框架已积累了丰富的性能调优经验。本文系统总结了包括性能数据采集与分析在内的方法论，旨在帮助开发者更高效地运用MindStudio工具链，实现强化学习场景下的性能优化。 + +### 强化学习计算流程概述 + +1. **Rollout**：策略（actor）模型基于输入的prompt序列，推理生成回答（response序列） +2. **ref logprob**：基于prompt和生成的response，reference模型计算ref logprob用于KL散度计算 +3. **logprob**：基于prompt和生成的response，actor模型计算logprob用于重要性采样 +4. **reward**：基于prompt和生成的response，奖励模型评估奖励值R_N。 +5. **update**：基于计算得到的R_N、ref logprob、logprob计算优化函数和策略梯度，对actor模型进行更新 + +![rl_data_stream](https://github.com/chengminhua/verl_data/raw/main/MindStudio_Insight_use/rl_data_stream.png) + +## profiling工具使能 + +### 使能方法 + +使能和配置教程可参考：[verl/docs/ascend_tutorial/profiling/ascend_profiling_zh.rst at main · verl-project/verl](https://github.com/verl-project/verl/blob/main/docs/ascend_tutorial/profiling/ascend_profiling_zh.rst) + +## 性能分析方法论 + +### 整体性能概览分析 + +#### 1. 长耗时任务与资源空泡分析 + +- **操作**：使用MindStudio Insight加载profiling数据，自动识别不同计算阶段，通过RL页签流水图定位长耗时任务与NPU资源空泡 +- **价值**：快速掌握不同阶段耗时占比 +- **效果展示**： + +![Bubble_analysis](https://github.com/chengminhua/verl_data/raw/main/MindStudio_Insight_use/Bubble_analysis.png) + +#### 2. 负载均衡分析 + +- **操作**：通过MindStudio Insight直接查看MSTX打点数据，观察Rollout阶段不同DP Rank的负载均衡情况 +- **价值**：快速识别负载不均问题 +- **效果展示：** + +![Load_Balancing_Analysis](https://github.com/chengminhua/verl_data/raw/main/MindStudio_Insight_use/Load_Balancing_Analysis.gif) + +#### 3. 集群整体性能分析 + +- **操作**：结合MSTT的rl_analysis功能，生成集群Timeline缩略图，观察各阶段整体耗时 +- **价值**：宏观掌握集群性能瓶颈 +- **操作指南**：[rl_analysis使用文档](https://gitcode.com/Ascend/mstt/raw/pre-research/profiler/msprof_analyze/docs/features/rl_analysis.md) +- **效果展示**： + +![Cluster%20Performance%20Analysis](https://github.com/chengminhua/verl_data/raw/main/MindStudio_Insight_use/Cluster%20Performance%20Analysis.png) + +### 细粒度分析 + +#### 性能分析 + +- **操作**：可通过 MindStudio Insight Windows 或 Linux 版本加载 Profiling 数据 +- **价值**：MindStudio Insight 支持分析任务调度效率、算子执行性能、计算资源利用率、集合通信性能等。其 Timeline 视图具备任务拆解与 Overlap 分析功能（**为 MindStudio 独有核心特性，在 NV 及其他竞品中不具备，是 AI 调优的必备工具**），并支持鼠标交互式分析。 +- **效果展示**： + +![performance%20analysis](https://github.com/chengminhua/verl_data/raw/main/MindStudio_Insight_use/performance%20analysis.png) + +#### 内存分析 + +##### **通过 Profiling 结合调用栈分析系统内存变化** + +- **操作**：采集数据时开启调用栈和内存视图功能。 +- **价值**：观察框架、CANN内存申请释放情况，可结合调用栈跟踪到前端python代码。 +- **效果展示**：结合调用栈进行内存变化分析。效果如下所示： + +![in-memory%20analytics](https://github.com/chengminhua/verl_data/raw/main/MindStudio_Insight_use/in-memory%20analytics.gif) + +##### **使用 msleaks 工具进行深层次内存分析** + +- **操作步骤**：参考 [msleaks 工具使用指南](https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/83RC1alpha003/devaids/msleaks/atlas_msleaks_0001.html)。 +- **价值**：可以查看框架内存申请总量折线图/内存块图，并直接对应调用栈，可深层次分析框架内存使用情况。 +- **效果展示**： + +![msleaks](https://github.com/chengminhua/verl_data/raw/main/MindStudio_Insight_use/msleaks.gif) + +## 性能分析案例 + +要做具体的性能分析，profiling要开启**level1**，否则算子的关键信息会缺失。 + +### 1.host bound诊断 + +host bound是指CPU任务量综合大于NPU，导致NPU执行出现空泡的现象。可以通过看Host2Device的同步连线来判断，如果连线都是歪的，那证明这里的set信号早于wait信号，NPU一ready就执行了，那也是device bound： + +![host_bound_1](https://github.com/chengminhua/verl_data/raw/main/MindStudio_Insight_use/host_bound_1.png) + +如果确诊为host bound，那么我们可以打开CPU侧，找出各算子的下发耗时。注意找的时候需要找出所有CPU耗时的累加值，而不能找单层，因为首次调用的耗时是很长的。例如下图的GmmSwigluQuant，CPU上首次调用需要1ms，后续每次只需要200us。 + +![host_bound_2](https://github.com/chengminhua/verl_data/raw/main/MindStudio_Insight_use/host_bound_2.png) + +此时有的算子在负重前行，有的算子拖了后腿，后者多于了前者。我们优先**找出来host耗时大于device的top算子，这些算子是拖后腿的**，可以交予算子团队重点分析。 + +### 2.组网合理性分析 + +有的时候，模型组网没有按照最高效的方式来，这一点在profiling中是非常易于识别的，下面会介绍一下分析思路并给出例子。 + +通常来讲，LLM中的大的热点算子是Attention和FFN中的矩阵乘计算，二者加起来在prefill下可能达到计算耗时的70%+，decode下可能达到50%+。如果整体的耗时比例不符合预期，或者profiling中出现了一些新面孔，或者拼接类算子太多了，这都值得我们去分析一下模型组网，是不是使用算子的方式错了？尤其是拼接类算子，是值得我们逐一分析的。 + +对于slice/split/concat这样的拼接类算子，还有transpose/cast这种转换算子，他们的存在往往是前后算子不直接配套造成的。如果前一个算子可以直接对输出做好尾处理，往往可以节省一个算子的启动开销和一次冗余读写。但这样的改变不一定符合算子的基本设计原则。 + +举一个正例，对于某次Matmul的输出shape为[m, n0 + n1]，在这后面我们接了两个slice，输入均为这个[m, n0 + n1]的tensor，输出分别为[m, n0]和[m, n1]。第一个优化的思路是将两个slice改为一个split，这样耗时可以基本减半，[m, n0 + n1]的显存也可以尽早释放。进一步优化的思路是将矩阵乘的权重从[k, n0 + n1]分割为[k, n0]和[k, n1]，将原来的矩阵乘任务分成两个（前提是这两个的耗时加起来不比之前的劣化太多，分核策略不能出问题），从而彻底消除这个slice/split操作。 + +![network_1](https://github.com/chengminhua/verl_data/raw/main/MindStudio_Insight_use/network_1.png) + +举一个反例，Rmsnorm(fp16)+Cast(fp16->fp32)+Matmul(fp32)，Rmsnorm虽然输入输出都是fp16，但考虑到累加运算的精度，内部是fp32做计算的。如果将Cast融到Rmsnorm内，本就内部使用fp32做计算的Rmsnorm就可以省去一个末尾fp32->fp16的cast，加上我们干掉的Cast，总共节省两个cast的同时避免了一次精度丢失。虽然这样看起来精度性能双收了，但fp16进，fp32出的Rmsnorm是反原则的（核心输入和输出需要是同数据类型），除非我们能在广大开源模型中频繁找到这样的结构，证明它的普适性，否则算子团队是不允许做这样的算子的。 + +![network_2](https://github.com/chengminhua/verl_data/raw/main/MindStudio_Insight_use/network_2.png) + +### 3.算子性能初诊 + +需要利用 `".\ASCEND_PROFILER_OUTPUT\operator_details.csv"`来做分析，从而判断算子是否有性能问题。 + +Profiling工具会统计这些流水线在不同核上的平均繁忙时间（xxx_time），与最慢核的完整kernel耗时（task_duration）做除法，得到流水线利用率（xxx_ratio）。这些流水线之间虽然互有依赖，且搬运类流水线会互抢带宽，但算子只要设计得当，是可以做到互相掩盖的。因此我们可以初步认为，**当算子的执行耗时大到一定程度上，算子应当在某一条流水线上形成bound**，即利用率要高到一定程度。经验上，在单算子耗时达到50μs时，就可以认为算子应当在bound流水线上，达成80%+的占用率了。 + +以下图为例，第一行是一个FA算子，第二行是一个Matmul算子，FA在vec流水线上达到了88.1%的利用率，Matmul算子在mac流水线上达到了89.8%的利用率，他们的性能可以认为是合格的。 + +![Operator%20performance](https://github.com/chengminhua/verl_data/raw/main/MindStudio_Insight_use/Operator%20performance.png) + +### 4.亲和shape调整 + +对于一个模型而言，超参是我们控制不了的，但我们可以控制并发度、权重格式、切分策略等因素来迎合算子，使其发挥出最大的性能，这一节主要从算子搬运效率和负载均衡两个方面出发，讨论模型侧值得尝试的调整方向。 + +#### 4.1 搬运效率亲和的shape + +mte2是一个自身效率严重受shape影响的流水线。要想让mte2保证最大搬运效率，我们需要保障如下两个条件至少满足其一： + +**（1）被搬运的矩阵使用nz作为format（最优） +（2）被搬运的矩阵的尾轴512B对齐，且不为16KB的整数倍（近似最优）** + +对于权重矩阵来说，推理阶段尤其是decode，我们通常满足（1），训练阶段我们通常满足（2）。**如果我们做不到（1），我们就要迎合（2）**。典型的手段有： + +1，如果没达成B的矩阵的首轴是亲和的而尾轴不亲和，那么对它做transpose +2，调整TP切分策略，避免出现不亲和的尾轴 + +#### 4.2 负载均衡亲和的shape + +在算子shape不大时，受制于算子语义，我们有可能不能把所有核都利用起来，或者即使开满核，负载均衡却很差。这一小节主要是对decode阶段的小shape做分析。 + +首先，我们明确出当前NPU卡是多少核的，如果不清楚，跑出来的profiling里都是20，40这样的数，就说明是20核，反之是24核。这里我的24核其实是代表了一个cube和两个vector组成的小组，我们可以认为是一个cube作为主核，带了两个vector作为从核。如果一个算子是纯vector算子，那么就不再有组的概念，40或48个vector核会作为主核直接独立去拿逻辑任务。 + +对于LLM中的vector算子，它的一种常见分核策略有可能是分在最高维，也就是batch维，常见于对低维（也叫尾轴）有规约操作的norm类、动态量化类等算子；另一种是整体拍平，允许算子切分的非常细的算子，如elementwise算子。对于第一种，我们就可以在模型侧关注它的负载均衡问题。例如我们打48batch，而硬件却是个40个vector核，那这40个核会循环2次，第二次有多数的核会无事可做，这个batch数就可以认为是不友好的。如果将batch打到64或80，性能可以预见会是无损的。同样的情况下，如果是48核的卡，那我们可以认为这就是个非常友好的batch数。 + +对于cube类算子，它常见的分核策略是以base快去切分M和N（K轴是累加轴，对它分核会引入确定性问题）。最常见的分块是baseM=128，baseN=256。在decode阶段，我们的耗时基本可以看做都是在搬权重，这是因为激活的M极小，M方向大概率只分了一块，那么右矩阵就只需要搬一次。所以我们在M≤128的范围内可以尽情提高M，对性能都基本是无损的，如果M大于128，可以认为(128, 256]是下一个性能分档。 +除了M外，N轴切分的任务也影响算子亲和性，以deepseekR1中的MLA预处理为例，它会使用同一个激活（shape为[batch_size, 7168]）与两个权重做矩阵乘(shape为[7168, 1536]和[7168, 576])。在batch_size打不大的情况下，即使baseN缩短为128，N轴都不能用满核数，所以此时这两个矩阵乘各自的耗时，会约等于将他们权重N轴拼起来乘(shape为[7168, 2112])的矩阵乘的耗时。如果仅考虑模型竞争力，我们更希望对这两个权重做合并，否则两个小的矩阵乘带宽利用率都会非常差。 + +对于Attention算子，它常见的分核策略是q_seqlen、batch_size和kv_headnum。增量阶段q_seqlen会以MTP和GQA倍数做合并，但是通常也不会大过128，划分不出第二个任务，那么并行度基本就是batch_size * kv_headnum。 + +总的来说，我们可以依据shape信息和算子类别，对算子是否有负载均衡问题作出识别，从而对我们切分策略选择，最高吞吐量的batch策略作出预判。 diff --git a/verl/docs/ascend_tutorial/dev_guide/performance/ascend_profiling_en.rst b/verl/docs/ascend_tutorial/dev_guide/performance/ascend_profiling_en.rst new file mode 100644 index 0000000000000000000000000000000000000000..e30c0ebfc513755fcdc59964abdb2408107abb44 --- /dev/null +++ b/verl/docs/ascend_tutorial/dev_guide/performance/ascend_profiling_en.rst @@ -0,0 +1,390 @@ +Profiling Data Collection Guide +========================================================================================== + +Last updated: 12/20/2025. + +This is a tutorial for data collection using the GRPO or DAPO algorithm +based on FSDP or MindSpeed(Megatron) on Ascend devices. + +Configuration +------------- + +Leverage two levels of configuration to control data collection: + +- **Global profiler control**: Use parameters in ``verl/trainer/config/ppo_trainer.yaml`` (FSDP) or ``verl/trainer/config/ppo_megatron_trainer.yaml`` (MindSpeed) to control the collection mode and steps. +- **Role profile control**: Use parameters in each role's ``profile`` field to control various parameters. + +Global collection control +~~~~~~~~~~~~~~~~~~~~~~~~~ + +Use parameters in ppo_trainer.yaml to control the collection mode +and steps. + +- global_profiler: Control the ranks and mode of profiling + + - tool: The profiling tool to use, options are nsys, npu, torch, + torch_memory. + - steps: This parameter can be set as a list that has + collection steps, such as [2, 4], which means it will collect steps 2 + and 4. If set to null, no collection occurs. + - save_path: The path to save the collected data. Default is + "outputs/profile". + + +Role collection control +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +In each role's ``profiler`` field, you can control the collection mode for that role. + +- enable: Whether to enable profiling for this role. +- all_ranks: Whether to collect data from all ranks. +- ranks: A list of ranks to collect data from. If empty, no data is collected. +- tool_config: Configuration for the profiling tool used by this role. + +Use parameters in each role's ``profiler.tool_config.npu`` to control npu profiler behavior: + +- level: Collection level—options are level_none, level0, level1, and + level2 + + - level_none: Disables all level-based data collection (turns off profiler_level). + - level0: Collect high-level application data, underlying NPU data, and operator execution details on NPU. After balancing data volume and analytical capability, Level 0 is recommended as the default configuration. + - level1: Extends level0 by adding CANN-layer AscendCL data and AI Core performance metrics on NPU. + - level2: Extends level1 by adding CANN-layer Runtime data and AI CPU metrics. + +- contents: A list of options to control the collection content, such as + npu, cpu, memory, shapes, module, stack. + + - npu: Whether to collect device-side performance data. + - cpu: Whether to collect host-side performance data. + - memory: Whether to enable memory analysis. + - shapes: Whether to record tensor shapes. + - module: Whether to record framework-layer Python call stack information. It is recommended to use 'module' instead of 'stack' for recording call stack information, as it costs less performance overhead. + - stack: Whether to record operator call stack information. + +- analysis: Enables automatic data parsing. +- discrete: Whether to enable discrete mode. + + +Examples +-------- + +Disabling collection +~~~~~~~~~~~~~~~~~~~~ + +.. code:: yaml + + global_profiler: + steps: null # disable profile + +End-to-End collection +~~~~~~~~~~~~~~~~~~~~~ + +.. code:: yaml + + global_profiler: + steps: [1, 2, 5] + save_path: ./outputs/profile + actor_rollout_ref: + actor: # Set actor role profiler collection configuration parameters + profiler: + enable: True + all_ranks: True + tool_config: + npu: + discrete: False + contents: [npu, cpu] # Control collection list, default cpu, npu, can configure memory, shapes, module, etc. + # rollout & ref follow actor settings + + +Discrete Mode Collection +~~~~~~~~~~~~~~~~~~~~~~~~ + +.. code:: yaml + + global_profiler: + steps: [1, 2, 5] + save_path: ./outputs/profile + actor_rollout_ref: + actor: + profiler: + enable: True # Set to True to profile training + all_ranks: False + ranks: [0] # Global Rank 0 + tool_config: + npu: + discrete: True + contents: [npu, cpu] + rollout: + profiler: + enable: True # Set to True to profile inference + all_ranks: False + ranks: [0] # In Agent Loop mode, this is the Replica Rank (e.g., 0-th instance) + tool_config: + npu: + discrete: True # Must be enabled in Agent Loop mode + # ref follow actor settings + +**Agent Loop Mode Description**: + +When Rollout runs in `Agent Loop <../advance/agent_loop.rst>`_ mode, performance data for the Rollout phase **must be collected using discrete mode**. In this case, the Profiler is triggered by the inference engine backend. + +1. Rank Definition: ranks in the Rollout configuration refers to Replica Rank (inference instance index), not Global Rank. + +2. Inference Engine Support: Currently, vLLM and SGLang engines are supported without additional settings. Specific details are as follows: + + - vLLM Engine: Automatically collects AsyncLLM scheduling stacks and inference process performance data. Does not support setting analysis (defaults to no analysis, requires offline analysis) and profiler_level (defaults to level1). + - SGLang Engine: Automatically collects inference process performance data. Does not support the memory option in contents. Does not support setting analysis (defaults to enabled) and profiler_level (defaults to level0). + + +Visualization +------------- + +Collected data is stored in the user-defined save_path and can be +visualized by using the `MindStudio Insight `_ tool. + +Additionally, in a Linux environment, the MindStudio Insight tool is provided in the form of a `JupyterLab Plugin `_ ，offering a more intuitive and highly interactive user interface. The advantages of the JupyterLab plugin are as follows: + +- Seamless integration: Supports running the MindStudio Insight tool directly within the Jupyter environment, eliminating the need to switch platforms or copy data from the server, enabling data to be collected and used immediately. +- Fast startup: Allows MindStudio Insight to be launched quickly via the JupyterLab command line or graphical interface. +- Smooth operation: In a Linux environment, launching MindStudio Insight through JupyterLab effectively alleviates performance lag compared to the full-package communication mode, significantly improving the user experience. +- Remote access: Supports remotely launching MindStudio Insight. Users can connect to the service via a local browser for direct visual analysis, reducing the difficulty of uploading and downloading data during large-model training or inference. + +If the analysis parameter is set to False, offline parsing is required after data collection: + +.. code:: python + + import torch_npu + # Set profiler_path to the parent directory of the "localhost.localdomain___ascend_pt" folder + torch_npu.profiler.profiler.analyse(profiler_path=profiler_path) + + +Advanced Guide: Fine-grained Collection +--------------------------------------- + +Background and Challenges +~~~~~~~~~~~~~~~~~~~~~~~~~ + +Although the configuration-based collection method mentioned above is convenient, it faces challenges in training scenarios with **long sequences (Long Context)** or **large global batch sizes (Large Global Batch Size)**. Within a complete training step (Step), model computation exhibits high-frequency and repetitive characteristics: + +1. **Rollout phase**: Sequence generation (Generate Sequence) is an autoregressive process involving thousands of forward computations of the Decoder model. +2. **Training phase**: To control peak memory usage, verl typically adopts a Micro-Batch strategy, dividing large data streams into multiple micro-batches for computation. + + - **compute_log_prob (Actor/Ref)**: Involves multiple rounds of pure forward propagation. + - **update_policy (Actor/Critic)**: Involves multiple rounds of forward and backward propagation. + +This characteristic leads to massive and repetitive operator records from full profiling. As shown in the image below: + +.. image:: https://raw.githubusercontent.com/mengchengTang/verl-data/master/verl_ascend_profiler.png + +Even with ``discrete`` mode enabled, performance data files for a single stage can still reach several TB, leading to **parsing failures** or **visualization tool lag**. + +Solution: Critical Path Sampling +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +To solve the above problems, we can adopt a **critical path sampling** strategy: Based on the API interface provided by `torch_npu.profiler `_, directly modify Python source code to collect only representative data segments (such as specific Decode Steps or the first Micro-Batch). + + **Important Notes** + + 1. This chapter involves direct source code modification. It is recommended to back up files before modification and restore them after debugging. + 2. When using code instrumentation for collection, be sure to **disable global collection** (``global_profiler: steps: null``) in ``ppo_trainer.yaml`` or ``ppo_megatron_trainer.yaml`` to avoid Profiler conflicts. + +1. Fine-grained Collection in Rollout Phase +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +For vLLM or SGLang inference engines, we can control the ``schedule`` parameter to collect model forward propagation performance data for specific tokens. + +**vLLM Engine** + +- **Reference Version**: vLLM v0.11.0, vLLM-Ascend v0.11.0rc1 +- **Modified File**: ``vllm-ascend/vllm_ascend/worker/worker_v1.py`` + +.. code-block:: diff + + class NPUWorker(WorkerBase): + + def __init__(self, *args, **kwargs): + # ... existing code ... + + + # Initialize profiler + + import torch_npu + + experimental_config = torch_npu.profiler._ExperimentalConfig( + + profiler_level=torch_npu.profiler.ProfilerLevel.Level1, + + export_type=torch_npu.profiler.ExportType.Db, # You can choose torch_npu.profiler.ExportType.Text format + + ) + + self.profiler_npu = torch_npu.profiler.profile( + + activities=[torch_npu.profiler.ProfilerActivity.CPU, torch_npu.profiler.ProfilerActivity.NPU], + + with_modules=False, # Collect call stack + + profile_memory=False, # Collect memory + + experimental_config=experimental_config, + + # Skip first step, warmup one step, collect 3 steps, repeat 1 time. If you want to collect decode steps 30~70, set schedule=torch_npu.profiler.schedule(wait=29, warmup=1, active=30, repeat=1) + + schedule=torch_npu.profiler.schedule(wait=1, warmup=1, active=3, repeat=1), + + on_trace_ready=torch_npu.profiler.tensorboard_trace_handler("./outputs/vllm_profile", analyse_flag=True) # Data save path and whether to parse online + + ) + + self.profiler_npu.start() + + # ... existing code ... + + def execute_model(self, scheduler_output=None, intermediate_tensors=None, **kwargs): + # ... existing code ... + output = self.model_runner.execute_model(scheduler_output, + intermediate_tensors) + + + self.profiler_npu.step() # Drive schedule to collect partial decode steps + + # ... existing code ... + +**SGLang Engine** + +- **Reference Version**: SGLang master branch +- **Modified File**: ``sglang/python/sglang/srt/model_executor/model_runner.py`` + +.. code-block:: diff + + # ... existing imports ... + + import torch_npu + + class ModelRunner: + + def __init__(self, *args, **kwargs): + # ... existing init code ... + + + # Initialize profiler (same configuration as above, omitted) + + experimental_config = torch_npu.profiler._ExperimentalConfig(...) + + self.profiler_npu = torch_npu.profiler.profile( + + # ... + + # Skip first step, warmup one step, collect 3 steps, repeat 1 time. + + schedule=torch_npu.profiler.schedule(wait=1, warmup=1, active=3, repeat=1), + + on_trace_ready=torch_npu.profiler.tensorboard_trace_handler("./outputs/sglang_profile", analyse_flag=True) + + ) + + self.profiler_npu.start() + + def forward(self, forward_batch, **kwargs): + # ... existing code ... + + + self.profiler_npu.step() # Drive schedule to collect partial decode steps + return output + +2. Fine-grained Collection in compute_log_prob (Actor & Ref) Phase +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +This phase computes probability distributions for new and old policies. +With the unified model engine, both actor and reference log-prob +computation go through ``TrainingWorker.infer_batch`` which dispatches +to ``BaseEngine.infer_batch`` on the underlying backend engine. + +**FSDP Backend** + +The FSDP backend allows fine-grained control at the Micro-Batch level. +Instrument the micro-batch loop inside the FSDP engine's forward pass. + +- **Modified File**: ``verl/workers/engine/fsdp/transformer_impl.py`` + (``FSDPEngineWithLMHead.forward_backward_batch`` / ``forward_step``) + +.. code-block:: diff + + # ... import dependencies ... + + import torch_npu + + class FSDPEngineWithLMHead(FSDPEngine): + + def forward_backward_batch(self, data: TensorDict, loss_function, forward_only=False): + + + role = "Ref" if forward_only and not self.optimizer_config else "Actor" + + # Prepare profiler (same configuration as above, omitted) + + experimental_config = torch_npu.profiler._ExperimentalConfig(...) + + self.prof_npu = torch_npu.profiler.profile( + + # ... + + # wait=0, warmup=0, active=1: directly collect first micro-batch + + schedule=torch_npu.profiler.schedule(wait=0, warmup=0, active=1, repeat=1), + + on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(f"./outputs/{role}_compute_log_prob", analyse_flag=True) + + ) + + + # forward_backward_batch is shared by ref and actor. Use the role flag to + + # distinguish; to collect actor_compute_log_prob instead, switch to role=="Actor": + + if role == "Ref": + + self.prof_npu.start() + + for micro_batch in micro_batches: + + # ... original computation logic ... + with torch.no_grad(): + output = self.forward_step(micro_batch, loss_function, forward_only=True) + + + # Drive schedule to collect micro batch + + if role == "Ref": + + self.prof_npu.step() + + # ... + + +**Megatron Backend** + +The Micro-Batch scheduling in the Megatron backend is managed internally +by Megatron's pipeline-parallel ``forward_backward_func`` and does not +currently support fine-grained collection at the Micro-Batch level +through simple code instrumentation. It is recommended to use the global +profiler configuration for collection. + +3. Fine-grained Collection in update_policy (Actor & Critic) Phase +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The Update phase includes forward and backward propagation. In the +unified engine, mini-batch iteration is driven by +``TrainingWorker.train_mini_batch`` in ``verl/workers/engine_workers.py``, +which then calls ``train_batch`` for each mini-batch. + +**FSDP Backend** + +The FSDP backend supports collection at both Mini-Batch and Micro-Batch +granularities. For Mini-Batch scope, instrument ``train_mini_batch`` in +``TrainingWorker``; for Micro-Batch scope, instrument the per-micro-batch +loop inside the FSDP engine's ``forward_backward_batch``. + +- **Modified File**: ``verl/workers/engine_workers.py`` + (``TrainingWorker.train_mini_batch``) for Mini-Batch granularity, or + ``verl/workers/engine/fsdp/transformer_impl.py`` + (``FSDPEngineWithLMHead.forward_backward_batch`` for Micro-Batch) + +.. code-block:: diff + + # ... import dependencies ... + + import torch_npu + + class TrainingWorker(Worker, DistProfilerExtension): + + def train_mini_batch(self, data: TensorDict) -> TensorDict: + + + # Prepare profiler (same configuration as above, omitted) + + experimental_config = torch_npu.profiler._ExperimentalConfig(...) + + self.prof_npu = torch_npu.profiler.profile( + + # ... + + # Only collect first Mini Batch (including all Micro-Batch computations and one optimizer update) + + schedule=torch_npu.profiler.schedule(wait=0, warmup=0, active=1, repeat=1), + + on_trace_ready=torch_npu.profiler.tensorboard_trace_handler("./outputs/fsdp_actor_update_profile", analyse_flag=True) + + ) + + self.prof_npu.start() + + # ... Mini Batch loop over the dataloader ... + for batch_idx, mini_batch_td in enumerate(dataloader): + # ... calls self.train_batch(mini_batch_td), which in turn runs + # Forward & Backward on every micro-batch and a single optimizer step + # inside the engine ... + actor_output = self.train_batch(mini_batch_td) + + + # Drive schedule to collect mini batch; for micro-batch granularity, move + + # self.prof_npu.step() into the micro-batch loop inside + + # FSDPEngineWithLMHead.forward_backward_batch. + + self.prof_npu.step() + + +**Megatron Backend** + +The Megatron backend supports collection at the Mini-Batch granularity. +The same ``TrainingWorker.train_mini_batch`` entry point applies – the +Megatron engine internally runs Megatron's pipeline-parallel forward / +backward schedule and the optimizer step. + +- **Modified File**: ``verl/workers/engine_workers.py`` + (``TrainingWorker.train_mini_batch``) — identical to the FSDP snippet + above; the output path should be renamed (e.g. ``./outputs/megatron_actor_update_profile``) + to distinguish traces from different backends. \ No newline at end of file diff --git a/verl/docs/ascend_tutorial/dev_guide/performance/ascend_profiling_zh.rst b/verl/docs/ascend_tutorial/dev_guide/performance/ascend_profiling_zh.rst new file mode 100644 index 0000000000000000000000000000000000000000..052b279a667351604239a733baedef505566419e --- /dev/null +++ b/verl/docs/ascend_tutorial/dev_guide/performance/ascend_profiling_zh.rst @@ -0,0 +1,375 @@ +Profiling采集指导 +================================================================================== + +Last updated: 12/20/2025. + +这是一份在昇腾设备上基于FSDP或MindSpeed(Megatron)后端，使用GRPO或DAPO算法进行数据采集的教程。 + +配置 +---- + +使用两级profile设置来控制数据采集 + +- 全局采集控制：使用verl/trainer/config/ppo_trainer.yaml(FSDP)，或verl/trainer/config/ppo_megatron_trainer.yaml(MindSpeed)中的配置项控制采集的模式和步数。 +- 角色profile控制：通过每个角色中的配置项控制等参数。 + +全局采集控制 +~~~~~~~~~~~~ + +通过 ppo_trainer.yaml 中的参数控制采集步数和模式： + +- global_profiler: 控制采集的rank和模式 + + - tool: 使用的采集工具，选项有 nsys、npu、torch、torch_memory。 + - steps: 此参数可以设置为包含采集步数的列表，例如 [2, 4]，表示将采集第2步和第4步。如果设置为 null，则不进行采集。 + - save_path: 保存采集数据的路径。默认值为 "outputs/profile"。 + +角色profiler控制 +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +在每个角色的 ``profiler`` 字段中，您可以控制该角色的采集模式。 + +- enable: 是否为此角色启用性能分析。 +- all_ranks: 是否从所有rank收集数据。 +- ranks: 要收集数据的rank列表。如果为空，则不收集数据。 +- tool_config: 此角色使用的性能分析工具的配置。 + +通过每个角色的 ``profiler.tool_config.npu`` 中的参数控制具体采集行为： + +- level: 采集级别—选项有 level_none、level0、level1 和 level2 + + - level_none: 禁用所有基于级别的数据采集（关闭 profiler_level）。 + - level0: 采集高级应用数据、底层NPU数据和NPU上的算子执行详情。在权衡数据量和分析能力后，level0是推荐的默认配置。 + - level1: 在level0基础上增加CANN层AscendCL数据和NPU上的AI Core性能指标。 + - level2: 在level1基础上增加CANN层Runtime数据和AI CPU指标。 + +- contents: 控制采集内容的选项列表，例如 + npu、cpu、memory、shapes、module、stack。 + + - npu: 是否采集设备端性能数据。 + - cpu: 是否采集主机端性能数据。 + - memory: 是否启用内存分析。 + - shapes: 是否记录张量形状。 + - module: 是否记录框架层Python调用栈信息。相较于stack，更推荐使用module记录调用栈信息，因其产生的性能膨胀更低。 + - stack: 是否记录算子调用栈信息。 + +- analysis: 启用自动数据解析。 +- discrete: 使用离散模式。 + +示例 +---- + +禁用采集 +~~~~~~~~~~~~~~~~~~~~ + +.. code:: yaml + + global_profiler: + steps: null # disable profile + +端到端采集 +~~~~~~~~~~~~~~~~~~~~~ + +.. code:: yaml + + global_profiler: + steps: [1, 2, 5] + save_path: ./outputs/profile + actor_rollout_ref: + actor: # 设置 actor role 的 profiler 采集配置参数 + profiler: + enable: True + all_ranks: True + tool_config: + npu: + discrete: False + contents: [npu, cpu] # 控制采集列表，默认cpu、npu，可配置memory、shapes、module等 + + # rollout & ref follow actor settings + + +离散模式采集 +~~~~~~~~~~~~~~~~~~~~~~~~ + +.. code:: yaml + + global_profiler: + steps: [1, 2, 5] + save_path: ./outputs/profile + actor_rollout_ref: + actor: + profiler: + enable: True # 设置为 True 以采集训练阶段 + all_ranks: False + ranks: [0] # 全局 Rank 0 + tool_config: + npu: + discrete: True + contents: [npu, cpu] + rollout: + profiler: + enable: True # 设置为 True 以采集推理阶段 + all_ranks: False + ranks: [0] # 在 Agent Loop 模式下，此处指推理实例的 Replica Rank (例如第 0 个实例) + tool_config: + npu: + discrete: True # Agent Loop 模式下必须开启离散模式 + # ref follow actor settings + +**Agent Loop 模式说明**： + +在 `Agent Loop <../advance/agent_loop.rst>`_ 模式下，Rollout 阶段的性能数据 **必须使用离散模式** 采集，此时 Profiler 由推理引擎后端触发。 + +1. Rank 定义：Rollout 配置中的 ranks 指代 Replica Rank（推理实例索引），而非全局 Rank。 + +2. 推理引擎支持：当前支持vLLM和SGLang引擎，无需额外设置。具体说明如下： + + - vLLM 引擎：自动采集 AsyncLLM 调度栈及推理进程性能数据。不支持设置 analysis（默认不解析，需离线解析）和 profiler_level（默认 level1）。 + - SGLang 引擎：自动采集推理进程性能数据。不支持 contents 中的 memory 配置项。不支持设置 analysis（默认解析）和 profiler_level（默认 level0）。 + +可视化 +------ + +采集后的数据存放在用户设置的save_path下，可通过 `MindStudio Insight `_ 工具进行可视化。 + +另外在Linux环境下，MindStudio Insight工具提供了 `JupyterLab插件 `_ 形态，提供更直观和交互式强的操作界面。JupyterLab插件优势如下： + +- 无缝集成：支持在Jupyter环境中直接运行MindStudio Insight工具，无需切换平台，无需拷贝服务器上的数据，实现数据即采即用。 +- 快速启动：通过JupyterLab的命令行或图形界面，可快速启动MindStudio Insight工具。 +- 运行流畅：在Linux环境下，通过JupyterLab环境启动MindStudio Insight，相较于整包通信，有效解决了运行卡顿问题，操作体验显著提升。 +- 远程访问：支持远程启动MindStudio Insight，可通过本地浏览器远程连接服务直接进行可视化分析，缓解了大模型训练或推理数据上传和下载的困难。 + +如果analysis参数设置为False，采集之后需要进行离线解析： + +.. code:: python + + import torch_npu + # profiler_path请设置为"localhost.localdomain___ascend_pt"目录的上一级目录 + torch_npu.profiler.profiler.analyse(profiler_path=profiler_path) + + +进阶指南：精细化采集 +-------------------- + +背景与挑战 +~~~~~~~~~~ + +上述基于配置文件的采集方式虽然便捷，但在 **长序列 (Long Context)** 或 **大全局批量 (Large Global Batch Size)** 的训练场景中面临挑战。 +在一个完整的训练步 (Step) 内，模型计算呈现出高频次、重复性的特征： + +1. Rollout 阶段：序列生成 (Generate Sequence) 是一个自回归过程，涉及成千上万次 Decoder 模型的前向计算。 +2. Training 阶段：为了控制显存峰值，verl 通常采用 Micro-Batch 策略，将庞大的数据流切分为多个微批次进行计算。 + + - compute_log_prob (Actor/Ref)：涉及多轮纯前向传播。 + - update_policy (Actor/Critic)：涉及多轮前向与反向传播。 + +这种特性会导致全量 Profiling 产生海量且重复的算子记录。如下图所示： + +.. image:: https://raw.githubusercontent.com/mengchengTang/verl-data/master/verl_ascend_profiler.png + +即使使用了 ``discrete`` 模式，单个阶段的性能数据文件仍可能达到数 TB，导致 **解析失败** 或 **可视化工具卡顿** 。 + +解决方案：关键路径采样 +~~~~~~~~~~~~~~~~~~~~~~ + +为了解决上述问题，我们可以采用 **关键路径采样** 策略：基于 `torch_npu.profiler `_ 提供的API接口，直接修改 Python 源码，仅采集具有代表性的数据片段（如特定 Decode Step 或首个 Micro-Batch）。 + + **重要提示** + + 1. 本章节涉及直接修改源码。建议修改前备份文件，调试完成后恢复。 + 2. 使用代码插桩采集时，请务必在 ``ppo_trainer.yaml`` 或 ``ppo_megatron_trainer.yaml`` 中**禁用全局采集** (``global_profiler: steps: null``)，以避免 Profiler 冲突。 + +1. Rollout 阶段精细化采集 +~~~~~~~~~~~~~~~~~~~~~~~~~ + +对于 vLLM 或 SGLang 推理引擎，我们可以通过控制 ``schedule`` 参数来控制采集模型在特定token的前向传播性能数据。 + +**vLLM 引擎** + +- **参考版本**：vLLM v0.11.0, vLLM-Ascend v0.11.0rc1 +- **修改文件**：``vllm-ascend/vllm_ascend/worker/worker_v1.py`` + +.. code-block:: diff + + class NPUWorker(WorkerBase): + + def __init__(self, *args, **kwargs): + # ... existing code ... + + + # Initialize profiler + + import torch_npu + + experimental_config = torch_npu.profiler._ExperimentalConfig( + + profiler_level=torch_npu.profiler.ProfilerLevel.Level1, + + export_type=torch_npu.profiler.ExportType.Db, # 可选择torch_npu.profiler.ExportType.Text格式 + + ) + + self.profiler_npu = torch_npu.profiler.profile( + + activities=[torch_npu.profiler.ProfilerActivity.CPU, torch_npu.profiler.ProfilerActivity.NPU], + + with_modules=False, # 采集调用栈 + + profile_memory=False, # 采集内存 + + experimental_config=experimental_config, + + # 跳过第一步，warmup一步，采集3步，重复1次。如果想采集第30~70个decode step，可以设置为schedule=torch_npu.profiler.schedule(wait=29, warmup=1, active=30, repeat=1) + + schedule=torch_npu.profiler.schedule(wait=1, warmup=1, active=3, repeat=1), + + on_trace_ready=torch_npu.profiler.tensorboard_trace_handler("./outputs/vllm_profile", analyse_flag=True) # 采集数据保存路径，是否在线解析 + + ) + + self.profiler_npu.start() + + # ... existing code ... + + def execute_model(self, scheduler_output=None, intermediate_tensors=None, **kwargs): + # ... existing code ... + output = self.model_runner.execute_model(scheduler_output, + intermediate_tensors) + + + self.profiler_npu.step() # 驱动 schedule，对部分decode step进行采集 + + # ... existing code ... + +**SGLang 引擎** + +- **参考版本**：SGLang master 分支 +- **修改文件**：``sglang/python/sglang/srt/model_executor/model_runner.py`` + +.. code-block:: diff + + # ... existing imports ... + + import torch_npu + + class ModelRunner: + + def __init__(self, *args, **kwargs): + # ... existing init code ... + + + # Initialize profiler (配置同上，略) + + experimental_config = torch_npu.profiler._ExperimentalConfig(...) + + self.profiler_npu = torch_npu.profiler.profile( + + # ... + + # 跳过第一步，warmup一步，采集3步，重复1次。 + + schedule=torch_npu.profiler.schedule(wait=1, warmup=1, active=3, repeat=1), + + on_trace_ready=torch_npu.profiler.tensorboard_trace_handler("./outputs/sglang_profile", analyse_flag=True) + + ) + + self.profiler_npu.start() + + def forward(self, forward_batch, **kwargs): + # ... existing code ... + + + self.profiler_npu.step() # 驱动 schedule，对部分decode step进行采集 + return output + +2. compute_log_prob (Actor & Ref) 阶段精细化采集 +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +该阶段计算新旧策略的概率分布。统一模型引擎下，actor 和 ref 的 log-prob +计算都走 ``TrainingWorker.infer_batch``，最终分发到对应后端引擎的 +``BaseEngine.infer_batch`` 上。 + +**FSDP 后端** + +FSDP 后端允许在 Micro-Batch 级别进行精细控制，可在 FSDP 引擎 forward 过程 +的 micro-batch 循环内插桩。 + +- **修改文件**：``verl/workers/engine/fsdp/transformer_impl.py`` + （``FSDPEngineWithLMHead.forward_backward_batch`` / ``forward_step``） + +.. code-block:: diff + + # ... 引入依赖 ... + + import torch_npu + + class FSDPEngineWithLMHead(FSDPEngine): + + def forward_backward_batch(self, data: TensorDict, loss_function, forward_only=False): + + + role = "Ref" if forward_only and not self.optimizer_config else "Actor" + + # 准备 profiler (配置同上，略) + + experimental_config = torch_npu.profiler._ExperimentalConfig(...) + + self.prof_npu = torch_npu.profiler.profile( + + # ... + + # wait=0, warmup=0, active=1: 直接采集第一个 micro-batch + + schedule=torch_npu.profiler.schedule(wait=0, warmup=0, active=1, repeat=1), + + on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(f"./outputs/{role}_compute_log_prob", analyse_flag=True) + + ) + + + # forward_backward_batch 被 ref 和 actor 共用，通过 role 标志位区分； + + # 如需采集 actor_compute_log_prob，可改为 role == "Actor": + + if role == "Ref": + + self.prof_npu.start() + + for micro_batch in micro_batches: + + # ... 原始计算逻辑 ... + with torch.no_grad(): + output = self.forward_step(micro_batch, loss_function, forward_only=True) + + + # 驱动 schedule，对micro batch进行采集 + + if role == "Ref": + + self.prof_npu.step() + + # ... + + +**Megatron 后端** + +Megatron 后端的 Micro-Batch 调度由 Megatron 的流水并行 +``forward_backward_func`` 内部管理，暂不支持通过简单的代码插桩进行 +Micro-Batch 级别的精细化采集。建议使用全局 profiler 配置进行采集。 + +3. update_policy (Actor & Critic) 阶段精细化采集 +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Update 阶段包含前向和反向传播。统一模型引擎下，mini-batch 循环由 +``verl/workers/engine_workers.py`` 中的 ``TrainingWorker.train_mini_batch`` +驱动，它会对每个 mini-batch 调用 ``train_batch``。 + +**FSDP 后端** + +FSDP 后端支持设置对 Mini-Batch 和 Micro-Batch 的粒度进行采集。 +Mini-Batch 级别请插桩 ``TrainingWorker.train_mini_batch``； +Micro-Batch 级别请插桩 FSDP 引擎的 ``forward_backward_batch`` 中的 +micro-batch 循环。 + +- **修改文件**：``verl/workers/engine_workers.py`` + （``TrainingWorker.train_mini_batch``，Mini-Batch 粒度）或 + ``verl/workers/engine/fsdp/transformer_impl.py`` + （``FSDPEngineWithLMHead.forward_backward_batch``，Micro-Batch 粒度） + +.. code-block:: diff + + # ... 引入依赖 ... + + import torch_npu + + class TrainingWorker(Worker, DistProfilerExtension): + + def train_mini_batch(self, data: TensorDict) -> TensorDict: + + + # 准备 profiler (配置同上，略) + + experimental_config = torch_npu.profiler._ExperimentalConfig(...) + + self.prof_npu = torch_npu.profiler.profile( + + # ... + + # 仅采集第一个 Mini Batch（包含所有 Micro-Batch 的计算和一次优化器更新） + + schedule=torch_npu.profiler.schedule(wait=0, warmup=0, active=1, repeat=1), + + on_trace_ready=torch_npu.profiler.tensorboard_trace_handler("./outputs/fsdp_actor_update_profile", analyse_flag=True) + + ) + + self.prof_npu.start() + + # ... Mini Batch 循环（遍历 dataloader） ... + for batch_idx, mini_batch_td in enumerate(dataloader): + # ... 内部调用 self.train_batch(mini_batch_td)，后者在引擎内部 + # 对每个 micro-batch 执行 Forward & Backward，并完成一次优化器更新 ... + actor_output = self.train_batch(mini_batch_td) + + + # 驱动 schedule，对 mini batch 进行采集；如需 micro-batch 粒度， + + # 请将 self.prof_npu.step() 移动到 + + # FSDPEngineWithLMHead.forward_backward_batch 中的 micro-batch 循环内。 + + self.prof_npu.step() + + +**Megatron 后端** + +Megatron 后端支持以 Mini-Batch 的粒度进行采集，入口同样是 +``TrainingWorker.train_mini_batch``：Megatron 引擎内部会调用 Megatron 的 +流水并行 forward/backward 调度并执行一次优化器 step。 + +- **修改文件**：``verl/workers/engine_workers.py`` + （``TrainingWorker.train_mini_batch``）—— 与上方 FSDP 代码片段完全一致， + 建议将输出目录改名（例如 ``./outputs/megatron_actor_update_profile``） + 以区分不同后端的 trace。 diff --git a/verl/docs/ascend_tutorial/dev_guide/performance/perf_tuning_on_ascend.rst b/verl/docs/ascend_tutorial/dev_guide/performance/perf_tuning_on_ascend.rst new file mode 100644 index 0000000000000000000000000000000000000000..166dd1b174d57d92439116002fb28c422e4b2324 --- /dev/null +++ b/verl/docs/ascend_tutorial/dev_guide/performance/perf_tuning_on_ascend.rst @@ -0,0 +1,258 @@ +Performance Tuning Guide on Ascend +==================================== + +Last updated: 01/29/2026. + +Author: `Xiaobo Hu `_, `Haozhe Li `_ + +`Perf Tuning `_ 中介绍的性能调优方法在昇腾设备中同样适用。本文重点介绍了昇腾特有的一些调优手段，包括融合算子优化、特定硬件配置和昇腾亲和特性等。 + +融合算子 +-------------------------- + +常用融合算子列表 +********************************** + +融合算子的优化原理为，通过数学意义上的等价替换，将多个算子融为一个算子的计算，减少冗余计算，同时减少下发次数，从而提高性能。几个典型的NPU融合算子列举如下，目前均已在 npu_patch.py 中对 Qwen2、Qwen3 系列模型完成替换。 + +当前verl中使用的全量融合算子请查阅 `npu_patch.py `_ + +Matrix Computation-Communication operator fusion (MC2) +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +MC2 是 CANN 中一系列计算通信融合算子的统称，这些算子将原本串行的通信和计算操作融合在一起，通过内部的切分和流水线并行执行来优化性能。 + +在 vllm-ascend 中，可以通过指定环境变量： + +.. code-block:: sh + + export VLLM_ASCEND_ENABLE_MATMUL_ALLREDUCE=1 + +在前向计算的 ``RowParallelLinear`` 中使能 ``torch_npu.npu_mm_all_reduce_base`` ，将分离的 ``matmul`` 和 ``allreduce`` 合并为一个融合算子。 + +`RotaryMul&RotaryMulGrad `_ +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +torch_npu 接口: ``torch_npu.npu_rotary_mul(x, r1, r2)`` + +参数说明: + +- x: q，k，shape要求输入为4维，一般为 ``[B, N, S, D]`` 或 ``[B, S, N, D]`` 或 ``[S, B, N, D]`` 。 + +- r1: cos值 ，shape要求输入为4维，一般为 ``[1, 1, S, D]`` 或 ``[1, S, 1, D]`` 或 ``[S, 1, 1, D]`` 。 + +- r2: sin 值，shape要求输入为4维，一般为 ``[1, 1, S, D]`` 或 ``[1, S, 1, D]`` 或 ``[S, 1, 1, D]`` 。 + +`RmsNorm&RmsNormGrad `_ +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +torch_npu 接口: ``torch_npu.npu_rms_norm(self, gamma, epsilon=1e-06) -> (Tensor, Tensor)`` +参数说明: + +- self: Tensor 类型，shape 支持 1-8 维。 + +- gamma: Tensor 类型，通常为weight，shape 要求与 self 的后几维保持一致。 + +- epsilon: Float 数据类型，用于防止除 0 错误。 + +输出说明: + +- 第 1 个输出为 Tensor，计算公式的最终输出y。 + +- 第 2 个输出为 Tensor， rms_norm 的中间结果 rstd ，用于反向计算。 + +`Swiglu `_ +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +torch_npu 接口: ``torch_npu.npu_swiglu(Tensor self, int dim=-1) -> (Tensor)`` + +参数说明: + +- self: Tensor 类型，shape支持 1-8 维。 + +- dim: Int 类型，默认为 -1。 + +输出说明: + +- 输出为 Tensor，计算公式的最终输出 y。 + +`GroupMatMul `_ +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +函数原型： + +.. code:: python + + npu_grouped_matmul( + x, + weight, + *, + bias=None, + scale=None, + offset=None, + antiquant_scale=None, + antiquant_offset=None, + per_token_scale=None, + group_list=None, + activation_input=None, + activation_quant_scale=None, + activation_quant_offset=None, + split_item=0, group_type=None, + group_list_type=0, + act_type=0, + output_dtype=None, + tuning_config=None + ) -> List[Tensor] + +详细使用方法见标题文档链接 + +FSDP后端融合算子使用方法 +********************************** + +在 ``verl/models/transformers/npu_patch.py`` 目录中，已经把可用的融合算子通过 patch 的形式进行替换，无需进行其他操作即可默认进行使用 + +Megatron后端融合算子使用方法 +********************************** + +Megatron 的融合算子集成在 MindSpeed 中，需要添加特定参数开启: + +1. **Flash Attention（必须开启）** + :: + + +actor_rollout_ref.actor.megatron.override_transformer_config.use_flash_attn=True + ++actor_rollout_ref.ref.megatron.override_transformer_config.use_flash_attn=True + +2. **RotaryMul** + :: + + +actor_rollout_ref.actor.megatron.override_transformer_config.apply_rope_fusion=True + +actor_rollout_ref.actor.megatron.override_transformer_config.use_fused_rotary_pos_emb=True + +3. **RMSNorm** + :: + + +actor_rollout_ref.actor.megatron.override_transformer_config.use_fused_rmsnorm=True + +4. **GroupMatMul** + :: + + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_grouped_gemm=True + +5. **Swiglu** + :: + + +actor_rollout_ref.actor.megatron.override_transformer_config.use_fused_swiglu=True + +6. **Permute/Unpermute** + :: + + +actor_rollout_ref.actor.megatron.override_transformer_config.fused_permute_unpermute=True + +7. **MC2** + :: + + +actor_rollout_ref.actor.megatron.override_transformer_config.use_ascend_mc2=True + +昇腾通用配置 +-------------------------- + +`算子下发 `_ +************************************************************************************************************************************************************************************************************ + +通过 ``TASK_QUEUE_ENABLE`` 可配置 task_queue 算子下发队列优化等级，默认为 Level 1 优化。该配置可以减少host下发时间，可用于缓解由下发导致的整体free过大问题。 + +.. image :: https://github.com/verl-project/verl-data/blob/main/images/ascend/perf_tuning_task_queue.png + :width: 500px + +Level 0 : 不开启下发流水优化。 + +Level 1 : \ 将算子下发任务分为两段，一部分任务（主要是 aclnn 算子的调用）放在新增的二级流水上，一、二级流水通过算子队列传递任务，相互并行，通过部分掩盖减少整体的下发耗时，提升端到端性能。 + +Level 2 : \ 基于 Level 1 的优化进一步平衡了一、二级流水的任务负载，主要是将 workspace 相关任务迁移至二级流水，掩盖效果更好，性能收益更大。该配置仅在二进制场景生效，建议配置值为 Level 2 优化。 + +`通讯算法编排展开 `_ +************************************************************************************************************************************************************************************************************ +使用环境变量 ``HCCL_OP_EXPANSION_MODE=AIV`` 用于配置通信算法的编排展开位置，支持如下取值: + +- **AI_CPU:** 代表通信算法的编排展开位置在 Device 侧的 AI CPU，Device 侧根据硬件型号自动选择相应的调度器。 + +- **AIV:** 代表通信算法的编排展开位置在 Device 侧的 Vector Core，执行也在 Vector Core。 + +- **HOST:** 代表通信算法的编排展开位置为 Host 侧 CPU，Device 侧根据硬件型号自动选择相应的调度器。 + +- **HOST_TS:** 代表通信算法的编排展开位置为 Host 侧 CPU，Host 向 Device 的 Task Scheduler 下发任务，Device 的 Task Scheduler 进行任务调度执行。 + +推理阶段调优 +-------------------------- + +Chunked Prefill in V1 +*************************** + +VLLM 当前版本已默认启用 VLLM V1，使用以下配置启用 Chunked Prefill： + +.. code-block:: sh + + actor_rollout_ref.rollout.enable_chunked_prefill=True + +原理参考 `VLLM 官方文档 `_。 + +Graph Mode +*************************** + +与 CUDA 类似，NPU 通过以下配置启用 **ACL Graph**： + +.. code-block:: sh + + actor_rollout_ref.rollout.enforce_eager=False + +文档：`ACL Graph `_ + +.. note:: + ACL Graph 与 ``taskqueue Level 2`` 原理冲突，**二者无法同时开启**。 + +训练阶段调优 +-------------------------- + +FSDP +********************************** + +.. csv-table:: + :header: "FSDP", "说明" + :widths: 30, 60 + + "/","仅切分优化器(Zero-1)" + SHARD_GRAD_OP,切分梯度和优化器(Zero-2) + "HYBRID_SHARD","切分权重、梯度和优化器(Zero-3)" + "2D device_mesh+HYBRID_SHARD","又称HSDP（FSDP+DDP）例如device_mesh=[2,8], 每8个rank为一个FSDP组，组内进行FSDP切分，共有两个组，两个组间进行DDP，通过allreduce同步梯度。" + "2D device_mesh+HYBRID_SHARD_ZERO2","HSDP的Zero2版本" + NO_SHARD,DDP + +FSDP 不支持 Zero-1， VeRL中会根据卡数和 ``actor_rollout_ref.actor.fsdp_config.fsdp_size`` 来决定 device mesh 的取值，默认使用 Zero-3 进行切分；如果模型较小（建议小于 7B 时），可以通过控制参数 ``actor_rollout_ref.actor.fsdp_config.reshard_after_forward`` 为 ``True`` 在 FSDP/FSDP2 上使用 Zero-2 来优化性能. + +Megatron +********************************** + +在模型较大时，使用 Megatron 作为训练后端可以更灵活的进行性能调优。 + +当 DP 并行显存无法容纳模型时，优先开启 TP 来切分模型权重，如果模型仍然过大，再开启 PP 来进一步切分；如果序列过长导致激活太大，则可以开启 CP 和 SP 来进行优化；在 MoE 模型中则可以额外开启 EP 来控制对专家的切分，如果专家过小，为了避免将权重切的果味细碎，则可以开启 ETP 来避免 MoE 部分的 TP 切分，而将多个完整的专家分布到 DP 和 TP 上。 + +TP、PP、EP、ETP和 Megatron 使用方式一样，CP 和 SP 在 NPU 上开启方式: + +- SP: ``Sequence Parallel`` 在 Tensor Parallel 的基础上进一步提高计算效率，是一种通过将输入数据的序列维度进行切分的并行计算方式。在 NPU 上通过 MindSpeed 来调用SP: + :: + + actor_rollout_ref.actor.megatron.override_transformer_config.sequence_parallel=True + +- CP: ``Context Parallel`` 是一种在多个 GPU/NPU 上并行处理神经网络激活值的方法，他通过在序列维度上对输入张量进行划分来实现。在 NPU 上通过 MindSpeed 来调用 CP （两个参数必须同时添加）: + :: + + actor_rollout_ref.actor.megatron.context_parallel_size + actor_rollout_ref.actor.megatron.override_transformer_config.context_parallel_size + +Megatron-distributed optimizer +********************************** + +在面对较大尺寸模型时，通常需要将优化器分片到一个 DP 域内的每张卡上来节省显存。Megatron 后端下在 NPU 上开启分布式优化器: + +:: + + +actor_rollout_ref.actor.megatron.override_transformer_config.use_distributed_optimizer=True diff --git a/verl/docs/ascend_tutorial/dev_guide/precision_analysis/precision_alignment_zh.md b/verl/docs/ascend_tutorial/dev_guide/precision_analysis/precision_alignment_zh.md new file mode 100644 index 0000000000000000000000000000000000000000..1046e1e51b61034353db681d8ab62c93bfdcb775 --- /dev/null +++ b/verl/docs/ascend_tutorial/dev_guide/precision_analysis/precision_alignment_zh.md @@ -0,0 +1,141 @@ +# Precision Alignment + +在 VeRL 框架中进行强化学习（RL）训练时，**精度对齐**是确保训练过程可复现、可调试的关键环节。 + +本文档总结了在 VeRL 上对NPU和GPU进行精度对齐的方法以供参考。 + +Last updated: 05/09/2026. + +## 1. 环境与权重对齐 + +### 1.1 依赖版本对齐 + +VeRL、transformers版本需要进行强对齐，否则会直接影响精度结果。 + +其他关键依赖（torch、megatron、vllm）如无法进行强对齐，需优先保持一致或相近。 + +### 1.2 模型权重对齐 + +检查模型的权重和config.json文件是否完全一致 + + +## 2. 输入数据对齐 + +在verl训练启动脚本中添加如下配置： + +```bash +data.shuffle=False +data.validation_shuffle=False +``` + + +## 3. 配置对齐 + +在NPU与GPU做精度对齐时，需检查配置是否完全对齐。包含： +1. 直接对比脚本写入配置 +2. 运行过程中保存日志，收集日志打屏中的配置进行对比，可比较默认参数配置是否一致，需保证关键参数对齐 + + +## 4. 固定确定性 + +### 4.1 固定随机种子 + +在环境中安装 `msprobe` : + +```bash +pip install mindstudio-probe +``` + +在 worker 文件开头添加确定性函数： + +```python +from msprobe.pytorch import seed_all +seed_all(mode=True) +``` + +### 4.2 固定通信环境变量 + +在多卡通信情况下： + +- HCCL通信下(默认场景): + + - export CLOSE_MATMUL_K_SHIFT=1 + - export ATB_MATMUL_SHUFFLE_K_ENABLE=0 + - export HCCL_DETERMINISTIC="true" + - export VLLM_ENABLE_V1_MULTIPROCESSING=0 + +- LCCL通信下(通过export HCCL_OP_EXPANSION_MODE="AIV"使能): + + - export CLOSE_MATMUL_K_SHIFT=1 + - export ATB_MATMUL_SHUFFLE_K_ENABLE=0 + - export LCCL_DETERMINISTIC=1 + - export ATB_LLM_LCOC_ENABLE=0 + - export VLLM_ENABLE_V1_MULTIPROCESSING=0 + +在单卡无通信情况下： + + - export CLOSE_MATMUL_K_SHIFT=1 + - export ATB_MATMUL_SHUFFLE_K_ENABLE=0 + - export VLLM_ENABLE_V1_MULTIPROCESSING=0 + + + +## 5. 验证训练精度 + +### 5.1 训练打桩 + +**打桩**即保留当前阶段的输入输出数据，便于从结果上对比分析。在进行精度问题排查时，需要进行打桩辅助问题定位。常见的打桩方式是直接将rollout阶段的数据直接dump下来。 + +**第一步：在 GPU 环境生成基准数据** + +先跑一次GPU脚本，开启如下配置： + +```bash +trainer.rollout_data_dir='/path/dump/data_json' +``` +可以保存每步推理结果为jsonl文件。 + +**第二步：在 NPU 环境复现验证** + +NPU上开启如下参数，复用上一步生成的序列，端到端运行： + +```bash +actor_rollout_ref.rollout.skip_rollout=True \ +actor_rollout_ref.rollout.skip_dump_dir="/path/dump/data_json" \ +``` + +**第三步：对比指标** + +在打桩输入相同推理结果，训练配置保持一致，并且固定随机性的情况下，比较NPU与GPU的rewards/pg_loss/grad_norm值是否存在差异。 + + +## 6. 验证推理精度 + +### 6.1 resharding + +在推理正式开始前，vllm会进行**dummy run**，通过推理一个 token 来评估推理时的显存占用，进而分配显存。可以在 vLLM 的 LLM 初始化时指定参数 load_format 来指定 dummy run 的权重是随机初始化的（dummy）还是真实权重（safetensors）。在 VeRL 中，通过参数 **actor_rollout_ref.rollout.load_format** 指定该参数。 + +当出现推理乱码现象时，如果引擎初始化方式为**load_format=dummy**，则sharding高概率存在问题，即使换成了safetensors后吐字正常，sharding也是存在问题的，需要对比前向。 + + +### 6.2 推理结果对齐 + +```bash +trainer.rollout_data_dir='/path/dump/data_json' +``` + +保存每步推理结果为jsonl文件，可以直接打开jsonl文件快速确认整网推理结果是否乱码，用于推理精度问题定界。 + + +在dump推理数据之前，若复现推理精度问题占用的资源较多，可以先尝试缩小推理精度问题复现成本，减少复现的规模，减少需要dump和对比的数据。在多batch、长序列的场景下，可以通过发送单batch请求，减少序列长度尝试复现。 + + +## 7. dump对比 + +[精度调试工具](./precision_debugger_zh.md)，定位到问题出现的阶段之后，可以通过msprobe工具进行数据dump来细致定位。 + +在推理或训练过程中，模型可能出现输出偏离预期、生成异常、甚至产生 NaN/Inf 等数值不稳定问题。要定位根因，需要对模型执行路径进行精细化监控，采集中间特征、权重、激活值以及各关键层的输入输出，并记录提示词、张量 dtype、硬件配置等上下文信息。通过捕获这些核心张量及元数据，可以系统性地追踪精度退化或数值错误的来源。 + + + + diff --git a/verl/docs/ascend_tutorial/dev_guide/precision_analysis/precision_debugger_zh.md b/verl/docs/ascend_tutorial/dev_guide/precision_analysis/precision_debugger_zh.md new file mode 100644 index 0000000000000000000000000000000000000000..ebb97e875e513f3d2139d56a58192f09fd427036 --- /dev/null +++ b/verl/docs/ascend_tutorial/dev_guide/precision_analysis/precision_debugger_zh.md @@ -0,0 +1,373 @@ +# Precision Debugger (msprobe) in verl + +Last updated: 04/13/2026. + +This guide explains how to collect precision data in verl using the +`msprobe` PrecisionDebugger. + +## Prerequisites + +* Install `msprobe` in the training environment: + +```bash +pip install mindstudio-probe +``` + +* Prepare a `config.json` for msprobe (see examples below). +* Enable profiler for the roles you want to collect. + +Reference: +* `https://gitcode.com/Ascend/msprobe.git` + +## Configuration + +PrecisionDebugger is integrated through verl's unified profiler interface. +Use a minimal two-part setup: + +* `global_profiler` selects the tool and config file. +* role `profiler.enable=True` turns on profiling for that role. + +### Global profiling control + +In `global_profiler`, set the profiler tool to `precision_debugger` and +configure the msprobe-specific options under `global_tool_config`. + +```yaml +global_profiler: + tool: precision_debugger + steps: [1, 2, 5] + save_path: "outputs/profile" + global_tool_config: + precision_debugger: + _target_: verl.utils.profiler.config.PrecisionDebuggerToolConfig + config_path: /path/to/config.json + stages: + - actor_update + - actor_compute_log_prob + - ref_compute_log_prob + - compute_values + - critic_update + - compute_rm_score + strict: False +``` + +Notes: + +* `global_profiler.steps` is the only step filter for PrecisionDebugger. +* Dumps are written under `global_profiler.save_path`. +* Actual dump path is `{global_profiler.save_path}/step_{global_step}/{stage}`. +* Do not set `dump_path` in `config.json`; output path is controlled by verl. + +### Role profiling control + +Enable profiling for the roles you want to collect: + +```yaml +actor_rollout_ref: + actor: + profiler: + enable: True + ref: + profiler: + enable: True +critic: + profiler: + enable: True +``` + +## Supported stages + +PrecisionDebugger collects data from the following stages: + +* `actor_update` +* `actor_compute_log_prob` +* `ref_compute_log_prob` +* `compute_values` +* `critic_update` +* `compute_rm_score` + +Rollout generation is intentionally skipped (`rollout_generate` is ignored). + +The current integration is designed for training-side stages. In a typical PPO +run, the most common useful combinations are: + +* actor/ref only: + `actor_compute_log_prob`, `ref_compute_log_prob`, `actor_update` +* actor/ref/critic: + `actor_compute_log_prob`, `ref_compute_log_prob`, `compute_values`, + `critic_update`, `actor_update` + +## msprobe config.json common examples + +### `statistics` mode + +```json +{ + "task": "statistics", + "rank": [], + "step": [], + "level": "L1", + "async_dump": false, + "statistics": { + "scope": [], + "list": [], + "tensor_list": [], + "data_mode": ["all"], + "summary_mode": "statistics" + } +} +``` + +### `tensor` mode + +```json +{ + "task": "tensor", + "rank": [], + "step": [], + "level": "L1", + "async_dump": false, + "tensor": { + "scope": [], + "list": [], + "data_mode": ["all"], + "summary_mode": "statistics" + } +} +``` + +## Minimal example + +The following example enables PrecisionDebugger on steps `1` and `2`. +If you need rank filtering, configure it only in msprobe `config.json`. + +```yaml +global_profiler: + tool: precision_debugger + steps: [1, 2] + global_tool_config: + precision_debugger: + _target_: verl.utils.profiler.config.PrecisionDebuggerToolConfig + config_path: /path/to/dump_config.json + stages: + - actor_compute_log_prob + - ref_compute_log_prob + - actor_update + strict: False + +actor_rollout_ref: + actor: + profiler: + enable: True + ref: + profiler: + enable: True +``` + +## Minimal CLI example + +Use only the required flags: + +```bash +python3 -m verl.trainer.main_ppo \ + global_profiler.tool=precision_debugger \ + global_profiler.steps='[1,2]' \ + global_profiler.save_path=outputs/profile \ + +global_profiler.global_tool_config.precision_debugger.config_path=/path/to/config.json \ + actor_rollout_ref.actor.profiler.enable=True \ + actor_rollout_ref.ref.profiler.enable=True +``` + +Optional stage filter: + +```bash ++global_profiler.global_tool_config.precision_debugger.stages='[actor_compute_log_prob,ref_compute_log_prob,actor_update]' +``` + +## Output layout + +Verl organizes PrecisionDebugger output by training global step and stage. +Inside each stage directory, msprobe creates its own `step*/rank*` layout. + +Example: + +```text +outputs/profile/ + step_1/ + actor_compute_log_prob/step0/rank0/dump.json + actor_update/step0/rank0/dump.json + ref_compute_log_prob/step0/rank0/dump.json + step_2/ + actor_compute_log_prob/step0/rank0/dump.json + actor_update/step0/rank0/dump.json + ref_compute_log_prob/step0/rank0/dump.json +``` + +Observed output from a real run: + +* Outer `step_` directories are created by verl. +* Inner `step0/rank0/dump.json` directories are created by msprobe. +* With the current integration, each profiled stage is collected in an + independent dump session, so stage-local output typically lands in `step0`. + +## How results are written + +The verl integration wraps each profiled stage with: + +* `debugger.start(model=...)` +* execute the stage +* `debugger.stop()` +* `service.reset_status()` if the msprobe runtime exposes it + +Verl does **not** manually call `debugger.step()` in the current integration. +Instead, each stage writes to its own dump directory and resets msprobe runtime +status after `stop()` to avoid stale `dump.json` cache growth across stages. + +For L0 collection, PrecisionDebugger must bind to the actual model used in the +stage. The profiler resolves the model inside +`verl/utils/profiler/precision_debugger_profile.py` and supports both legacy +workers and the newer model-engine worker path. + +## Overhead and disk usage + +Below are measurements from a real PPO run on Ascend with: + +* model: `Qwen2-0.5B` +* profiled steps: `[1, 2]` +* rank: `0` +* stages: + * L1: `actor_compute_log_prob`, `ref_compute_log_prob`, `actor_update` + * L0: `actor_compute_log_prob`, `ref_compute_log_prob`, `compute_values`, + `critic_update`, `actor_update` + +### Time overhead + +| Run | Model | Profiled steps | Measured step time | +|---|---|---:|---:| +| Baseline | `Qwen2-0.5B` | None | about `16-18 s/step` in steady state | +| L0 | `Qwen2-0.5B` | `step 1` | `66.81 s` | +| L0 | `Qwen2-0.5B` | `step 2` | `48.78 s` | +| L0 | `Qwen2-0.5B` | non-profiled later steps | about `17 s/step` | +| L1 | `Qwen2-0.5B` | `step 1` | `177.35 s` | +| L1 | `Qwen2-0.5B` | `step 2` | `161.80 s` | +| L1 | `Qwen2-0.5B` | non-profiled later steps | about `17 s/step` | + +In this experiment, profiled L0 steps were about `3x-4x` slower than the +baseline steady-state step time, and profiled L1 steps were about `9x-10x` +slower. Non-profiled later steps remained close to baseline in both cases. + +In general, PrecisionDebugger should be treated as a heavy-weight precision +debugging tool rather than a lightweight profiler. In larger models or broader +stage coverage, it is common to observe `tens-X` performance inflation for +profiled steps. + +### Disk usage + +| Level | Model | Stages | Scope | Disk usage | +|---|---|---|---|---:| +| L1 | `Qwen2-0.5B` | `actor_compute_log_prob`, `ref_compute_log_prob`, `actor_update` | total for `step_1` and `step_2` | `21 MB` | +| L1 | `Qwen2-0.5B` | `actor_compute_log_prob`, `ref_compute_log_prob`, `actor_update` | per step | about `11 MB` | +| L1 | `Qwen2-0.5B` | `actor_update` | per step | about `5.1-5.2 MB` | +| L1 | `Qwen2-0.5B` | `actor_compute_log_prob` | per step | about `2.6 MB` | +| L1 | `Qwen2-0.5B` | `ref_compute_log_prob` | per step | about `2.6 MB` | +| L0 | `Qwen2-0.5B` | `actor_compute_log_prob`, `ref_compute_log_prob`, `actor_update` | total for `step_1` and `step_2` | `8.8 MB` | +| L0 | `Qwen2-0.5B` | `actor_compute_log_prob`, `ref_compute_log_prob`, `actor_update` | per step | about `4.4 MB` | +| L0 | `Qwen2-0.5B` | `actor_update` | per step | about `2.5 MB` | +| L0 | `Qwen2-0.5B` | `actor_compute_log_prob` | per step | about `1.1 MB` | +| L0 | `Qwen2-0.5B` | `ref_compute_log_prob` | per step | about `0.86-0.87 MB` | + +In this experiment, total L1 disk usage was about `2.4x` the L0 disk usage for +the measured actor/ref stage set. + +These numbers depend on: + +* selected stages +* number of profiled steps +* dump level and task +* model shape and sequence length + +## How to analyze results + +At minimum, check: + +* which `step_` directory was generated +* which stage directories exist under that step +* whether `dump.json` exists under `step0/rank0` + +For downstream analysis, use standard msprobe tools such as: + +* `msprobe compare` +* `msprobe visualization` + +Example compare usage: + +```bash +msprobe compare \ + --target-path /path/to/target_dump/dump.json \ + --golden-path /path/to/golden_dump/dump.json +``` + +You can compare: + +* the same stage across two runs +* different global steps of the same stage +* different ranks when multi-rank collection is enabled + +For more advanced analysis workflows, refer to the official msprobe +documentation for compare and visualization commands. + +## Usage notes + +* Verl integrates PrecisionDebugger through `DistProfiler.annotate` wrappers on + training stages. +* PrecisionDebugger is automatically discrete: each profiled stage is + collected in an independent `start -> stop -> reset_status` session. It does + not currently expose the unified profiler `discrete` configuration used by + tools such as `nsys` or `npu`. +* `global_steps` is read from batch `meta_info` or from worker attributes. +* If `strict` is `True`, missing msprobe or unknown stages raise errors. +* If a stage prints `PrecisionDebugger model not resolved`, that stage ran + normally but no dump was collected because verl could not bind msprobe to a + valid model object. +* Because dump cost is high, prefer collecting a small number of representative + steps first, then narrow the stage set if necessary. + +## Quality checklist + +Use this checklist to verify your setup is complete and reproducible: + +* `global_profiler.tool=precision_debugger` +* `global_profiler.steps` includes the target step +* `+global_profiler.global_tool_config.precision_debugger.config_path=...` is set +* role `profiler.enable=True` is set for the stages you need +* `msprobe` is importable in the runtime environment +* output exists under `{global_profiler.save_path}/step_//...` + +## Troubleshooting + +### No dump directory is generated + +Check: + +* `global_profiler.tool=precision_debugger` +* `global_profiler.steps` contains the target step +* role profiler is enabled for the target role +* msprobe is installed in the training environment + +### `PrecisionDebugger model not resolved` + +This means the stage was reached, but verl could not find the actual model used +by that worker. The stage itself still runs, but dump is skipped. This usually +indicates: + +* a new worker path was introduced and profiler model resolution needs to be + updated +* the role or engine backend differs from the paths currently supported by the + resolver + +### `dump.json` keeps growing unexpectedly + +If `stop()` is called without resetting msprobe runtime state, cached dump data +may continue to accumulate across stage invocations. The current verl +integration resets msprobe runtime status after `stop()` when the service API +supports it. diff --git a/verl/docs/ascend_tutorial/faq/faq.rst b/verl/docs/ascend_tutorial/faq/faq.rst new file mode 100644 index 0000000000000000000000000000000000000000..23a130c681451657a4b5673f1a90869202f6800f --- /dev/null +++ b/verl/docs/ascend_tutorial/faq/faq.rst @@ -0,0 +1,205 @@ +NPU 常见问题解答 +================ + +Last updated: 05/13/2026. + +本文档总结了在 NPU 上执行 VERL 训练和推理时遇到的常见问题及解决方案。 + +环境配置问题 +------------ + +### Q1: NPU 设备不可见怎么办？ + +**问题现象**：torch_npu.npu.is_available() 返回 False + +**解决方案**： + +.. code-block:: bash + + # 检查设备可见性 + echo $ASCEND_RT_VISIBLE_DEVICES + + # 设置可见设备并禁用ray自动设置 + export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 + export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 + + # 检查驱动状态 + npu-smi info + +调试和诊断 +---------- + +### Q1: 如何启用 NPU 性能分析？ + +使用 VERL 内置的 profiler： + +.. code-block:: shell + + actor_rollout_ref.actor.profiler.tool_config.npu.discrete=true \ + actor_rollout_ref.actor.profiler.tool_config.npu.contents=npu,cpu \ + actor_rollout_ref.actor.profiler.tool_config.npu.level=1 \ + actor_rollout_ref.actor.profiler.tool_config.npu.analysis=true + +### Q2: 如何排查 NPU 训练失败的问题？ + +**排查步骤**： + +1. 检查环境变量配置 +2. 验证设备可见性 +3. 检查 CANN 版本兼容性 +4. 查看日志中的具体错误信息 +5. 使用最小化示例复现问题 + +**启用详细日志**： + +.. code-block:: bash + + # VERL 框架日志 + export VERL_LOGGING_LEVEL=DEBUG + + # 昇腾 NPU 日志（0=DEBUG, 1=INFO, 2=WARNING, 3=ERROR） + export ASCEND_GLOBAL_LOG_LEVEL=0 + export ASCEND_SLOG_PRINT_TO_STDOUT=1 + + # HCCL 通信日志 + export HCCL_DEBUG=INFO + +常见错误信息 +------------ + +### Q1: "torch_npu detected, but NPU device is not available or visible" + +**原因**：NPU 驱动未正确安装或设备不可见 + +**解决方案**：检查驱动安装状态和 ASCEND_RT_VISIBLE_DEVICES 设置 + +### Q2: "KeyError: decoder.layers.0.self_attention.q_layernorm.weight" + +**原因**：MindSpeed版本过低 + +**解决方案**：切换MindSpeed至 2.3.0_core_r0.12.1 + +### Q3: "AssertionError: Weight ... is too large to fit in the bucket" + +**问题现象**：在分布式训练权重同步时，出现如下错误： + +.. code-block:: text + + AssertionError: Weight model.embed_tokens.weight(torch.Size([151936, 4096]), torch.float32) is too large to fit in the bucket. + Please increase rollout.update_weights_bucket_megabytes(2048 MB). + +**原因**：模型某个权重张量的大小超过了权重传输 bucket 的默认容量（2048 MB）。在 verl 框架中，模型权重通过 bucket（缓冲区）进行分块打包传输。当单个权重张量超过 bucket 大小时，断言检查失败。 + +**权重大小计算方法**： + +权重张量的内存占用（字节）= 各维度大小的乘积 × 每个元素的字节数 + +其中数据类型对应的字节数为： + +- ``torch.float32`` → 4 字节 +- ``torch.float16`` / ``torch.bfloat16`` → 2 字节 +- ``torch.int8`` → 1 字节 + +以本例中的 ``model.embed_tokens.weight`` 为例： + +.. code-block:: text + + 张量形状: torch.Size([151936, 4096]) + 数据类型: torch.float32 (4 字节) + 权重大小 = 151936 × 4096 × 4 = 2,483,027,968 字节 ≈ 2369 MB + + 默认 bucket 大小 = 2048 MB < 2369 MB → 触发断言失败 + +**解决方案**：在启动训练时增加 ``update_weights_bucket_megabytes`` 参数，使 bucket 容量大于最大权重张量的内存占用： + +.. code-block:: bash + + actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=4096 + +**参数值选择建议**： + +1. **计算模型中最大权重张量的内存占用**：遍历模型所有参数，找出 ``nbytes`` 最大的那个，将其转换为 MB（除以 1024²）。 + +2. **向上取整到 2 的幂次**：为便于内存分配和管理，建议将计算结果向上取整到最近的 2 的幂次（如 2048、4096、8192 等）。例如最大权重为 2369 MB，则取 4096 MB。 + +3. **预留适当余量**：考虑到内存对齐和运行时开销，建议 bucket 大小至少为最大权重大小的 1.2~1.5 倍，再向上取整到 2 的幂次。 + +4. **注意内存限制**：bucket 大小会直接影响 worker 节点的内存占用，设置过大会导致 OOM。应在满足权重传输需求的前提下，尽量选择较小的值。 + +**常见模型的推荐值**： + +.. list-table:: + :header-rows: 1 + + * - 模型规模 + - 典型最大权重形状 + - 推荐 bucket 大小 + * - 7B (Qwen2 等) + - [151936, 4096] float32 + - 4096 MB + * - 14B + - [152064, 5120] float32 + - 4096 MB + * - 72B + - [152064, 8192] float32 + - 8192 MB + +### Q4: 非共享存储下 checkpoint 加载失败，找不到 common.pt / .metadata / metadata.json + +**问题现象**：使用 verl + Megatron 后端在**非共享存储**的多机环境下，保存 checkpoint 正常，但重新加载时报错，提示找不到以下文件： + +.. code-block:: text + + FileNotFoundError: common.pt + FileNotFoundError: .metadata + FileNotFoundError: metadata.json + +**原因**：当前 checkpoint 机制对非共享存储的支持不完善。具体表现为： + +- **分布式训练权重是分节点保存的**，每个节点只保存自己负责的分片权重，不会只在主节点保存全部权重。 +- 但 ``common.pt``、``.metadata``、``metadata.json`` 等元数据文件**仅保存在执行保存操作的节点上**（通常是 rank 0 所在节点），其他节点本地没有这些文件。 +- 加载 checkpoint 时，每个节点都需要读取这些元数据文件来还原模型状态，但非共享存储下其他节点本地路径中不存在这些文件，导致加载失败。 + +**临时解决方案**：手动将元数据文件从保存节点复制到所有其他节点： + +.. code-block:: bash + + # 假设 checkpoint 保存在 rank 0 节点的 /path/to/ckpt/ 目录下 + # 将元数据文件从 rank 0 节点复制到其他所有节点 + + # 需要复制的文件 + /path/to/ckpt/common.pt + /path/to/ckpt/.metadata + /path/to/ckpt/metadata.json + + # 示例：使用 scp 复制到其他节点 + scp /path/to/ckpt/common.pt node1:/path/to/ckpt/ + scp /path/to/ckpt/.metadata node1:/path/to/ckpt/ + scp /path/to/ckpt/metadata.json node1:/path/to/ckpt/ + + # 对所有节点重复上述操作 + +**注意事项**： + +- 每次保存 checkpoint 后都需要重新复制元数据文件，因为保存操作可能会更新这些文件的内容。 +- 如果训练过程中频繁保存 checkpoint（如按步数自动保存），建议编写脚本在保存后自动触发复制，避免遗漏。 +- 长期方案应等待框架层面支持非共享存储的 checkpoint 加载，使元数据文件能自动同步到所有节点。 + +参考资料 +-------- + +- `NPU 性能优化指南 <../perf/perf_tuning_on_ascend.rst>`_ +- `NPU 快速开始指南 <../start/install.rst>`_ +- `NPU CI 指南 <../contribution_guide/ascend_ci_guide_zh.rst>`_ +- Ascend NPU 文档: https://www.hiascend.com/document +- CANN 工具包文档: https://www.hiascend.com/software/cann + +获取更多帮助 +------------ + +如果以上 FAQ 无法解决您的问题，请： + +1. 查看完整的错误日志 +2. 在 GitHub Issues 中搜索类似问题 +3. 提供详细的错误信息和环境配置 +4. 提供最小可复现示例 \ No newline at end of file diff --git a/verl/docs/ascend_tutorial/feature_support/ascend_backend_features.md b/verl/docs/ascend_tutorial/feature_support/ascend_backend_features.md new file mode 100644 index 0000000000000000000000000000000000000000..3fb7b1b2e1a5026360e0fc297068436494dcce8c --- /dev/null +++ b/verl/docs/ascend_tutorial/feature_support/ascend_backend_features.md @@ -0,0 +1,276 @@ +# Ascend Backend Features Guide +================================================================================== + +Last updated: 03/03/2026. + +昇腾全面支持verl生态建设，本文将介绍NPU上对于verl的适配工作及后端特性支持供开发者进行参考 + +--- + +## 推理后端 + +当前verl支持vllm/sglang这两种主流推理后端，均可在昇腾NPU上运行。 + +### 1. vllm: + +昇腾通过vllm-ascend插件来支持vllm推理后端，该插件是 vLLM 社区支持 Ascend 后端的推荐方法。它遵循[[RFC]](https://github.com/vllm-project/vllm/issues/11162)，提供了一个可插拔接口，将 Ascend NPU 与 vLLM 解耦。 + +##### 参数特性支持 + +| vllm参数| verl对应通用参数 | 简介| +| --- | --- | --- | +| `model_path` | `actor_rollout_ref.model.path` |模型权重文件的路径| +| `gpu_memory_utilization` | `actor_rollout_ref.rollout.gpu_memory_utilization` |用于控制每个阶段可使用的 GPU 内存量。它被指定为一个介于 0.0 和 1.0 之间的分数，其中：- 0.8 表示 GPU 总内存的 80%- 1.0 表示 GPU 总内存的 100%（不推荐，没有预留缓冲）| +| `enforce_eager`| `actor_rollout_ref.rollout.enforce_eager` |禁用图模式，verl默认为False| +| `enable_chunked_prefill`| `actor_rollout_ref.rollout.enable_chunked_prefill` | 分块预填充允许将大预填充分块成更小的块，并将它们与解码请求一起批处理。| +| `free_cache_engine`| `actor_rollout_ref.rollout.free_cache_engine` |在部署生成阶段之后卸载 KVCache，默认值为 True。| +| `max_model_len` | `actor_rollout_ref.rollout.max_model_len` | 模型能够处理的最大序列长度。它限制了单个输入序列的最大长度 | +| `tp_size`| `actor_rollout_ref.rollout.tensor_model_parallel_size * data_parallel_size`|TP并行度| +| `dp_size`| `actor_rollout_ref.rollout.data_parallel_size`|DP并行度| +| `ep_size`| `actor_rollout_ref.rollout.expert_parallel_size`|EP并行度| +| `node_rank`| `无，根据实际实例和卡数自动计算` |实例中的节点排序| +| `load_format`| `actor_rollout_ref.rollout.load_format` |要加载的模型权重格式| +| `disable_log_stats`| `actor_rollout_ref.rollout.disable_log_stats`|控制是否记录 rollout 统计日志 | +| `nnodes `| `无，根据实际实例和卡数自动计算` | 每个实例包含的节点数量` | +| `trust_remote_code`| `actor_rollout_ref.model.trust_remote_code`|是否允许在 Hub 上定义自定义模型，并将其写入自己的建模文件中| +| `max_num_seqs` | `actor_rollout_ref.rollout.max_num_seqs` |正在运行的请求的最大数量| +| `max_num_batched_tokens`| `actor_rollout_ref.rollout.max_num_batched_tokens` |在一次批处理（batch）中可以处理的最大总Token数| +| `skip_tokenizer_init`| `actor_rollout_ref.rollout.skip_tokenizer_init` |跳过初始化分词器并将 input_ids 传递到推理请求中| +| `enable_prefix_caching` | `actor_rollout_ref.rollout.enable_prefix_caching`|`用于启用自动前缀缓存` | +| `quantization`| `actor_rollout_ref.rollout.quantization，默认为None`|`量化方法`| +| `enforce_eager`|`actor_rollout_ref.rollout.enforce_eager`|标志用于强制使用PyTorch的eager执行模式，而非默认的图执行模式| + +### 2. sglang: + +对于sglang推理后端，昇腾通过直接向sglang社区进行持续建设与维护来支持相关功能。 +此外在verl中使用sglang还涉及以下组件, 我们在[quick start](https://github.com/verl-project/verl/blob/main/docs/ascend_tutorial/quick_start/ascend_sglang_quick_start.rst)中提供详细说明与一键安装脚本。 + +| 组件| 描述| +| --- | --- | +| [sgl_kernel_npu](https://github.com/sgl-project/sgl-kernel-npu/blob/main/python/sgl_kernel_npu/README.md) | Ascend NPU SGL 优化推理内核集合，包括注意力机制、归一化、激活函数、LoRA 适配器等。 | +| [deepep](https://github.com/sgl-project/sgl-kernel-npu/blob/main/python/deep_ep/README.md) | DeepEP的 Ascend 实现，为MoE模型提供高度优化的专家并行 (EP) 通信内核 | + +##### 参数特性支持 + +verl中通过rollout config管理推理后端参数使能，包含通用参数和engine_kwargs自定义传参。 +以下列举在verl中常见设置的sglang特性参数，更多参数介绍请参考 [sglang社区NPU特性支持](https://docs.sglang.io/platforms/ascend_npu_support_features.html) + +| sglang参数| verl对应通用参数 | 简介| +| --- | --- | --- | +| model_path | actor_rollout_ref.model.path|模型权重文件的路径| +| mem_fraction_static| actor_rollout_ref.rollout.gpu_memory_utilization |用于静态分配（模型权重和键值缓存内存池）的内存比例| +| disable_cuda_graph| actor_rollout_ref.rollout.enforce_eager|禁用图模式，verl默认为False| +| enable_memory_saver| 无，verl中默认设置为True | 允许使用 release_memory_occupation 和 resume_memory_occupation 来节省内存 +| base_gpu_id| 无，根据实际实例和卡数自动计算 |用于分配每个实例上计算卡资源时的的初始ID +| gpu_id_step| 无，默认设置为1| 使用的连续计算卡ID 之间的差值 +| tp_size| actor_rollout_ref.rollout.tensor_model_parallel_size * data_parallel_size|TP并行度| +| dp_size| actor_rollout_ref.rollout.data_parallel_size|DP并行度| +| ep_size| actor_rollout_ref.rollout.expert_parallel_size|EP并行度| +| node_rank| 无，根据实际实例和卡数自动计算 |实例中的节点排序| +| load_format| actor_rollout_ref.rollout.load_format|要加载的模型权重格式| +| dist_init_addr| 无，自动计算|用于初始化分布式后端的主机地址| +| nnodes| 无，根据实际实例和卡数自动计算|每个实例包含的节点数量| +| trust_remote_code| actor_rollout_ref.model.trust_remote_code|是否允许在 Hub 上定义自定义模型，并将其写入自己的建模文件中| +| max_running_requests| actor_rollout_ref.rollout.max_num_seqs |正在运行的请求的最大数量| +| log_level| 无，默认设置为error |日志记录器的日志级别| +| skip_tokenizer_init| actor_rollout_ref.rollout.skip_tokenizer_init |跳过初始化分词器并将 input_ids 传递到推理请求中| +| skip_server_warmup| 无，默认设置为True |跳过预热| +| quantization| actor_rollout_ref.rollout.quantization，默认为None|量化方法| +| attention_backend|actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend|attention内核,NPU应该设置为ascend| + +--- + +## 训练后端 + +### 1. FSDP + +昇腾通过torch_npu提供FSDP相关支持能力，当前pytorch api支持度参照[版本说明](https://www.hiascend.com/document/detail/zh/Pytorch/730/apiref/PyTorchNativeapi/docs/zh/native_apis/pytorch_2-7-1/torch-distributed-fsdp.md)。 + +#### FSDP1 +##### 参数特性支持 +| verl参数 | 简介| +| --- | --- | +| `actor_rollout_ref.actor.fsdp_config.param_offload` |是否卸载模型权重到CPU，默认值为False| +| `actor_rollout_ref.actor.fsdp_config.optimizer_offload` |是否卸载优化器状态到CPU，默认值为False| +| `actor_rollout_ref.actor.fsdp_config.reshard_after_forward` |控制前向计算后的参数行为，平衡内存与通信。默认值为True：前向后重新分片参数，反向时重新全收集| +| `actor_rollout_ref.actor.fsdp_config.fsdp_size` | 每个FSDP分片组中的NPU数量；默认值-1表示自动。| +| `actor_rollout_ref.actor.fsdp_config.forward_prefetch` |在前向计算完成前预取下一次前向传播的 all-gather，仅用于FSDP1，默认值为False| +| `actor_rollout_ref.actor.fsdp_config.use_orig_params` | FSDP是否会使用module的原始参数来初始化，仅用于FSDP1，默认值为False| +| `actor_rollout_ref.actor.ulysses_sequence_parallel_size`|Ulysses序列并行大小| +| `actor_rollout_ref.actor.entropy_from_logits_with_chunking`|通过分块计算熵以减少显存峰值，默认值为False| +| `actor_rollout_ref.actor.fsdp_config.entropy_checkpointing`|在训练时对熵计算启用重计算,降低显存峰值，默认值为False| +| `actor_rollout_ref.actor.fsdp_config.forward_only` |是否只进行前向计算，默认值为False| + +#### FSDP2 +##### 参数特性支持 +| verl参数 | 简介| +| --- | --- | +| `actor_rollout_ref.actor.fsdp_config.param_offload` |是否卸载模型权重到CPU，默认值为False| +| `actor_rollout_ref.actor.fsdp_config.optimizer_offload` |是否卸载优化器状态到CPU，默认值为False| +| `actor_rollout_ref.actor.fsdp_config.reshard_after_forward` |控制前向计算后的参数行为，平衡内存与通信。默认值为True：前向后重新分片参数，反向时重新全收集| +| `actor_rollout_ref.actor.fsdp_config.fsdp_size` | 每个FSDP分片组中的NPU数量；默认值-1表示自动。| +| `actor_rollout_ref.actor.ulysses_sequence_parallel_size`|Ulysses序列并行大小| +| `actor_rollout_ref.actor.entropy_from_logits_with_chunking`|通过分块计算熵以减少显存峰值，默认值为False| +| `actor_rollout_ref.actor.fsdp_config.entropy_checkpointing`|在训练时对熵计算启用重计算,降低显存峰值，默认值为False| +| `actor_rollout_ref.actor.fsdp_config.forward_only` |是否只进行前向计算，默认值为False| + + + +### 2. Megatron + +Megatron 是 NVIDIA 推出的一个专注于模型并行的训练框架仓库。如果一个仓库（例如 Verl）的训练后端使用了 Megatron，同时又希望在 NPU 上运行该仓库，那么就需要额外安装 MindSpeed 来提供底层支持。下文将介绍 MindSpeed 是如何实现无感替换 Megatron 中的关键组件，从而使其能够适配 NPU 的。 + +MindSpeed 底层的替换原理采用了 Monkey Patch 技术 + +* MindSpeed Moneky Patch框架 + +在verl里面通过`from mindspeed.megatron_adaptor import repatch `触发patch，调用栈如下： + +~~~ +from mindspeed.megatron_adaptor import repatch +├── 执行 megatron_adaptor.py 模块导入 +├── 导入 features_manager 模块 +├── 执行 mindspeed/features_manager/__init__.py +├── @AutoExecuteFunction 装饰器触发 +├── patch_features() 自动执行 +└── 进行`apply_features_pre_patches`和`apply_features_patches`操作 +~~~ + +`Patch`类是整个patch系统的核心，实现了函数/类的动态替换 + +~~~python +class Patch +~~~ + +`parse_path`方法实现了动态模块导入和创建 + +~~~python +def parse_path(module_path, function_name, create_dummy) +~~~ + +patch系统支持多层装饰器叠加 + +~~~ +def apply_patch(self): + final_patch_func = self.orig_func + if self.patch_func is not None: + final_patch_func = self.patch_func + + # 应用所有装饰器 + for wrapper in self.wrappers: + final_patch_func = wrapper(final_patch_func) +~~~ + +* MindSpeedPatchesManager类 + +`MindSpeedPatchesManager`作为全局单例管理所有patch + +~~~python +class MindSpeedPatchesManager: + patches_info: Dict[str, Patch] = {} +~~~ + +* Feature集成模式 + +各个Feature通过继承`MindSpeedFeature`基类集成patch系统 + +~~~python +class MindSpeedFeature: + """Base class for mindspeed features.""" + + def __init__(self, feature_name: str, optimization_level: int = 2): + self.feature_name = feature_name.lower().strip().replace('-', '_') + self.optimization_level = optimization_level + self.default_patches = self.optimization_level == 0 + + def is_need_apply(self, args): + """Check the feature is need to apply.""" + return (self.optimization_level <= args.optimization_level and getattr(args, self.feature_name, None)) \ + or self.default_patches + + def register_args(self, parser: ArgumentParser): + """Register cli arguments to enable the feature.""" + pass + + def pre_validate_args(self, args: Namespace): + """Validate the arguments of mindspeed before megatron args validation + and store some arguments of the mindspeed temporarily, + incase that megatron validate faile. + for example: + ```python + origin_context_parallel_size = args.context_parallel_size + args.context_parallel_size = 1 + ``` + """ + pass + + def validate_args(self, args: Namespace): + """Restore the arguments of the mindspeed. + + for example: + ```python + args.context_parallel_size = origin_context_parallel_size + ``` + """ + pass + + def post_validate_args(self, args: Namespace): + """validate mindspeed arguments after megatron arguments validation.""" + pass + + def pre_register_patches(self, patch_manager: MindSpeedPatchesManager, args: Namespace): + """Register all patch functions before import megatron""" + pass + + def register_patches(self, patch_manager: MindSpeedPatchesManager, args: Namespace): + """Register all patch functions the feature is related.""" + pass + + def incompatible_check(self, global_args, check_args): + """Register all incompatible functions the feature is related.""" + if getattr(global_args, self.feature_name, None) and getattr(global_args, check_args, None): + raise AssertionError('{} and {} are incompatible.'.format(self.feature_name, check_args)) + + def dependency_check(self, global_args, check_args): + """Register all dependency functions the feature is related.""" + if getattr(global_args, self.feature_name, None) and not getattr(global_args, check_args, None): + raise AssertionError('{} requires {}.'.format(self.feature_name, check_args)) + + @staticmethod + def add_parser_argument_choices_value(parser, argument_name, new_choice): + """Add a new choice value to the existing choices of a parser argument.""" + for action in parser._actions: + exist_arg = isinstance(action, argparse.Action) and argument_name in action.option_strings + if exist_arg and action.choices is not None and new_choice not in action.choices: + action.choices.append(new_choice) +~~~ + +##### 参数特性支持 +| verl参数 | 简介| +| --- | --- | +| `actor_rollout_ref.actor.megatron.optimizer_offload` |是否卸载模型优化器到CPU，默认值为False| +| `actor_rollout_ref.actor.megatron.use_mbridge` |是否使用mbridge进行权重转换| +| `actor_rollout_ref.actor.megatron.param_offload` |是否卸载模型权重到CPU，默认值为False| +| `actor_rollout_ref.actor.megatron.tensor_model_parallel_size` | 张量并行大小；默认值为1。| +| `actor_rollout_ref.actor.megatron.pipeline_model_parallel_size` |流水并行大小，默认值为1| +| `actor_rollout_ref.actor.megatron.expert_model_parallel_size` | 专家并行大小，默认值为1| +| `actor_rollout_ref.actor.megatron.expert_tensor_parallel_size`|TP拓展EP大小，默认值为null| +| `actor_rollout_ref.actor.context_parallel_size`|序列并行大小，默认值为False| +| `actor_rollout_ref.actor.megatron.override_transformer_config.deallocate_pipeline_outputs`|张量在发送到下一个pp stage后,输出数据被释放，降低显存峰值，默认值为False| +| `actor_rollout_ref.actor.megatron.override_transformer_config.persist_layer_norm` |是否使用持久化 LayerNorm，默认值为False| +| `actor_rollout_ref.actor.megatron.override_transformer_config.moe_grouped_gemm` |是否使用持Group GEMM，默认值为False| +| `actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype` |用于路由和专家输出加权平均的数据类型。使用 fp32 或 fp64 可以提高稳定性，尤其是在专家数量较多时，默认值为fp32| +| `actor_rollout_ref.actor.megatron.override_transformer_config.account_for_loss_in_pipeline_split` |如果设置为 True，在流水线并行的划分和放置策略中，loss 层会被视为一个标准的 Transformer 层来处理。默认为False。| +| `actor_rollout_ref.actor.megatron.override_transformer_config.account_for_embedding_in_pipeline_split` |如果设置为 True，在流水线并行的划分和放置策略中，输入embedding 层会被视为一个标准的 Transformer 层来处理。默认为False。| +| `actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity` |重新计算激活的粒度，可选项为'full', 'selective' and 'none'。其中full代表重新计算整个transformer layer，selective代表只计算transformer layer中的核心注意力部分。默认为'none'。| +| `actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method` |该参数需将recompute_granularity设置为'full'才生效，可选项为'uniform', 'block'。默认为None。| +| `actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers` |该参数需将recompute_granularity设置为'full'才生效，默认为None。若recompute_method设置为uniform，该参数含义为每个均匀划分的重新计算单元的transformer layers数量。例如你可以指定为--recompute_granularity full --recompute_method uniform --recompute_num_layers 4。recompute_num_layers越大，显存占用越小，计算成本越大。注意：当前进程中的模型层数需能被recompute_num_layers整除。默认为None。| +| `actor_rollout_ref.actor.megatron.use_dist_checkpointing` |是否使用分布式权重，默认值为False| +| `actor_rollout_ref.actor.megatron.dist_checkpointing_path` |分布式权重路径，默认值为null| +| `actor_rollout_ref.actor.megatron.override_transformer_config.use_flash_attn` |是否使用fa，默认值为true| +| `actor_rollout_ref.actor.megatron.override_transformer_config.use_fused_rotary_pos_emb` |是否使用融合旋转位置编码，默认值为False| +| `actor_rollout_ref.actor.megatron.override_transformer_config.use_fused_swiglu` |是否使用融合swiglu，默认值为False| +| `actor_rollout_ref.actor.megatron.override_transformer_config.num_layers_in_first_pipeline_stage` |第一个pipeline stage 的层数，默认值为none| +| `actor_rollout_ref.actor.megatron.override_transformer_config.num_layers_in_last_pipeline_stage` |最后一个pipeline stage 的层数，默认值为none| + +注：`actor_rollout_ref.actor.megatron.use_mbridge` 与 `actor_rollout_ref.actor.megatron.virtual_pipeline_model_parallel_size` (VPP) 暂不支持同时开启。由于 verl 默认开启 mbridge, 使用 VPP 参数时请手动将 `actor_rollout_ref.actor.megatron.use_mbridge` 置为 False。 \ No newline at end of file diff --git a/verl/docs/ascend_tutorial/feature_support/npu_advance_features.md b/verl/docs/ascend_tutorial/feature_support/npu_advance_features.md new file mode 100644 index 0000000000000000000000000000000000000000..6548595d4f6ef5e2d8b6b4914fe135e58df9c5ff --- /dev/null +++ b/verl/docs/ascend_tutorial/feature_support/npu_advance_features.md @@ -0,0 +1,239 @@ +# NPU 高级特性指南 + +> 本文档介绍昇腾 NPU 在 verl 生态中的高级特性与优化能力，供开发者参考。 +> +Last updated: 05/13/2026. + +--- + +## 目录 + +- [1. 推理后端高级特性](#1-推理后端高级特性) + - [1.1 vLLM 推理后端](#11-vllm-推理后端) + - [1.2 SGLang 推理后端](#12-sglang-推理后端) +- [2. 训练后端高级特性](#2-训练后端高级特性) + - [2.1 FSDP 训练后端](#21-fsdp-训练后端) + - [2.2 Megatron 训练后端](#22-megatron-训练后端) +- [3. 性能优化特性](#3-性能优化特性) + - [3.1 内存优化](#31-内存优化) + - [3.2 计算加速](#32-计算加速) + - [3.3 并行策略](#33-并行策略) +- [4. 混合专家模型 (MoE) 特性](#4-混合专家模型-moe-特性) +- [5. 限制与注意事项](#5-限制与注意事项) + +--- + +## 1. 推理后端高级特性 + +当前 verl 支持 vLLM 和 SGLang 两种主流推理后端，均可在昇腾 NPU 上运行。以下列出各后端支持的高级特性参数。 + +### 1.1 vLLM 推理后端 + +昇腾通过 **vllm-ascend 插件** 支持 vLLM 推理后端。该插件遵循 [RFC](https://github.com/vllm-project/vllm/issues/11162)，提供可插拔接口将 Ascend NPU 与 vLLM 解耦。 + +--- + +### 1.2 SGLang 推理后端 + +昇腾通过向 SGLang 社区持续建设与维护来支持相关功能，涉及以下核心组件： + +| 组件 | 描述 | +|:---|:---| +| [sgl_kernel_npu](https://github.com/sgl-project/sgl-kernel-npu/blob/main/python/sgl_kernel_npu/README.md) | Ascend NPU 优化推理内核集合，含注意力机制、归一化、激活函数、LoRA 适配器等 | +| [deepep](https://github.com/sgl-project/sgl-kernel-npu/blob/main/python/deep_ep/README.md) | DeepEP 的 Ascend 实现，为 MoE 模型提供高度优化的专家并行 (EP) 通信内核 | + +#### 高级参数配置 + +| SGLang 参数 | verl 对应通用参数 | 功能说明 | +|:---|:---|:---| +| `attention_backend` | `actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend` | **注意力后端选择** — NPU 上应设置为 `ascend` 以调用昇腾优化内核 | +| `quantization` | `actor_rollout_ref.rollout.quantization` | **量化支持** — 支持模型量化加载与推理 | + + +> 更多 SGLang NPU 特性参数请参考 [sglang 社区 NPU 特性支持文档](https://docs.sglang.io/platforms/ascend_npu_support_features.html) + +--- + +## 2. 训练后端高级特性 + +### 2.1 FSDP 训练后端 + +昇腾通过 `torch_npu` 提供 FSDP 相关支持能力。 + +### 2.2 Megatron 训练后端 + +Megatron 是 NVIDIA 推出的模型并行训练框架。在 NPU 上运行需额外安装 **MindSpeed** 提供底层支持。MindSpeed 采用 **Monkey Patch** 技术无感替换 Megatron 关键组件，实现 NPU 适配。 + +#### MindSpeed Monkey Patch 框架原理 + +**触发入口:** +```python +from mindspeed.megatron_adaptor import repatch +``` + +**调用链:** +``` +repatch +├── 执行 megatron_adaptor.py 模块导入 +├── 导入 features_manager 模块 +├── 执行 mindspeed/features_manager/__init__.py +├── @AutoExecuteFunction 装饰器触发 +├── patch_features() 自动执行 +└── 进行 apply_features_pre_patches 和 apply_features_patches 操作 +``` + +**核心组件:** + +| 组件 | 职责 | +|:---|:---| +| `Patch` 类 | 实现函数/类的动态替换，支持多层装饰器叠加 | +| `parse_path()` | 动态模块导入和创建 | +| `MindSpeedPatchesManager` | 全局单例管理所有 patch 注册 | +| `MindSpeedFeature` | Feature 基类，各特性通过继承集成 patch 系统 | + +#### Megatron 高级参数配置 + +##### 内存与计算优化 + +| verl 参数 | 功能说明 | +|:---|:---| +| `actor_rollout_ref.actor.megatron.override_transformer_config.deallocate_pipeline_outputs` | **流水线输出释放** — 张量发送到下一 PP stage 后释放输出数据，降低显存峰值，默认 `False` | +| `actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity` | **重计算粒度控制** — 可选 `full` / `selective` / `none`。`full` 重算整个 Transformer 层，`selective` 仅重算注意力核心部分，默认 `none` | +| `actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method` | **重计算方法** — 需 `recompute_granularity=full`，可选 `uniform` / `block`，默认 `None` | +| `actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers` | **重计算层数** — 需 `recompute_granularity=full`，值越大显存占用越小、计算成本越高，需能被当前进程模型层数整除 | + +##### 融合算子加速 + +| verl 参数 | 功能说明 | +|:---|:---| +| `actor_rollout_ref.actor.megatron.override_transformer_config.use_flash_attn` | **Flash Attention** — 是否使用 Flash Attention 加速注意力计算，默认 `true` | +| `actor_rollout_ref.actor.megatron.override_transformer_config.use_fused_rotary_pos_emb` | **融合旋转位置编码** — 使用融合算子加速 RoPE 计算，默认 `False` | +| `actor_rollout_ref.actor.megatron.override_transformer_config.use_fused_swiglu` | **融合 SwiGLU** — 使用融合算子加速 SwiGLU 激活函数，默认 `False` | +| `actor_rollout_ref.actor.megatron.override_transformer_config.persist_layer_norm` | **持久化 LayerNorm** — 使用持久化策略优化 LayerNorm，默认 `False` | + +##### 流水线并行优化 + +| verl 参数 | 功能说明 | +|:---|:---| +| `actor_rollout_ref.actor.megatron.override_transformer_config.account_for_loss_in_pipeline_split` | **Loss 层流水线划分** — 将 loss 层视为标准 Transformer 层参与划分，默认 `False` | +| `actor_rollout_ref.actor.megatron.override_transformer_config.account_for_embedding_in_pipeline_split` | **Embedding 层流水线划分** — 将输入 embedding 层视为标准 Transformer 层参与划分，默认 `False` | +| `actor_rollout_ref.actor.megatron.override_transformer_config.num_layers_in_first_pipeline_stage` | **首 stage 层数** — 指定第一个 pipeline stage 的层数，默认 `none` | +| `actor_rollout_ref.actor.megatron.override_transformer_config.num_layers_in_last_pipeline_stage` | **末 stage 层数** — 指定最后一个 pipeline stage 的层数，默认 `none` | + +##### 权重管理 + +| verl 参数 | 功能说明 | +|:---|:---| +| `actor_rollout_ref.actor.megatron.use_mbridge` | **MBridge 权重转换** — 启用 mbridge 进行权重格式转换 | +| `actor_rollout_ref.actor.megatron.use_dist_checkpointing` | **分布式 checkpoint** — 使用分布式格式保存/加载权重，默认 `False` | +| `actor_rollout_ref.actor.megatron.dist_checkpointing_path` | **分布式权重路径** — 分布式 checkpoint 加载路径，默认 `null` | + +--- + +## 3. 性能优化特性 + +### 3.1 内存优化 + +| 特性 | 推理/训练 | 说明 | +|:---|:---|:---| +| KV Cache 动态释放 (`free_cache_engine`) | 推理 (vLLM) | 生成阶段后自动卸载 KV Cache，默认启用 | +| 内存节省模式 (`enable_memory_saver`) | 推理 (SGLang) | 支持显存动态释放/恢复，verl 默认 `True` | +| 参数 CPU 卸载 (`param_offload`) | 训练 (FSDP/Megatron) | 将模型权重卸载到 CPU | +| 优化器 CPU 卸载 (`optimizer_offload`) | 训练 (FSDP/Megatron) | 将优化器状态卸载到 CPU | +| 分块熵计算 (`entropy_from_logits_with_chunking`) | 训练 (FSDP) | 分块计算熵值降低显存峰值 | +| 熵计算重计算 (`entropy_checkpointing`) | 训练 (FSDP) | 对熵计算启用重计算 | +| 流水线输出释放 (`deallocate_pipeline_outputs`) | 训练 (Megatron) | PP 场景下释放已传递的张量 | +| 激活重计算 (`recompute_granularity`) | 训练 (Megatron) | 支持 full/selective/none 三级粒度控制 | + +### 3.2 计算加速 + +| 特性 | 推理/训练 | 说明 | +|:---|:---|:---| +| 分块预填充 (`enable_chunked_prefill`) | 推理 (vLLM) | 大预填充分块并与解码 batch 处理 | +| 前缀缓存 (`enable_prefix_caching`) | 推理 (vLLM) | 自动缓存共享前缀，减少重复计算 | +| Flash Attention | 训练 (Megatron) | 使用 Flash Attention 加速注意力计算，默认启用 | +| 融合旋转位置编码 (`use_fused_rotary_pos_emb`) | 训练 (Megatron) | 融合算子加速 RoPE | +| 融合 SwiGLU (`use_fused_swiglu`) | 训练 (Megatron) | 融合算子加速 SwiGLU 激活函数 | +| 持久化 LayerNorm (`persist_layer_norm`) | 训练 (Megatron) | 优化 LayerNorm 执行策略 | +| Group GEMM (`moe_grouped_gemm`) | 训练 (Megatron) | MoE 场景下的 Group GEMM 优化 | + +### 3.3 并行策略 + +| 并行类型 | vLLM | SGLang | FSDP | Megatron | 说明 | +|:---|:---|:---|:---|:---|:---| +| 数据并行 (DP) | ✅ | ✅ | ✅ | ✅ | 数据维度并行 | +| 张量并行 (TP) | ✅ | ✅ | — | ✅ | 层内张量切分 | +| 流水线并行 (PP) | — | — | — | ✅ | 层间流水线切分 | +| 专家并行 (EP) | ✅ | ✅ | — | ✅ | MoE 专家维度并行 | +| 序列并行 (SP/Ulysses) | ✅ | ✅ | ✅ | ✅ | 序列维度切分，支持长序列 | +| 上下文并行 (CP) | ✅ | — | — | ✅ | 上下文并行处理 | + +--- + +## 4. 混合专家模型 (MoE) 特性 + +### vLLM/SGLang 推理 MoE 支持 + +- **专家并行 (EP)** — 通过 `ep_size` 参数配置，将不同专家分配到不同 NPU 设备 +- SGLang 通过 [deepep](https://github.com/sgl-project/sgl-kernel-npu/blob/main/python/deep_ep/README.md) 提供高度优化的 EP 通信内核 + +### Megatron 训练 MoE 支持 + +| verl 参数 | 功能说明 | +|:---|:---| +| `actor_rollout_ref.actor.megatron.expert_model_parallel_size` | 专家并行 (EP) 大小，默认 `1` | +| `actor_rollout_ref.actor.megatron.expert_tensor_parallel_size` | TP 拓展 EP 大小，默认 `null` | +| `actor_rollout_ref.actor.megatron.override_transformer_config.moe_grouped_gemm` | **Group GEMM** — MoE 场景下使用 Group GEMM 优化专家计算，默认 `False` | +| `actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype` | **路由数据类型** — 路由与专家输出加权平均的数据类型，可选 `fp32`/`fp64`，默认 `fp32`，提高多专家场景稳定性 | + +--- + +## 5. 限制与注意事项 + +1. **mbridge 与 VPP 互斥** + - `actor_rollout_ref.actor.megatron.use_mbridge` 与 `actor_rollout_ref.actor.megatron.virtual_pipeline_model_parallel_size` (VPP) **暂不支持同时开启** + - 由于 verl 默认开启 mbridge，使用 VPP 时需手动将 `use_mbridge` 置为 `False` + +2. **FSDP1 vs FSDP2 差异** + - `forward_prefetch` 和 `use_orig_params` 仅适用于 FSDP1 + - FSDP2 为默认推荐版本，API 支持度参照 [昇腾 PyTorch 版本说明](https://www.hiascend.com/document/detail/zh/Pytorch/730/apiref/PyTorchNativeapi/docs/zh/native_apis/pytorch_2-7-1/torch-distributed-fsdp.md) + +3. **重计算参数依赖关系** + - `recompute_method` 需 `recompute_granularity='full'` 才生效 + - `recompute_num_layers` 需 `recompute_granularity='full'` 才生效 + - 当 `recompute_method='uniform'` 时，`recompute_num_layers` 表示每个重计算单元的 Transformer 层数，需能被当前进程模型层数整除 + +4. **SGLang NPU 特有配置** + - `attention_backend` 必须设置为 `ascend` 以调用昇腾优化内核 + - `enable_memory_saver` 在 verl 中默认启用，无需额外配置 + +--- + +## 附录：参数速查表 + +### 推理后端参数速查 + +| 参数类别 | vLLM 参数 | SGLang 参数 | verl 通用参数 | +|:---|:---|:---|:---| +| 模型路径 | `model_path` | `model_path` | `actor_rollout_ref.model.path` | +| 显存控制 | `gpu_memory_utilization` | `mem_fraction_static` | `actor_rollout_ref.rollout.gpu_memory_utilization` | +| 图模式 | `enforce_eager` | `disable_cuda_graph` | `actor_rollout_ref.rollout.enforce_eager` | +| 量化 | `quantization` | `quantization` | `actor_rollout_ref.rollout.quantization` | +| 最大序列长度 | `max_model_len` | — | `actor_rollout_ref.rollout.max_model_len` | +| 最大并发数 | `max_num_seqs` | `max_running_requests` | `actor_rollout_ref.rollout.max_num_seqs` | +| 分词器 | `skip_tokenizer_init` | `skip_tokenizer_init` | `actor_rollout_ref.rollout.skip_tokenizer_init` | +| 远程代码 | `trust_remote_code` | `trust_remote_code` | `actor_rollout_ref.model.trust_remote_code` | +| TP 并行 | `tp_size` | `tp_size` | `actor_rollout_ref.rollout.tensor_model_parallel_size` | +| DP 并行 | `dp_size` | `dp_size` | `actor_rollout_ref.rollout.data_parallel_size` | +| EP 并行 | `ep_size` | `ep_size` | `actor_rollout_ref.rollout.expert_parallel_size` | + +### 训练后端参数速查 + +| 参数类别 | FSDP 参数 | Megatron 参数 | +|:---|:---|:---| +| 参数卸载 | `fsdp_config.param_offload` | `megatron.param_offload` | +| 优化器卸载 | `fsdp_config.optimizer_offload` | `megatron.optimizer_offload` | +| 序列并行 | `ulysses_sequence_parallel_size` | `context_parallel_size` | +| Flash Attention | — | `override_transformer_config.use_flash_attn` | +| 重计算粒度 | — | `override_transformer_config.recompute_granularity` | +| 分布式 Checkpoint | — | `use_dist_checkpointing` | diff --git a/verl/docs/ascend_tutorial/get_start/dockerfile_build_guidance.rst b/verl/docs/ascend_tutorial/get_start/dockerfile_build_guidance.rst new file mode 100644 index 0000000000000000000000000000000000000000..e0a8903765812078f84722703b7f2ccf41a06c70 --- /dev/null +++ b/verl/docs/ascend_tutorial/get_start/dockerfile_build_guidance.rst @@ -0,0 +1,146 @@ +Ascend Dockerfile Build Guidance +=================================== + +Last updated: 05/19/2026. + + +镜像获取 & 公开镜像地址 +-------------------------- + +昇腾在 `quay.io/ascend/verl `_ 中托管每日构建的 A2/A3 镜像，基于下方“Dockerfile构建镜像脚本清单”中列出的 Dockerfile 构建。 + +每日构建镜像名格式：verl-{CANN版本}-{NPU设备类型}-{操作系统版本}-{python版本}-latest + +verl release版本镜像名格式：verl-{CANN版本}-{NPU设备类型}-{操作系统版本}-{python版本}-{verl release版本号} + + + +镜像硬件支持 +----------------------------------- + +Atlas 200T A2 Box16 + +Atlas 900 A2 PODc + +Atlas 800T A3 + + +镜像内各组件版本信息清单 +-------------------------- + +================= ============ +组件 版本 +================= ============ +基础镜像 Ubuntu 22.04 +Python 3.11 +CANN 8.5.0 +torch 2.9.0 +torch_npu 2.9.0 +torchvision 0.24.0 +vLLM 0.18.0 +vLLM-ascend 0.18.0 +Megatron-LM v0.12.1 +MindSpeed 2.3.0_core_r0.12.1 +triton-ascend 3.2.0 +mbridge 0.15.1 +SGLang v0.5.10 +sgl-kernel-npu 2026.02.01 +================= ============ + + + +.. _ascend-dockerfile-list: + +Dockerfile构建镜像脚本清单 +--------------------------- + +============== ================== ============== =========================================================================================================== +设备类型 CANN基础镜像版本 推理后端 参考文件 +============== ================== ============== =========================================================================================================== +A2 8.5.1 + 8.5.2组件 vLLM `Dockerfile.ascend_8.5.2_a2_qwen3-5 `_ +A2 8.5.0 vLLM `Dockerfile.ascend_8.5.0_a2 `_ +A2 8.5.0 vLLM `Dockerfile.ascend_8.5.0_a2_v0.7.1 `_ +A2 8.3.RC1 vLLM `Dockerfile.ascend_8.3.rc1_a2 `_ +A2 8.2.RC1 vLLM `Dockerfile.ascend_8.2.rc1_a2 `_ +A2 8.5.0 SGLang `Dockerfile.ascend.sglang_8.5.0_a2 `_ +A2 8.3.RC1 SGLang `Dockerfile.ascend.sglang_8.3.rc1_a2 `_ +A3 8.5.1 + 8.5.2组件 vLLM `Dockerfile.ascend_8.5.2_a3_qwen3-5 `_ +A3 8.5.0 vLLM `Dockerfile.ascend_8.5.0_a3 `_ +A3 8.5.0 vLLM `Dockerfile.ascend_8.5.0_a3_v0.7.1 `_ +A3 8.3.RC1 vLLM `Dockerfile.ascend_8.3.rc1_a3 `_ +A3 8.2.RC1 vLLM `Dockerfile.ascend_8.2.rc1_a3 `_ +A3 8.5.0 SGLang `Dockerfile.ascend.sglang_8.5.0_a3 `_ +A3 8.3.RC1 SGLang `Dockerfile.ascend.sglang_8.3.rc1_a3 `_ +============== ================== ============== =========================================================================================================== + +**说明：** + +* 推理后端为 ``vLLM`` 镜像中，vLLM、vLLM-ascend、MindSpeed、Megatron-LM、verl 为源码安装，源码位于镜像根目录 ``/`` 下。 +* 推理后端为 ``SGLang`` 镜像中，SGLang、MindSpeed、verl 为源码安装，源码位于镜像根目录 ``/`` 下。 + + +镜像构建命令示例 +-------------------- + +.. code:: bash + + # Navigate to the directory containing the Dockerfile + cd {verl-root-path}/docker/ascend + + # Build the image + # vLLM + docker build -f Dockerfile.ascend_8.5.0_a2 -t verl-ascend:8.5.0-a2 . + # SGLang + docker build -f Dockerfile.ascend.sglang_8.5.0_a2 -t verl-ascend-sglang:8.5.0-a2 . + + # Query local images after build + docker images + +**说明：** + +* 以 vLLM 的镜像为例，``Dockerfile.ascend_8.5.0_a2`` 为 Dockerfile 文件名，``verl-ascend:8.5.0-a2`` 中， verl-ascend 为自定义的镜像名称，8.5.0-a2 为自定义的镜像标签 + +容器启动命令模板 +---------------- + +.. code:: bash + + docker run -dit \ + --ipc=host \ + --network host \ + --name {your_docker_name} \ + --privileged \ + -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \ + -v /usr/local/Ascend/firmware:/usr/local/Ascend/firmware \ + -v /usr/local/sbin:/usr/local/sbin \ + -v /usr/sbin:/usr/sbin \ + -v /home:/home \ + -v /data:/data \ + {image_name}:{tag} \ + /bin/bash + +**说明：** + +* 如需挂载其他本地路径到容器，请自行添加 ``-v <宿主机路径>:<容器内路径>`` +* 建议将 ``{your_docker_name}`` 替换为具有实际意义的容器名称 +* ``--privileged`` 参数授予容器扩展权限，请根据实际安全需求评估是否必要 +* ``{image_name}:{tag}`` 请换成容器构建时对应的镜像名称与标签 + +启动容器 +-------- + +.. code:: bash + + docker start {your_docker_name} + +进入正在运行的容器 +------------------ + +.. code:: bash + + docker exec -it {your_docker_name} bash + + +声明 +-------------------- +verl中提供的ascend相关Dockerfile、镜像皆为参考样例，可用于尝鲜体验，如在生产环境中使用请通过官方正式途径沟通，谢谢。 diff --git a/verl/docs/ascend_tutorial/get_start/install_guidance.rst b/verl/docs/ascend_tutorial/get_start/install_guidance.rst new file mode 100644 index 0000000000000000000000000000000000000000..5bd7a235b11f9e18553364ffe146e573c94c05b8 --- /dev/null +++ b/verl/docs/ascend_tutorial/get_start/install_guidance.rst @@ -0,0 +1,462 @@ +Ascend Install Guidance +================= + +Last updated: 05/13/2026. + +关键更新 +-------- + +- 2026/05/13：vLLM 路线已按 `PR + #6291 `__\ 将 vLLM / + vLLM-Ascend 从 ``0.13.0`` 更新为 ``0.18.0``\ ，vLLM + 路线对应基础环境版本同步调整为 torch ``2.9.0``\ 、torch_npu + ``2.9.0``\ 。 +- 2025/12/11：verl 存量场景目前支持自动识别 NPU 设备类型。原则上，GPU + 脚本在昇腾上运行时不再需要显式设置 + ``trainer.device=npu``\ ；新增特性仍可通过设置 ``trainer.device`` + 优先指定设备类型。 + +.. + + [说明] 自动识别 NPU 设备类型的前提，是运行程序所在环境包含 + ``torch_npu`` 软件包。如环境中不包含 ``torch_npu``\ ，仍需显式指定 + ``trainer.device=npu``\ 。 + +硬件支持 +-------- + +Atlas 200T A2 Box16 + +Atlas 900 A2 PODc + +Atlas 800T A3 + +后端拆分说明 +------------ + +verl 在昇腾 NPU 上通常涉及两类后端：\ **rollout +后端**\ 和\ **训练后端**\ 。二者负责的阶段不同，安装时不要混淆。 + +.. list-table:: + :header-rows: 1 + + * - 后端类型 + - 可选项 + - 作用 + * - rollout 后端 + - vLLM-Ascend / SGLang + - 负责生成阶段，给 prompt 生成 response + * - 训练后端 + - FSDP / Megatron 路线 + - 负责训练阶段，包括 actor/ref 计算、logprob、loss、反向传播、参数更新和 checkpoint + +说明： + +- vLLM-Ascend 和 SGLang 是 rollout 后端，二者主要负责推理生成。 +- 训练侧主要包括 FSDP 路线和 Megatron 路线。 +- FSDP 是基础训练路线，通常不需要额外安装 Megatron-LM、MindSpeed 或 + mbridge。 +- Megatron 路线在昇腾上包含两种常见实现：Megatron + MindSpeed，以及基于 + Megatron/MindSpeed 体系的 MindSpeed-LLM。 +- 使用 Megatron 作为训练后端时，需要额外安装 Megatron-LM、MindSpeed 和 + mbridge。 + +本文中的 Megatron 后端，对应 verl 中的 +``actor_rollout_ref.actor.strategy=megatron``\ 。其中，Megatron-LM +提供大模型并行训练能力；MindSpeed 提供昇腾 NPU 上的 Megatron-LM +适配和优化；mbridge 用于 Megatron-LM / MindSpeed +相关模型桥接、权重加载和转换。 + +- 使用 MindSpeed-LLM 时，需要额外安装 MindSpeed-LLM，并配置对应的 + MindSpeed、Megatron-LM 和 mbridge。MindSpeed-LLM 本质上是基于 + Megatron/MindSpeed 体系的昇腾 LLM 训练后端实现。 + +安装流程 +-------- + +DockerFile 镜像构建、获取和使用 +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +如需要通过 DockerFile 构建镜像，或希望使用基于 verl 构建的镜像，请参考 +`dockerfile_build_guidance `__\ 。 + +如果想直接获取镜像，可前往 +`quay.io/ascend/verl `__ +获取，镜像中通常已包含基础环境和依赖软件包。 + +vLLM 路线 +~~~~~~~~~ + +本路线适用于使用 vLLM-Ascend 作为 rollout 后端的场景。该路线只说明生成侧 +vLLM-Ascend 的安装；训练侧可继续使用 FSDP，或按需安装 Megatron 路线相关依赖 +（Megatron + MindSpeed 或 MindSpeed-LLM）。 + +关键支持版本 +^^^^^^^^^^^^ + +============= ======================================= +软件 版本 +============= ======================================= +Python ``>=3.10, <3.12``\ ，推荐 ``3.11`` +CANN ``8.5.0`` +NNAL / ATB ``8.5.0`` +torch ``2.9.0`` +torch_npu ``2.9.0`` +torchvision ``0.24.0`` +torchaudio ``2.9.0`` +triton-ascend ``3.2.0`` +transformers ``>=4.57.4, <5.0.0``\ ，推荐 ``4.57.6`` +vLLM ``0.18.0`` +vLLM-Ascend ``0.18.0`` +============= ======================================= + +.. + + [说明] vLLM-Ascend ``0.18.0`` 的 `release + note `__ + 中提到，因已知问题可手动升级到 ``torch_npu==2.9.0.post1+git4c901a4`` + 和 ``triton-ascend==3.2.0.dev20260322``\ 。如环境中已升级到 CANN + ``9.0.0``\ ，需要同步升级对应的 ``torch_npu`` 和 ``triton-ascend`` + 版本。 + +安装基础环境 +^^^^^^^^^^^^ + +基础环境涉及以下软件包，请参考 +`文档 `__ 安装。 + +========= ================= +软件 版本 +========= ================= +Python ``>=3.10, <3.12`` +CANN ``8.5.0`` +torch ``2.9.0`` +torch_npu ``2.9.0`` +========= ================= + +可创建 Python 环境： + +.. code:: bash + + conda create -n verl-vllm-npu python=3.11 -y + conda activate verl-vllm-npu + +在 x86 平台安装时，pip 需要配置额外的源，指令如下： + +.. code:: bash + + pip config set global.extra-index-url "https://download.pytorch.org/whl/cpu/" + +安装其他软件包 +^^^^^^^^^^^^^^ + +基础环境准备完毕后，需要通过指令安装以下软件包： + +.. code:: bash + + # 清理环境上可能存在的历史 triton / triton-ascend 软件包残留 + pip uninstall -y triton triton-ascend + + # 安装与 vLLM-Ascend 0.18.0 对应的软件包 + pip install torchvision==0.24.0 + pip install torchaudio==2.9.0 + pip install triton-ascend==3.2.0 + pip install "transformers>=4.57.4,<5.0.0" + +.. _安装-vllm--vllm-ascend: + +安装 vLLM & vLLM-Ascend +^^^^^^^^^^^^^^^^^^^^^^^ + +需确保CANN ascend-toolkit 和 nnal 环境变量被激活，对于CANN默认安装路径 +/usr/local/Ascend 而言，激活指令如下： + +.. code:: bash + + source /usr/local/Ascend/ascend-toolkit/set_env.sh + source /usr/local/Ascend/nnal/atb/set_env.sh + +vLLM 源码安装指令： + +.. code:: bash + + git clone --depth 1 --branch v0.18.0 https://github.com/vllm-project/vllm.git + cd vllm + VLLM_TARGET_DEVICE=empty pip install -v -e . + cd .. + +vLLM-Ascend 源码安装指令： + +.. code:: bash + + git clone --depth 1 --branch v0.18.0 https://github.com/vllm-project/vllm-ascend.git + cd vllm-ascend + git submodule update --init --recursive + pip install -v -e . + cd .. + +安装 verl +^^^^^^^^^ + +.. code:: bash + + git clone --recursive https://github.com/verl-project/verl.git + cd verl + pip install -r requirements-npu.txt + pip install -v -e . + + # （可选）提示：为了更佳的使用体验，最好将 recipe 子模块更新至最新 commit + cd recipe + git checkout main + cd .. + +SGLang 路线 +~~~~~~~~~~~ + +本路线适用于使用 SGLang 作为 rollout 后端的场景。该路线只说明生成侧 +SGLang 的安装；训练侧可继续使用 FSDP，或按需安装 Megatron 路线相关依赖 +（Megatron + MindSpeed 或 MindSpeed-LLM）。 + +.. _关键支持版本-1: + +关键支持版本 +^^^^^^^^^^^^ + +========= ================= +软件 版本 +========= ================= +Python ``3.11`` +HDK ``>=25.3.RC1`` +CANN ``>=8.3.RC1`` +torch ``>=2.7.1`` +torch_npu ``>=2.7.1.post2`` +SGLang ``v0.5.8`` +========= ================= + +从Docker镜像进行安装 +^^^^^^^^^^^^^^^^^^^^ + +我们提供了DockerFile进行构建,详见 +`dockerfile_build_guidance `__ +，请根据设备自行选择对应构建文件 + +从自定义环境安装 +^^^^^^^^^^^^^^^^ + +1. 安装 HDK & CANN 依赖并激活。 + +异构计算架构 CANN 是昇腾针对 AI +场景推出的异构计算架构。为了使训练和推理引擎能够利用更好、更快的硬件支持，需要安装以下 +`先决条件 `__\ ： + +.. code:: text + + HDK >=25.3.RC1 + CANN >=8.3.RC1 + +安装完成后请激活环境： + +.. code:: bash + + source /usr/local/Ascend/ascend-toolkit/set_env.sh + source /usr/local/Ascend/nnal/atb/set_env.sh + +2. 创建 conda 环境。 + +.. code:: bash + + conda create -n verl-sglang python==3.11 + conda activate verl-sglang + +3. 执行 verl 中提供的 SGLang 安装脚本。 + +.. code:: bash + + git clone https://github.com/verl-project/verl.git + + # Make sure you have activated verl conda env + # NPU_DEVICE=A3 or A2 depends on your device + # USE_MEGATRON=1 if you need to install megatron backend + NPU_DEVICE=A3 USE_MEGATRON=1 bash verl/scripts/install_sglang_mcore_npu.sh + +.. + + [说明] 如果在此步骤中遇到错误，请检查 + ``verl/scripts/install_sglang_mcore_npu.sh``\ ，并手动按照脚本中的步骤操作。 + +4. 安装 verl。 + +.. code:: bash + + cd verl + pip install --no-deps -e . + pip install -r requirements-npu.txt + +SGLang 环境变量 +^^^^^^^^^^^^^^^ + +当前 NPU 上支持 SGLang 后端必须添加以下环境变量： + +.. code:: bash + + # 支持 NPU 单卡多进程 + export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050 + export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050 + + # 规避 Ray 在 device 侧调用无法根据 is_npu_available 接口识别设备可用性 + export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 + + # 根据当前设备和需要卡数定义 + export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 + + # 使能推理 EP 时需要 + export SGLANG_DEEPEP_BF16_DISPATCH=1 + +8 卡环境： + +.. code:: bash + + export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 + +16 卡环境： + +.. code:: bash + + export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15 + +Megatron 后端训练依赖 +~~~~~~~~~~~~~~~~~~~~~ + +使用 Megatron 作为训练后端时，需要额外安装 Megatron-LM、MindSpeed 和 +mbridge。 + +版本要求 +^^^^^^^^ + +=========== ====================== +软件 版本 +=========== ====================== +MindSpeed ``2.3.0_core_r0.12.1`` +Megatron-LM ``core_v0.12.1`` +=========== ====================== + +安装 MindSpeed +^^^^^^^^^^^^^^ + +MindSpeed 源码安装指令： + +.. code:: bash + + # 下载 MindSpeed，切换到指定 commit-id，并下载 Megatron-LM + git clone https://gitcode.com/Ascend/MindSpeed.git + cd MindSpeed && git checkout 2.3.0_core_r0.12.1 && cd .. + git clone --depth 1 --branch core_v0.12.1 https://github.com/NVIDIA/Megatron-LM.git + + # 安装 Megatron & MindSpeed + pip install -e Megatron-LM + pip install -e MindSpeed + + # 安装 mbridge + pip install mbridge + +MindSpeed 对应 Megatron-LM 后端使用场景，使用方式如下： + +1. 使能 verl worker 模型 ``strategy`` 配置为 ``megatron``\ ，例如 + ``actor_rollout_ref.actor.strategy=megatron``\ 。 +2. MindSpeed 自定义入参可通过 ``override_transformer_config`` + 参数传入，例如对 actor 模型开启 FA 特性可使用 + ``+actor_rollout_ref.actor.megatron.override_transformer_config.use_flash_attn=True``\ 。 +3. 更多特性信息可参考 `MindSpeed & verl + 文档 `__\ 。 + +MindSpeed-LLM 训练后端支持 +~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +如需使用基于 Megatron/MindSpeed 体系的 MindSpeed-LLM 训练后端，需要额外下载 +MindSpeed-LLM。需要注意的是，MindSpeed-LLM 训练后端依赖 MindSpeed-LLM +master 分支、MindSpeed master 分支以及 Megatron-LM ``core_v0.12.1`` +分支。 + +MindSpeed-LLM 及相关依赖的源码安装指令： + +.. code:: bash + + # 下载 MindSpeed-LLM、MindSpeed 和 Megatron-LM + git clone https://gitcode.com/Ascend/MindSpeed-LLM.git + git clone https://gitcode.com/Ascend/MindSpeed.git + git clone --depth 1 --branch core_v0.12.1 https://github.com/NVIDIA/Megatron-LM.git + + # 配置环境变量 + export PYTHONPATH=$PYTHONPATH:your path/Megatron-LM + export PYTHONPATH=$PYTHONPATH:your path/MindSpeed + export PYTHONPATH=$PYTHONPATH:your path/MindSpeed-LLM + + # 安装 mbridge + pip install mbridge + +MindSpeed-LLM 作为基于 Megatron/MindSpeed 体系的昇腾 LLM 训练后端使用时，使用方式如下： + +1. 使能 verl worker 模型 ``strategy`` 配置为 ``mindspeed``\ ，例如 + ``actor_rollout_ref.actor.strategy=mindspeed``\ 。 +2. MindSpeed-LLM 自定义入参可通过 ``llm_kwargs`` 参数传入，例如对 MOE + 模型开启 GMM 特性可使用 + ``+actor_rollout_ref.actor.mindspeed.llm_kwargs.moe_grouped_gemm=True``\ 。 +3. 更多特性信息可参考 `MindSpeed-LLM + 内的特性文档 `__\ 。 + +昇腾暂不支持生态库说明 +~~~~~~~~~~~~~~~~~~~~~~ + +verl 中昇腾暂不支持生态库如下： + ++------------------+--------------------------------------------------+ +| 软件 | 说明 | ++==================+==================================================+ +| ``flash_attn`` | 不支持通过独立 ``flash_attn`` 包使能 flash | +| | attention 加速，支持通过 transformers 使用 | ++------------------+--------------------------------------------------+ +| ``liger-kernel`` | 暂不支持 | ++------------------+--------------------------------------------------+ + +组件作用说明 +------------ + ++---------------+--------------------------------------------------------------------------+ +| 组件 | 作用 | ++===============+==========================================================================+ +| CANN | 昇腾 NPU 运行时、算子和通信基础环境 | ++---------------+--------------------------------------------------------------------------+ +| HDK | 昇腾硬件开发组件 | ++---------------+--------------------------------------------------------------------------+ +| NNAL / ATB | vLLM-Ascend 官方环境要求中包含的高级算子组件； | +| | 官方镜像通常已包含，缺 ``libatb.so`` 时再补装 | ++---------------+--------------------------------------------------------------------------+ +| torch | PyTorch 基础框架 | ++---------------+--------------------------------------------------------------------------+ +| torch_npu | PyTorch 昇腾 NPU 适配 | ++---------------+--------------------------------------------------------------------------+ +| torchvision | 与 torch 匹配的视觉相关依赖 | ++---------------+--------------------------------------------------------------------------+ +| torchaudio | 与 torch 匹配的音频相关依赖 | ++---------------+--------------------------------------------------------------------------+ +| triton-ascend | 昇腾侧 Triton 适配 | ++---------------+--------------------------------------------------------------------------+ +| transformers | 模型结构、配置和 tokenizer 相关依赖 | ++---------------+--------------------------------------------------------------------------+ +| vLLM | 高性能推理框架 | ++---------------+--------------------------------------------------------------------------+ +| vLLM-Ascend | vLLM 的昇腾适配后端 | ++---------------+--------------------------------------------------------------------------+ +| SGLang | 推理后端 | ++---------------+--------------------------------------------------------------------------+ +| verl | 强化学习和 SFT 训练框架 | ++---------------+--------------------------------------------------------------------------+ +| Megatron-LM | 大模型并行训练基础框架 | ++---------------+--------------------------------------------------------------------------+ +| MindSpeed | Megatron-LM 在昇腾 NPU 上的适配和优化组件 | ++---------------+--------------------------------------------------------------------------+ +| MindSpeed-LLM | 基于 Megatron/MindSpeed 体系的昇腾 LLM 训练后端实现； | +| | 使用时通过 ``actor_rollout_ref.actor.strategy=mindspeed`` 使能 | ++---------------+--------------------------------------------------------------------------+ +| mbridge | Megatron-LM / MindSpeed 相关模型桥接、权重加载和转换组件 | ++---------------+--------------------------------------------------------------------------+ \ No newline at end of file diff --git a/verl/docs/ascend_tutorial/get_start/quick_start.rst b/verl/docs/ascend_tutorial/get_start/quick_start.rst new file mode 100644 index 0000000000000000000000000000000000000000..1552f07e0c6d9c95de0a67aea19c4334dc02227f --- /dev/null +++ b/verl/docs/ascend_tutorial/get_start/quick_start.rst @@ -0,0 +1,172 @@ +Ascend Quickstart +================= + +**Last updated:** 05/13/2026. + +关键更新 +-------- + +- 2026/05/13：将 quick start 和 install guidance 分开。 +- 2025/12/11：verl 存量场景目前支持自动识别 NPU 设备类型，GPU 脚本在昇腾上运行，原则上不再需要显式设置 ``trainer.device=npu`` 参数，新增特性通过设置 ``trainer.device`` 仍可优先使用，逐步适配自动识别能力。 + +硬件支持 +-------- + +- Atlas 200T A2 Box16 +- Atlas 900 A2 PODc +- Atlas 800T A3 + +Ascend Quickstart with vLLM Backend +=================================== + +基础验证场景 +------------ + +如需快速验证环境和基础链路，可以使用 Qwen2.5-0.5B GRPO 场景。 + +该场景用于检查： + +- verl 入口是否可用； +- 数据是否可读取； +- actor、rollout、reference worker 是否能初始化； +- vLLM-Ascend rollout 是否能生成； +- 训练链路是否能完成首个 step。 + +准备 GSM8K 数据 +--------------- + +.. code-block:: bash + + python3 examples/data_preprocess/gsm8k.py --local_save_dir ~/data/gsm8k + +生成文件： + +.. code-block:: text + + ~/data/gsm8k/train.parquet + ~/data/gsm8k/test.parquet + +启动 Qwen2.5-0.5B GRPO 基础验证 +-------------------------------- + +.. code-block:: bash + + set -x + + python3 -m verl.trainer.main_ppo \ + algorithm.adv_estimator=grpo \ + data.train_files=$HOME/data/gsm8k/train.parquet \ + data.val_files=$HOME/data/gsm8k/test.parquet \ + data.train_batch_size=128 \ + data.max_prompt_length=512 \ + data.max_response_length=128 \ + data.filter_overlong_prompts=True \ + data.truncation='error' \ + actor_rollout_ref.model.path=Qwen/Qwen2.5-0.5B-Instruct \ + actor_rollout_ref.actor.optim.lr=5e-7 \ + actor_rollout_ref.model.use_remove_padding=False \ + actor_rollout_ref.actor.entropy_coeff=0.001 \ + actor_rollout_ref.actor.ppo_mini_batch_size=64 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=20 \ + actor_rollout_ref.actor.use_kl_loss=True \ + actor_rollout_ref.actor.kl_loss_coef=0.001 \ + actor_rollout_ref.actor.kl_loss_type=low_var_kl \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.fsdp_config.param_offload=False \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=40 \ + actor_rollout_ref.rollout.enable_chunked_prefill=False \ + actor_rollout_ref.rollout.tensor_model_parallel_size=2 \ + actor_rollout_ref.rollout.name=vllm \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \ + actor_rollout_ref.rollout.n=5 \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=40 \ + actor_rollout_ref.ref.fsdp_config.param_offload=True \ + algorithm.kl_ctrl.kl_coef=0.001 \ + trainer.critic_warmup=0 \ + trainer.logger=console \ + trainer.project_name='verl_grpo_example_gsm8k' \ + trainer.experiment_name='qwen2_5_0_5b_grpo_vllm_ascend' \ + trainer.n_gpus_per_node=8 \ + trainer.nnodes=1 \ + trainer.save_freq=-1 \ + trainer.test_freq=5 \ + trainer.total_epochs=1 $@ + +关键配置 +-------- + +.. list-table:: + :header-rows: 1 + :widths: 40 60 + + * - 配置项 + - 说明 + * - ``actor_rollout_ref.rollout.name=vllm`` + - 使用 vLLM 作为 rollout 后端 + * - ``actor_rollout_ref.rollout.tensor_model_parallel_size`` + - rollout 推理侧张量并行大小 + * - ``actor_rollout_ref.rollout.gpu_memory_utilization`` + - rollout 推理侧可使用的设备显存比例 + * - ``actor_rollout_ref.rollout.n`` + - 每个 prompt 生成的 response 数量 + * - ``trainer.n_gpus_per_node`` + - 昇腾场景中表示每节点使用的 NPU 数量 + * - ``trainer.nnodes`` + - 节点数量 + +Ascend Quickstart with SGLang Backend +===================================== + +最佳实践 +-------- + +我们提供 `最佳实践 `_ 作为参考。 + +环境变量与参数 +-------------- + +当前 NPU 上支持 SGLang 后端必须添加以下环境变量。 + +.. code-block:: bash + + # 支持 NPU 单卡多进程 + # https://www.hiascend.com/document/detail/zh/canncommercial/850/commlib/hcclug/hcclug_000091.html + export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050 + export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050 + + # 规避 Ray 在 device 侧调用无法根据 is_npu_available 接口识别设备可用性 + export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 + + # 根据当前设备和需要卡数定义 + export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15 + + # 使能推理 EP 时需要 + export SGLANG_DEEPEP_BF16_DISPATCH=1 + +当前 verl 已解析推理常见参数，详见 `async_sglang_server.py `_ 中 ``ServerArgs`` 初始化传参。 + +其他 `SGLang 参数 `_ 均可通过 ``engine_kwargs`` 进行参数传递。 + +vLLM 后端脚本转换为 SGLang +-------------------------- + +vLLM 后端推理脚本转换为 SGLang，需要添加或修改以下参数。 + +.. code-block:: bash + + # 必须 + actor_rollout_ref.rollout.name=sglang \ + +actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend="ascend" \ + + # 可选 + # 使能推理 EP，详细使用方法见： + # https://github.com/sgl-project/sgl-kernel-npu/blob/main/python/deep_ep/README_CN.md + ++actor_rollout_ref.rollout.engine_kwargs.sglang.deepep_mode="auto" \ + ++actor_rollout_ref.rollout.engine_kwargs.sglang.moe_a2a_backend="deepep" \ + + # MoE 模型多 DP 时必须设置为 True + +actor_rollout_ref.rollout.engine_kwargs.sglang.enable_dp_attention=False \ + + # chunked_prefill 默认关闭 + +actor_rollout_ref.rollout.engine_kwargs.sglang.chunked_prefill_size=-1 \ No newline at end of file diff --git a/verl/docs/ascend_tutorial/model_support/examples/ascend_retool_best_pratice.rst b/verl/docs/ascend_tutorial/model_support/examples/ascend_retool_best_pratice.rst new file mode 100644 index 0000000000000000000000000000000000000000..5938f9167931293a1be24d8e6df50ae98e18db38 --- /dev/null +++ b/verl/docs/ascend_tutorial/model_support/examples/ascend_retool_best_pratice.rst @@ -0,0 +1,266 @@ +Ascend Retool Best Practice +=================================== + +Last updated: 03/01/2026. + +引言 +---------------------------------- + +Retool论文参考([Retool](https://arxiv.org/pdf/2504.11536)) +集成代码解释器工具，通过多轮实时代码执行进行策略部署，并教会模型根据结果反馈学习何时以及如何调用工具。 + +1. 环境构建 +2. 模型训练 + +用例模型脚本以及其需要的硬件条件各自如下： + +=============== ============ ============ =============== +模型 NPU型号 节点数量 训推后端 +=============== ============ ============ =============== +``Qwen2.5-7B`` Atlas 900 A2 1 ``vllm + FSDP`` +=============== ============ ============ =============== + +环境构建 +----------------------------------- +1.从自定义Conda环境进行构建 + +============ ============================================================ +software version +============ ============================================================ +Python ``>=3.10, <3.12`` +CANN ``==8.3.RC1`` +torch ``==2.7.1`` +torch_npu ``==2.7.1`` +verl ``v0.6.1 commitId=d62da4950573d7a4b7ef2362337952e7ab59e78d`` +vllm ``v0.11.0`` +vllm-ascend ``v0.11.0-dev`` +transformers ``4.57.6`` +============ ============================================================ + +模型训练与评估 +----------------------------------- +1.模型数据准备 +^^^^^^^^^^^ +`Qwen2.5-7B` +^^^^^^^^^^^ +**下载模型权重** + +--local-dir: 模型保存路径 + +.. code-block:: bash + + git clone https://huggingface.co/Qwen/Qwen2.5-7B-Instruct + +**下载训练数据集** + +.. code-block:: bash + + git clone https://huggingface.co/datasets/BytedTsinghua-SIA/DAPO-Math-17k + +**下载评估数据集** + +.. code-block:: bash + + git clone https://huggingface.co/datasets/Maxwell-Jia/AIME_2024 + +**下载预训练数据集** + +.. code-block:: bash + + python3 recipe/retool/retool_sft_preprocess.py + +*注:自动下载ReTool-SFT，最后生成数据默认保存在~/ReTool-SFT/data目录下* + +**执行预训练脚本** + +.. code-block:: bash + + bash recipe/retool/run_qwen2_7b_sft_npu.sh # 需适配脚本中路径 + +**合并预训练权重生成checkpoint** + +.. code-block:: bash + + python3 -m verl.model_merger merge --backend fsdp \ + --local_dir /PATH/TO/checkpoint/multiturn-sft-qwen-2.5-7b-instruct/global_step_372 \ + --target_dir /PATH/TO/checkpoint/multiturn-sft-qwen-2.5-7b-instruct/global_step_372/huggingface + +2.代码沙箱准备 + +开源沙箱代码及部署参考 +https://github.com/bytedance/SandboxFusion + +**沙箱代码下载** + +.. code-block:: bash + + git clone -b main https://github.com/bytedance/SandboxFusion.git + +**沙箱安装** + +.. code-block:: bash + + cd SandboxFusion + conda create -n sandbox -y python=3.11 + conda activate sandbox + pip install poetry + poetry lock + poetry install + mkdir -p docs/build + cd runtime/python + bash install-python-runtime.sh + cd ../../ + make run-online + +3.训练 + +示例配置文件如下，在recipe/retool目录下创建一个run_qwen2.5_7b_dapo_npu.sh +根据开发者实际路径配置情况修改模型训练脚本中的以下参数 + +.. code-block:: bash + + set -x + + export VLLM_USE_V1=1 + export TORCHDYNAMO_DISABLE=1 + export VLLM_ASCEND_ENABLE_NZ=0 + export TASK_QUEUE_ENABLE=1 + export VLLM_ENABLE_GRAPH_MODE=1 + export HCCL_OP_EXPANSION_MODE="AIV" + export VLLM_ASCEND_ENABLE_MLP_OPTIMIZE=1 + export LD_PRELOAD=/usr/local/lib/libjemalloc.so.2 + + # ================= data/model/tool ================= + HDFS_ROOT=${HDFS_ROOT:-"${PWD}"} + DATA_ROOT=${DATA_ROOT:-"${PWD}"} + + dapo_math_17k=$DATA_ROOT/dataset/BytedTsinghua-SIA/DAPO-Math-17k + aime_2024=$DATA_ROOT/dataset/Maxwell-Jia/AIME_2024 + #aime_2025=$DATA_ROOT/dataset/yentinglin/aime_2025 + model_path=$DATA_ROOT/dataset/checkpoint/multiturn-sft-qwen-2.5-7b-instruct/global_step_372/huggingface + + train_files="['$dapo_math_17k']" + test_files="['$aime_2024']" + + # tool + tool_config_path=recipe/retool/sandbox_fusion_tool_config.yaml + + # wandb + project_name=retool + experiment_name=qwen2.5-7b_dapo + default_local_dir=$DATA_ROOT/checkpoint/$experiment_name + + # 创建日志文件 + export TIMESTAMP=$(date +%Y%m%d_%H%M%S) + LOG_DIR="$HDFS_ROOT/verl/logs/$project_name/$experiment_name" + # 判断路径是否存在 + if [ ! -d "$LOG_DIR" ]; then + # 路径不存在，创建路径 + mkdir -p "$LOG_DIR" + echo "Directory $LOG_DIR created." + else + echo "Directory $LOG_DIR already exists." + fi + + LOG_FILE="${LOG_DIR}/${TIMESTAMP}.log" + touch "$LOG_FILE" + echo "Log file $LOG_FILE created." + + # ================= algorithm ================= + adv_estimator=grpo + + use_kl_in_reward=False + kl_coef=0.0 + use_kl_loss=False + kl_loss_coef=0.0 + + clip_ratio_low=0.2 + clip_ratio_high=0.28 + + max_turns=16 + max_prompt_length=2048 + max_response_length=20480 + actor_lr=1e-6 + + train_batch_size=32 + ppo_mini_batch_size=16 + + n_resp_per_prompt=16 + n_resp_per_prompt_val=30 + + # ================= performance ================= + infer_tp=2 # vllm + train_sp=4 # train + offload=True + + actor_max_token_len_per_gpu=$(( (max_prompt_length + max_response_length) * 1 )) + log_prob_max_token_len_per_gpu=$(( actor_max_token_len_per_gpu * 4 )) + + PYTHONUNBUFFERED=1 python3 -m verl.trainer.main_ppo \ + algorithm.adv_estimator=$adv_estimator \ + algorithm.use_kl_in_reward=$use_kl_in_reward \ + algorithm.kl_ctrl.kl_coef=$kl_coef \ + data.train_files="$train_files" \ + data.val_files="$test_files" \ + data.return_raw_chat=True \ + data.train_batch_size=$train_batch_size \ + data.max_prompt_length=$max_prompt_length \ + data.max_response_length=$max_response_length \ + data.filter_overlong_prompts=True \ + data.truncation='error' \ + data.custom_cls.path=recipe/retool/retool.py \ + data.custom_cls.name=CustomRLHFDataset \ + custom_reward_function.path=recipe/retool/retool.py \ + custom_reward_function.name=compute_score \ + actor_rollout_ref.model.path=$model_path \ + actor_rollout_ref.model.use_remove_padding=True \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.use_kl_loss=$use_kl_loss \ + actor_rollout_ref.actor.kl_loss_coef=$kl_loss_coef \ + actor_rollout_ref.actor.clip_ratio_low=$clip_ratio_low \ + actor_rollout_ref.actor.clip_ratio_high=$clip_ratio_high \ + actor_rollout_ref.actor.clip_ratio_c=10.0 \ + actor_rollout_ref.actor.optim.lr=$actor_lr \ + actor_rollout_ref.actor.use_dynamic_bsz=True \ + actor_rollout_ref.actor.ppo_mini_batch_size=$ppo_mini_batch_size \ + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=$actor_max_token_len_per_gpu \ + actor_rollout_ref.actor.ulysses_sequence_parallel_size=$train_sp \ + actor_rollout_ref.actor.fsdp_config.param_offload=$offload \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=$offload \ + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=$log_prob_max_token_len_per_gpu \ + actor_rollout_ref.rollout.max_num_batched_tokens=$actor_max_token_len_per_gpu \ + actor_rollout_ref.rollout.name=vllm \ + actor_rollout_ref.rollout.mode=async \ + actor_rollout_ref.rollout.max_num_seqs=1024 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=$infer_tp \ + actor_rollout_ref.rollout.multi_turn.enable=True \ + actor_rollout_ref.rollout.multi_turn.max_user_turns=$max_turns \ + actor_rollout_ref.rollout.multi_turn.max_assistant_turns=$max_turns \ + actor_rollout_ref.rollout.multi_turn.tool_config_path=$tool_config_path \ + actor_rollout_ref.rollout.multi_turn.format=hermes \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.9 \ + actor_rollout_ref.rollout.n=$n_resp_per_prompt \ + actor_rollout_ref.rollout.val_kwargs.top_p=0.6 \ + actor_rollout_ref.rollout.val_kwargs.temperature=1.0 \ + actor_rollout_ref.rollout.val_kwargs.n=$n_resp_per_prompt_val \ + actor_rollout_ref.rollout.enable_chunked_prefill=True \ + actor_rollout_ref.rollout.enforce_eager=False \ + trainer.logger=['console'] \ + trainer.project_name=$project_name \ + trainer.experiment_name=$experiment_name \ + trainer.n_gpus_per_node=8 \ + trainer.val_before_train=False \ + trainer.log_val_generations=20 \ + trainer.nnodes=1 \ + trainer.save_freq=100 \ + trainer.default_local_dir=$default_local_dir \ + trainer.test_freq=20 \ + trainer.device=npu \ + actor_rollout_ref.actor.entropy_from_logits_with_chunking=True \ + actor_rollout_ref.ref.entropy_from_logits_with_chunking=True \ + actor_rollout_ref.actor.use_torch_compile=False \ + actor_rollout_ref.ref.use_torch_compile=False \ + actor_rollout_ref.actor.entropy_checkpointing=True \ + actor_rollout_ref.ref.entropy_checkpointing=True \ + actor_rollout_ref.ref.use_torch_compile=False \ + trainer.total_epochs=1 $@ > $LOG_FILE 2>&1 & diff --git a/verl/docs/ascend_tutorial/model_support/examples/ascend_sglang_best_practices.rst b/verl/docs/ascend_tutorial/model_support/examples/ascend_sglang_best_practices.rst new file mode 100644 index 0000000000000000000000000000000000000000..be5c93b2fe46e7d609d49c55a5b48dbac0a898d1 --- /dev/null +++ b/verl/docs/ascend_tutorial/model_support/examples/ascend_sglang_best_practices.rst @@ -0,0 +1,297 @@ +Ascend SGLang Best Practice +=================================== + +Last updated: 01/27/2026. + +.. _Qwen3-30B: https://github.com/verl-project/verl/blob/c98cb8cc/examples/grpo_trainer/run_qwen3moe-30b_sglang_megatron_npu.sh +.. _Qwen2.5-32B: https://github.com/verl-project/verl/blob/c98cb8cc/examples/grpo_trainer/run_qwen2-32b_sglang_fsdp_npu.sh +.. _doclink: https://github.com/verl-project/verl/blob/c98cb8cc/docs/ascend_tutorial/examples/ascend_sglang_best_practices.rst +引言 +---------------------------------- + +SGLang 是当前主流的高性能开源推理引擎, 昇腾已经全面原生支持该推理引擎在verl中使用, +仅需简单的构建流程，开发者即可完成环境构建，本文将提供两个经典用例来帮助开发者了解以下内容： + +1. 环境构建 +2. 模型训练与评估 +3. 性能采集 + +两个用例模型脚本以及其需要的硬件条件各自如下： +(注:verl近期进行了脚本清理与命名变更, 请根据实际commit id 对应文档`doclink`_ 及对应脚本进行构建避免链接失效) ++----------------------+---------------------+----------+------------------------+ +| 模型 | NPU型号 | 节点数量 | 训推后端 | ++======================+=====================+==========+========================+ +| `Qwen3-30B`_ | Atlas 800T A3 | 1 | SGLang + Megatron | ++----------------------+---------------------+----------+------------------------+ +| `Qwen2.5-32B`_ | Atlas 900 A2 | 2 | SGLang + FSDP | ++----------------------+---------------------+----------+------------------------+ + +环境构建 +----------------------------------- +我们在quickstart中提供了两种构建环境的方法, 1.从镜像文件DockerFile进行构建 2.从自定义Conda环境进行构建 + +在本实践中, 我们额外指定verl 的commit id 以避免引入其他问题 + +.. code-block:: bash + + cd verl + git checkout 772c224 +模型训练与评估 +----------------------------------- +1.模型数据准备 +^^^^^^^^^^^ +`Qwen3-30B`_ +^^^^^^^^^^^ +**下载模型权重** + +--local-dir: 模型保存路径 + +.. code-block:: bash + + export HF_ENDPOINT=https://hf-mirror.com + huggingface-cli download --resume-download Qwen/Qwen3-30B-A3B --local-dir /path/to/local_dir + +**下载数据集** + +.. code-block:: bash + + git clone https://www.modelscope.cn/datasets/AI-ModelScope/DAPO-Math-17k.git + +**HuggingFace To Megatron权重转换(可选)** + +.. code-block:: bash + + python scripts/converter_hf_to_mcore.py \ + --hf_model_path Qwen/Qwen3-30B-A3B \ + --output_path Qwen/Qwen3-30B-A3B-mcore \ + --use_cpu_initialization # Only work for MoE models +*注:verl当前已支持mbridge进行灵活的hf和mcore之间的权重转换,可以修改以下相关参数直接加载hf权重* + +.. code-block:: bash + + actor_rollout_ref.actor.megatron.use_dist_checkpointing=False + actor_rollout_ref.actor.megatron.use_mbridge=True + +`Qwen2.5-32B`_ +^^^^^^^^^^^ +**下载模型权重** + +--local-dir: 模型保存路径 + +.. code-block:: bash + + export HF_ENDPOINT=https://hf-mirror.com + huggingface-cli download --resume-download Qwen/Qwen2.5-32B --local-dir /path/to/local_dir + +**下载及处理数据集** + +.. code-block:: bash + + wget https://huggingface.co/datasets/agentica-org/DeepScaleR-Preview-Dataset/resolve/main/deepscaler.json + python recipe/r1_ascend/json_to_parquet.py --output_dir ./data/deepscaler --json_path path/to/deepscaler.json --train_data_ratio 0.9 + +2.训练 +^^^^^^^^^^^ +根据开发者实际路径配置情况修改模型训练脚本中的以下参数 + +.. code-block:: bash + + # Model Weights Paths + MODEL_PATH=Qwen/Qwen3-30B-A3B + MCORE_MODEL_PATH=Qwen/Qwen3-30B-A3B-mcore + RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"} + CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"} + + # File System Paths + TRAIN_FILE=$RAY_DATA_HOME/dataset/dapo-math-17k.parquet + TEST_FILE=$RAY_DATA_HOME/dataset/aime-2024.parquet + + #保存频率，-1默认不保存，如需评测请修改此参数 + trainer.save_freq=-1 + +对于单机任务 `Qwen3-30B`_ , 可以直接bash执行verl仓上示例脚本 + +.. code-block:: bash + + bash examples/grpo_trainer/run_qwen3moe-30b_sglang_megatron_npu.sh +对于多节点任务 `Qwen2.5-32B`_ ，我们推荐使用以下脚本进行大规模多节点训练拉起 + +.. code-block:: bash + + pkill -9 python + ray stop --force + rm -rf /tmp/ray + export RAY_DEDUP_LOGS=0 + export HYDRA_FULL_ERROR=1 + # TASK_QUEUE_ENABLE，下发优化，图模式设置为1，非图模式设置为2 + export TASK_QUEUE_ENABLE=1 + export HCCL_ASYNC_ERROR_HANDLING=0 + export HCCL_EXEC_TIMEOUT=3600 + export HCCL_CONNECT_TIMEOUT=3600 + + export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050 + export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050 + export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 + export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8 + # 修改为当前需要跑的用例路径 + DEFAULT_SH="./run_*.sh" + echo "Use $DEFAULT_SH" + + ulimit -n 32768 + mkdir logs + + NNODES=2 + NPUS_PER_NODE=8 + # 修改为对应主节点IP + MASTER_ADDR="IP FOR MASTER NODE" + # 修改为当前节点的通信网卡 + SOCKET_IFNAME="Your SOCKET IFNAME" + export HCCL_SOCKET_IFNAME="SOCKET IFNAME FOR CURRENT NODE" + export GLOO_SOCKET_IFNAME="SOCKET IFNAME FOR CURRENT NODE" + # 获取当前IP + CURRENT_IP=$(ifconfig $SOCKET_IFNAME | grep -Eo 'inet (addr:)?([0-9]{1,3}\.){3}[0-9]{1,3}' | awk '{print $NF}') + if [ "$MASTER_ADDR" = "$CURRENT_IP" ]; then + # 主节点启动 + ray start --head --port 6766 --dashboard-host=$MASTER_ADDR --node-ip-address=$CURRENT_IP --dashboard-port=8260 --resources='{"NPU": '$NPUS_PER_NODE'}' + + while true; do + ray_status_output=$(ray status) + npu_count=$(echo "$ray_status_output" | grep -oP '(?<=/)\d+\.\d+(?=\s*NPU)' | head -n 1) + npu_count_int=$(echo "$npu_count" | awk '{print int($1)}') + device_count=$((npu_count_int / $NPUS_PER_NODE)) + + # 判断device_count 是否与 NNODES 相等 + if [ "$device_count" -eq "$NNODES" ]; then + echo "Ray cluster is ready with $device_count devices (from $npu_count NPU resources), starting Python script." + ray status + bash $DEFAULT_SH + break + else + echo "Waiting for Ray to allocate $NNODES devices. Current device count: $device_count" + sleep 5 + fi + done + else + # 子节点尝试往主节点注册 ray 直到成功 + while true; do + # 尝试连接 ray 集群 + ray start --address="$MASTER_ADDR:6766" --resources='{"NPU": '$NPUS_PER_NODE'}' --node-ip-address=$CURRENT_IP + + # 检查连接是否成功 + ray status + if [ $? -eq 0 ]; then + echo "Successfully connected to the Ray cluster!" + break + else + echo "Failed to connect to the Ray cluster. Retrying in 5 seconds..." + sleep 5 + fi + done + fi + + sleep 600 + +DEFAULT_SH:修改为训练所用配置 sh 文件路径。在此案例中修改为 `Qwen2.5-32B`_ 路径。 + +NNODES 和 NPUS_PER_NODE:修改为使用节点数量和每个节点 NPU 数量。在此案例中分别为2和8。 + +MASTER_ADDR:修改为对应主节点 IP。即所有节点的 MASTER_ADDR 应该相同。 + +SOCKET_IFNAME, HCCL_SOCKET_IFNAME, GLOO_SOCKET_IFNAME: 修改为对应通信网卡，通信网卡可以通过以下命令获取： + +.. code-block:: bash + + ifconfig |grep "$(hostname -I |awk '{print $1}'|awk -F '.' '{print $0}')" -B 1|awk -F ':' '{print$1}' | head -1 | tail -1 + +3.模型评估 +^^^^^^^^^^^ + +不同模型步骤一致,仅以Qwen3-30b为例列举 + +我们通过 AISBenchmark 评估模型,该工具支持vllm/sglang多种推理后端的评估 + +**安装方法** + +.. code-block:: bash + + git clone https://gitee.com/aisbench/benchmark.git + cd benchmark + pip install -e . + +**下载评估数据集** + +.. code-block:: bash + + cd path/to/benchmark/ais_bench/datasets + wget http://opencompass.oss-cn-shanghai.aliyuncs.com/datasets/data/math.zip + unzip math.zip + rm math.zip + +**修改AISBench配置代码使能sglang推理评测** + +打开 benchmark/ais_bench/benchmark/configs/models/vllm_api/vllm_api_stream_chat.py 文件，这是推理配置文件 + +.. code-block:: bash + + from ais_bench.benchmark.models import VLLMCustomAPIChatStream + from ais_bench.benchmark.utils.model_postprocessors import extract_non_reasoning_content + from ais_bench.benchmark.clients import OpenAIChatStreamClient, OpenAIChatStreamSglangClient + + models = [ + dict( + attr="service", + type=VLLMCustomAPIChatStream, + abbr='sgl-api-stream-chat', + path="/path/to/Qwen3-30B", # 修改为 Qwen3-30B 模型路径 + model="qwen3-30b", + request_rate = 0, + max_seq_len=2048, + retry = 2, + host_ip = "localhost", # 推理服务的IP + host_port = 8005, # 推理服务的端口 + max_out_len = 8192, # 最大输出tokens长度 + batch_size=48, # 推理的最大并发数 + trust_remote_code=False, + custom_client=dict(type=OpenAIChatStreamSglangClient), #使用sglang客户端 + generation_kwargs = dict( + temperature = 0, + seed = 1234, + ), + pred_postprocessor=dict(type=extract_non_reasoning_content) + ) + ] + + +**启动sglang_server服务** + +.. code-block:: bash + + python -m sglang.launch_server --model-path "/path/to/Qwen3-30B" --tp-size 4 --dp-size 1 --port 8005 + +**启动sglang_client评测** + +.. code-block:: bash + + ais_bench --models vllm_api_stream_chat --datasets math500_gen_0_shot_cot_chat_prompt + +**评测结果** + +经过训练,模型在Math-500上的评分显著上升 + ++------+----------------------+---------+----------+------+----------------------+ +| iter | dataset | version | metric | mode | sgl-api-stream-chat | ++======+======================+=========+==========+======+======================+ +| 0 | math_prm800k_500 | c4b6f0 | accuracy | gen | 84.4 | ++------+----------------------+---------+----------+------+----------------------+ +| 150 | math_prm800k_500 | c4b6f0 | accuracy | gen | 91.7 | ++------+----------------------+---------+----------+------+----------------------+ + +性能采集 +----------------------------------- +关于NPU profiling的详细文档请参考 `ascend_profiling_zh `_ + +在 `Qwen3-30B`_ 的脚本中提供了基本的采集性能选项PROF_CONFIG，默认设置 global_profiler.steps=null 关闭采集， 开发者可根据实际需要进行参数修改 + +采集完成后，开发者可以使用 `MindStudio Insight `_ 进行数据解析 + +注: verl框架侧进行采集全量 Profiling 产生海量且重复的算子记录，可以根据文档修改代码仅采集关键阶段 \ No newline at end of file diff --git a/verl/docs/ascend_tutorial/model_support/examples/dapo_multi_model_optimization_practice.md b/verl/docs/ascend_tutorial/model_support/examples/dapo_multi_model_optimization_practice.md new file mode 100644 index 0000000000000000000000000000000000000000..1cde03f8861bb4343774b13b2d0e44dc40a016d7 --- /dev/null +++ b/verl/docs/ascend_tutorial/model_support/examples/dapo_multi_model_optimization_practice.md @@ -0,0 +1,357 @@ +# DAPO multi model optimization practice + +## DAPO 介绍 + +Last updated: 03/04/2026. + +DAPO的论文可以参考：[DAPO](https://arxiv.org/pdf/2503.14476)，其中包含以下几个关键技术。 + +* ​**Clip-Higher**​: 通过对重要性采样比的上限剪裁促进了系统的多样性并避免了熵坍缩（Entropy Collapse）。 +* ​**Dynamic Sampling**​: 提高了训练效率和稳定性。DAPO出了一种执行动态采样的策略，并过滤掉准确率等于1和0的提示组，从而保持批次间具有有效梯度的提示数量一致。 +* ​**Token-level Policy Gradient Loss**​: 在长链思维强化学习 (long-CoT RL) 场景中至关重要。 +* ​**Overlong Reward Shaping**​: 减少奖励噪声并稳定了训练。 + +在verl中，可以进行如下设置，从而进行DAPO算法的运行。 + +- **奖励模型的管理策略为 DAPO** + 在dapo算法中，必须配置成dapo。 + +``` +reward_model.reward_manager.name=dapo +``` + +- **Clip-Higher 更高裁剪** + `clip_ratio_low` 和 `clip_ratio_high` 用于指定 DAPO 目标函数中的 $\varepsilon_{\text {low }}$ 和 $\varepsilon_{\text {high }}$。 + +``` +clip_ratio_low=0.2 # 裁剪比例下限，默认值为0.2 +clip_ratio_high=0.28 # 裁剪比例上限，默认值为0.28 +``` + +- **动态采样的相关配置** + 将 `filter_groups.enable` 设置为 `True` 会过滤掉输出 `metric` 完全相同的组，例如对于 `acc` 指标，过滤掉输出准确率全部为 1 或 0 的组。 + 训练器会使用 `gen_batch_size` 进行重复采样，直到生成足够数量的符合条件的组，或者达到 `max_num_gen_batches` 所指定的上限为止。 + +``` +data.gen_batch_size=${gen_prompt_bsz} +algorithm.filter_groups.enable=${enable_filter_groups} # 动态采样开关 +algorithm.filter_groups.metric=${filter_groups_metric} # 使用准确率作为过滤标准 +algorithm.filter_groups.max_num_gen_batches=${max_num_gen_batches} # 最大生成批次数量,最多重复生成数据的次数 +``` + +- **Token-level Loss** + 将 `loss_agg_mode` 设置为 `token-mean` 意味着计算一个批次中所有序列内所有 token 的（策略梯度）损失的平均值。 + +``` +actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} +# 注意：“token-mean”是默认行为。 +``` + +- **奖励模型对超长回答的惩罚配置** + 将 `overlong_buffer.enable` 设置为 `True` 将对输出长度过长但仍未超过硬上下文限制的输出进行惩罚。具体来说，当输出的长度超过 `max_response_length - overlong_buffer.len` 且超出 `0` 到 `overlong_buffer.len` 个 token 时，惩罚值会从 `0` 线性增加到 `overlong_buffer.penalty_factor`。 + +``` +reward_model.overlong_buffer.enable=${enable_overlong_buffer} # 启用超长缓冲区惩罚,开启对超长输出的惩罚机制 +reward_model.overlong_buffer.len=${overlong_buffer_len} # 缓冲区长度,定义缓冲区的toke,最大惩罚强度 +reward_model.overlong_buffer.penalty_factor=${overlong_penalty_factor} #惩罚因子,最大惩罚强度 +``` + +相关参数涉及的代码可以参考：[Recipe: Decoupled Clip and Dynamic Sampling Policy Optimization (DAPO)](https://github.com/verl-project/verl-recipe/blob/main/dapo/README.md) + +## 硬件要求 + +当前支持Atlas 800T A3 与 Atlas 900 A3 SuperPoD。完成跑完本次最佳实践需要 1 台 Atlas 900 A3 SuperPoD。关键软件版本可以参考：[Ascend Quickstart](https://github.com/verl-project/verl/blob/02e059ea/docs/ascend_tutorial/quick_start/ascend_quick_start.rst) + +## 安装基础环境 + +| software | version | +| ------------ | ---------------------------------------------------------- | +| Python | >=3.10, <3.12 | +| CANN | ==8.5 | +| torch | ==2.8.0 | +| torch_npu | ==2.8.0 | +| verl | v0.7.1 commitId=02e059ea18f5adf9768c5d9c280456cdfdfeda01 | +| vllm | v0.13.0 | +| vllm-ascend | v0.13.0 | +| transformers | 4.57.6 | + +在本实践中, 我们通过指定 verl 的commit id 以避免引入其他问题 +``` +cd verl +git checkout release/v0.7.1 +# 指定相应的recipe版本 +git submodule update --init --recursive recipe +cd recipe +git checkout main +``` + +## 模型训练 + +### 数据集准备 + +Geometry3k 数据集是由加利福尼亚大学洛杉矶分校与浙江大学联合研发的几何领域专用数据集，核心面向视觉问答（VQA）任务展开研究与模型训练。该数据集总计包含 3002 个样本，采用图像和文本两种模态数据形式构建，其中文本模态涵盖各类几何问题描述，图像则以可视化图表呈现问题中的几何图形信息，包括三角形、圆形、四边形等基础几何形状，以及不同图形间的位置、嵌套、相交等关联关系。可以从Hugging Face库下载对应的原始数据集：[Geometry3k ](https://huggingface.co/datasets/hiyouga/geometry3k) + +```shell +# 下载原始数据并预处理 +python ./examples/data_preprocess/geo3k.py --local_dir=./data/geo3k +``` + +### 权重下载 + +从Hugging Face库下载对应的模型权重：[Qwen3-VL-30B-A3B-Instruct](https://huggingface.co/Qwen/Qwen3-VL-30B-A3B-Instruct/tree/main +) + +### jemalloc安装 + +为了确保 Ray 进程能够正常回收内存，需要安装并使能 jemalloc 库进行内存管理。 + +#### Ubuntu 操作系统 + +通过操作系统源安装jemalloc（注意： 要求ubuntu版本>=20.04）： + +```shell +sudo apt install libjemalloc2 +``` + +在启动任务前执行如下命令通过环境变量导入jemalloc，需先通过 **find /usr -name libjemalloc.so.2** 确认文件是否存在 ： + +```shell +# arm64架构 +export LD_PRELOAD=/usr/lib/aarch64-linux-gnu/libjemalloc.so.2 +# x86_64架构 +export LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libjemalloc.so.2 +``` + +#### OpenEuler 操作系统 + +执行如下命令通过操作系统源安装jemalloc + +```shell +yum install jemalloc +``` + +如果上述方法无法正常安装，可以通过源码编译安装 前往jemalloc官网下载最新稳定版本，官网地址:https://github.com/jemalloc/jemalloc/releases/ + +```shell +tar -xvf jemalloc-{version}.tar.bz2 +cd jemalloc-{version} +./configure --prefix=/usr/local +make +make install +``` + +### 全局变量导入 + +- 为了确保 Ray 进程能够正常回收内存，需要安装并使能 jemalloc 库进行内存管理，用于更好管理内存，避免长跑过程中内存 OOM。 + +``` +# 根据实际安装路径设置 jemalloc 环境变量，例如安装路径为:/usr/local/lib/libjemalloc.so.2(可通过 find /usr -name libjemalloc.so.2 确认文件是否存在) +export LD_PRELOAD=/usr/local/lib/libjemalloc.so.2 +``` + +- 某些模型是通过 vllm ascend 进行优化的。但在某些情况下，优化后的模型可能并不适用。此时，将此值设置为 0 即可禁用优化后的模型。 + +``` +export USE_OPTIMIZED_MODEL=0 +``` + +- 启用vLLM V1 + +``` +export VLLM_USE_V1=1 +``` + +- 昇腾多卡通信的兜底配置，延长连接超时时间，避免集群环境下训练启动因连接慢而失败 + +``` +export HCCL_CONNECT_TIMEOUT=5400 +``` + +- 控制 vLLM 在昇腾芯片上是否启用NZ优化 + +``` +export VLLM_ASCEND_ENABLE_NZ=0 +``` + +### 训练 +``` +# Model Weights Paths +MODEL_PATH=hf_weights/Qwen3-VL-30B-A3B-Instruct +RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"} +CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"} + +# File System Paths +TRAIN_FILE=$RAY_DATA_HOME/datasets/geo3k/train.parquet +TEST_FILE=$RAY_DATA_HOME/datasets/geo3k/test.parquet + +# 保存频率，-1默认不保存，如需评测请修改此参数 +trainer.save_freq=-1 +``` + +- 对于单机任务 Qwen3-VL-30B ，可以使用以下脚本将训练拉起。 + +``` +pkill -9 python +ray stop --force +rm -rf /tmp/ray +export VLLM_USE_V1=1 +export HCCL_CONNECT_TIMEOUT=5400 +export VLLM_ASCEND_ENABLE_NZ=0 +export LD_PRELOAD=/usr/local/lib/libjemalloc.so.2 +# Some models are optimized by vllm ascend. While in some case, e.g. rlhf training, +# the optimized model may not be suitable. In this case, set this value to 0 to disable the optimized model. +export USE_OPTIMIZED_MODEL=0 +export CPU_AFFINITY_CONF=2 +export HCCL_OP_EXPANSION_MODE="AIV" +export VLLM_VERSION="0.13.0" + +# 修改为对应主节点IP +MASTER_ADDR="IP FOR MASTER NODE" +# 每个节点中 NPU 数 +NPUS_PER_NODE=16 +ray start --head --port 6766 --dashboard-host=$MASTER_ADDR --dashboard-port=8260 --resources='{"NPU": '$NPUS_PER_NODE'}' + +bash recipe/dapo/run_dapo_qwen3_vl_30b_fsdp2_npu.sh +``` +- 对于多节点任务 Qwen3-VL-30B ，我们推荐使用以下脚本进行大规模多节点训练拉起，根据实际需要修改`NNODES`与`NPUS_PER_NODE` ，并修改配置脚本中参数`trainer.nnodes`和`trainer.n_gpus_per_node`与之相对应。 + +``` +pkill -9 python +ray stop --force +rm -rf /tmp/ray +export VLLM_USE_V1=1 +export HCCL_CONNECT_TIMEOUT=5400 +export VLLM_ASCEND_ENABLE_NZ=0 +export LD_PRELOAD=/usr/local/lib/libjemalloc.so.2 +# Some models are optimized by vllm ascend. While in some case, e.g. rlhf training, +# the optimized model may not be suitable. In this case, set this value to 0 to disable the optimized model. +export USE_OPTIMIZED_MODEL=0 +export CPU_AFFINITY_CONF=2 +export HCCL_OP_EXPANSION_MODE="AIV" +export VLLM_VERSION="0.13.0" + +# 修改为当前需要跑的用例路径 +DEFAULT_SH="./run_*.sh" +echo "Use $DEFAULT_SH" + +ulimit -n 32768 +mkdir logs + +# 集群中计算节点数 +NNODES=2 +# 每个节点中 NPU 数 +NPUS_PER_NODE=8 +# 修改为对应主节点IP +MASTER_ADDR="IP FOR MASTER NODE" +# 修改为当前节点的通信网卡 +SOCKET_IFNAME="Your SOCKET IFNAME" +export HCCL_SOCKET_IFNAME="SOCKET IFNAME FOR CURRENT NODE" +export GLOO_SOCKET_IFNAME="SOCKET IFNAME FOR CURRENT NODE" +# 获取当前IP +CURRENT_IP=$(ifconfig $SOCKET_IFNAME | grep -Eo 'inet (addr:)?([0-9]{1,3}\.){3}[0-9]{1,3}' | awk '{print $NF}') +if [ "$MASTER_ADDR" = "$CURRENT_IP" ]; then + # 主节点启动 + ray start --head --port 6766 --dashboard-host=$MASTER_ADDR --node-ip-address=$CURRENT_IP --dashboard-port=8260 --resources='{"NPU": '$NPUS_PER_NODE'}' + + while true; do + ray_status_output=$(ray status) + npu_count=$(echo "$ray_status_output" | grep -oP '(?<=/)\d+\.\d+(?=\s*NPU)' | head -n 1) + npu_count_int=$(echo "$npu_count" | awk '{print int($1)}') + device_count=$((npu_count_int / $NPUS_PER_NODE)) + + # 判断device_count 是否与 NNODES 相等 + if [ "$device_count" -eq "$NNODES" ]; then + echo "Ray cluster is ready with $device_count devices (from $npu_count NPU resources), starting Python script." + ray status + bash $DEFAULT_SH + break + else + echo "Waiting for Ray to allocate $NNODES devices. Current device count: $device_count" + sleep 5 + fi + done +else + # 子节点尝试往主节点注册 ray 直到成功 + while true; do + # 尝试连接 ray 集群 + ray start --address="$MASTER_ADDR:6766" --resources='{"NPU": '$NPUS_PER_NODE'}' --node-ip-address=$CURRENT_IP + + # 检查连接是否成功 + ray status + if [ $? -eq 0 ]; then + echo "Successfully connected to the Ray cluster!" + break + else + echo "Failed to connect to the Ray cluster. Retrying in 5 seconds..." + sleep 5 + fi + done +fi + +sleep 600 +``` +DEFAULT_SH: 修改为训练所用配置 sh 文件路径。在此案例中修改为 [Qwen3_VL_30B](https://github.com/verl-project/verl-recipe/blob/main/dapo/run_dapo_qwen3_vl_30b_fsdp2_npu.sh) 路径。 + +NNODES 和 NPUS_PER_NODE: 修改为使用节点数量和每个节点 NPU 数量。在此案例中分别为2和8。 + +MASTER_ADDR:修改为对应主节点 IP。即所有节点的 MASTER_ADDR 应该相同。 + +SOCKET_IFNAME, HCCL_SOCKET_IFNAME, GLOO_SOCKET_IFNAME: 修改为对应通信网卡，通信网卡可以通过以下命令获取： + +``` +ifconfig |grep "$(hostname -I |awk '{print $1}'|awk -F '.' '{print $0}')" -B 1|awk -F ':' '{print$1}' | head -1 | tail -1 +``` + +## 优化参考 + +- **启动动态批次大小** + 根据单 GPU 的最大 Token 总数（ppo_max_token_len_per_gpu）动态调整批次大小 + +``` +actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} +actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} +actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} +``` + +- **单个 GPU 能处理的最大 Token 总数** + 当`use_dynamic_bsz=True`时，单 GPU 在一个微批次中能处理的最大 Token 数量 + +``` +actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} +actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} +actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} +``` + +- **单个 GPU 微批次大小** + 当`use_dynamic_bsz=True`时，框架会以该值为​初始批次大小​，再根据`ppo_max_token_len_per_gpu`向上 / 向下调整 + +``` +actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 +actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=2 +actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=2 +``` + +- **启用 FSDP2 框架** + “将模型参数、梯度、优化器状态分片存储在不同 GPU 上”，避免单卡加载全量模型导致显存溢出。 + +``` +# 启用 FSDP2 框架 +actor_rollout_ref.actor.strategy=fsdp2 +actor_rollout_ref.ref.strategy=fsdp2 +critic.strategy=fsdp2 + +# 仅用于 FSDP2：前向传播后重新分片以减少内存占用。 +actor_rollout_ref.actor.fsdp_config.reshard_after_forward=True +# 仅用于 FSDP2：是否在模型前向传播后重新分片以节省内存。 +actor_rollout_ref.ref.fsdp_config.reshard_after_forward=True +``` + +- **启用专家并行配置** + 指定有多少个 GPU用于并行计算不同的专家网络 + +``` +# MoE 架构 Actor 模型的专家并行配置 +actor_rollout_ref.rollout.expert_parallel_size=8 +``` + + diff --git a/verl/docs/ascend_tutorial/model_support/examples/gspo_optimization_practice.md b/verl/docs/ascend_tutorial/model_support/examples/gspo_optimization_practice.md new file mode 100644 index 0000000000000000000000000000000000000000..24e60dd8d12f2140f6eec29c7237c647d9c4b9f9 --- /dev/null +++ b/verl/docs/ascend_tutorial/model_support/examples/gspo_optimization_practice.md @@ -0,0 +1,412 @@ +# NPU Qwen3-32B GSPO Optimization Practice + +Last updated: 02/26/2026. + +本文章对应脚本地址：[qwen3_8b_gspo_npu](https://github.com/verl-project/verl/blob/main/examples/gspo_trainer/run_qwen3_8b_fsdp.sh) + +## 算法适配 + +GSPO通过将优化颗粒度从**token级**提升到**sequence级**，规避了GRPO会遇到的**方差急剧增大**导致训练不稳定的情况，增加了训练的稳定性，同时该算法也在一定程度上提升了算法的收敛速度。 + +想要成功在verl仓库中成功调用到GSPO算法，需要进行如下的必要配置 + +```python +# 核心算法配置 +algorithm.adv_estimator=grpo \ # 使用GRPO优势估计器 +algorithm.use_kl_in_reward=False \ # 不在奖励中添加KL惩罚 +# GSPO策略损失模式 +actor_rollout_ref.actor.policy_loss.loss_mode=gspo \ # 启用GSPO策略损失 +# 极小裁剪范围（GSPO特色） +actor_rollout_ref.actor.clip_ratio_low=0.0003 \ # 裁剪下界，论文推荐值 +actor_rollout_ref.actor.clip_ratio_high=0.0004 \ # 裁剪上界，论文推荐值 +# KL配置（GSPO不使用KL loss） +actor_rollout_ref.actor.use_kl_loss=False \ # 禁用KL损失 +actor_rollout_ref.actor.kl_loss_coef=0.0 \ # KL损失系数设为0 +# 序列级损失聚合模式（GSPO核心） +actor_rollout_ref.actor.loss_agg_mode=seq-mean-token-mean \ # 序列级平均，GSPO论文推荐 +# 批次配置 +actor_rollout_ref.rollout.n=16 \ # 每个prompt生成16个响应（组采样） +``` + +一般选择入口函数为 `verl.trainer.main_ppo` + +## 基础环境 + +当前支持Atlas 800T A3 与 Atlas 900 A3 SuperPoD。完成跑完本次最佳实践需要 4台Atlas 800T A3。关键软件版本可以参考：[Ascend Quickstart](https://github.com/verl-project/verl/blob/main/docs/ascend_tutorial/quick_start/ascend_quick_start.rst) + +### 安装基础环境 + +| software | version | +| ------------ | ---------------------------------------------------------- | +| Python | >=3.10, <3.12 | +| CANN | ==8.3.RC1 | +| torch | ==2.7.1 | +| torch_npu | ==2.7.1 | +| verl | main分支 commitId=252d76908b903ad8fb6969eb3a5e5f873c95ea2b | +| vllm | v0.11.0 | +| vllm-ascend | v0.11.0-dev | +| transformers | 4.57.3 | + +在本实践中, 我们通过指定 verl 的commit id 以避免引入其他问题 + +```bash +cd verl +git checkout 252d76908b903ad8fb6969eb3a5e5f873c95ea2b +# 指定相应的recipe版本 +git submodule update --init --recursive recipe +``` + +### 权重获取 + +从Hugging Face库下载对应的模型权重：[Qwen/Qwen3-32B · Hugging Face](https://huggingface.co/Qwen/Qwen3-32B) + +### 数据集准备 + +```bash +# 下载math-17k数据集 +git clone https://huggingface.co/datasets/BytedTsinghua-SIA/DAPO-Math-17k + +# 下载AIME_2024测试数据集 +git clone https://huggingface.co/datasets/Maxwell-Jia/AIME_2024 +``` + +### jemalloc安装 + +为了确保 Ray 进程能够正常回收内存，需要安装并使能 jemalloc 库进行内存管理。 + +#### Ubuntu 操作系统 + +通过操作系统源安装jemalloc（注意： 要求ubuntu版本>=20.04）： + +```shell +sudo apt install libjemalloc2 +``` + +在启动任务前执行如下命令通过环境变量导入jemalloc，需先通过 **find /usr -name libjemalloc.so.2** 确认文件是否存在 ： + +```shell +# arm64架构 +export LD_PRELOAD=/usr/lib/aarch64-linux-gnu/libjemalloc.so.2 +# x86_64架构 +export LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libjemalloc.so.2 +``` + +#### OpenEuler 操作系统 + +执行如下命令通过操作系统源安装jemalloc + +```shell +yum install jemalloc +``` + +如果上述方法无法正常安装，可以通过源码编译安装 前往jemalloc官网下载最新稳定版本，官网地址:https://github.com/jemalloc/jemalloc/releases/ + +```shell +tar -xvf jemalloc-{version}.tar.bz2 +cd jemalloc-{version} +./configure --prefix=/usr/local +make +make install +``` + +在启动任务前执行如下命令通过环境变量导入jemalloc： + +```shell +#根据实际安装路径设置环境变量，例如安装路径为:/usr/local/lib/libjemalloc.so.2,可通过以下命令来设置环境变量(可通过 find /usr -name libjemalloc.so.2 确认文件是否存在) +export LD_PRELOAD=/usr/lib/aarch64-linux-gnu/libjemalloc.so.2 +``` + +### 多机任务拉起 + +针对本实践提供的多机任务，可用下面的脚本拉起 + +```bash +pkill -9 python +ray stop --force +rm -rf /tmp/ray + +export RAY_DEDUP_LOGS=0 +export HYDRA_FULL_ERROR=1 +export TASK_QUEUE_ENABLE=1 +export HCCL_EXEC_TIMEOUT=3600 +export HCCL_CONNECT_TIMEOUT=3600 +export HCCL_ASYNC_ERROR_HANDLING=0 +export CPU_AFFINITY_CONF=1 +export VLLM_USE_V1=1 +export VLLM_ATTENTION_BACKEND=XFORMERS +export VLLM_ASCEND_ENABLE_FLASHCOMM=1 +export VLLM_ASCEND_ENABLE_PREFETCH_MLP=1 +export VLLM_ASCEND_ENABLE_DENSE_OPTIMIZE=1 +export LD_PRELOAD=/usr/local/lib/libjemalloc.so.2 + +# 修改为当前需要跑的用例路径 +DEFAULT_SH="./run_*.sh" +echo "Use $DEFAULT_SH" + +ulimit -n 32768 +mkdir logs + +NNODES=4 +NPUS_PER_NODE=16 +# 修改为对应主节点IP +MASTER_ADDR="IP FOR MASTER NODE" +# 修改为当前节点的通信网卡 +SOCKET_IFNAME="Your SOCKET IFNAME" +export HCCL_SOCKET_IFNAME="SOCKET IFNAME FOR CURRENT NODE" +export GLOO_SOCKET_IFNAME="SOCKET IFNAME FOR CURRENT NODE" +# 获取当前IP +CURRENT_IP=$(ifconfig $SOCKET_IFNAME | grep -Eo 'inet (addr:)?([0-9]{1,3}\.){3}[0-9]{1,3}' | awk '{print $NF}') +if [ "$MASTER_ADDR" = "$CURRENT_IP" ]; then + # 主节点启动 + ray start --head --port 6766 --dashboard-host=$MASTER_ADDR --node-ip-address=$CURRENT_IP --dashboard-port=8260 --resources='{"NPU": '$NPUS_PER_NODE'}' + + while true; do + ray_status_output=$(ray status) + npu_count=$(echo "$ray_status_output" | grep -oP '(?<=/)\d+\.\d+(?=\s*NPU)' | head -n 1) + npu_count_int=$(echo "$npu_count" | awk '{print int($1)}') + device_count=$((npu_count_int / $NPUS_PER_NODE)) + + # 判断device_count 是否与 NNODES 相等 + if [ "$device_count" -eq "$NNODES" ]; then + echo "Ray cluster is ready with $device_count devices (from $npu_count NPU resources), starting Python script." + ray status + bash $DEFAULT_SH + break + else + echo "Waiting for Ray to allocate $NNODES devices. Current device count: $device_count" + sleep 5 + fi + done +else + # 子节点尝试往主节点注册 ray 直到成功 + while true; do + # 尝试连接 ray 集群 + ray start --address="$MASTER_ADDR:6766" --resources='{"NPU": '$NPUS_PER_NODE'}' --node-ip-address=$CURRENT_IP + + # 检查连接是否成功 + ray status + if [ $? -eq 0 ]; then + echo "Successfully connected to the Ray cluster!" + break + else + echo "Failed to connect to the Ray cluster. Retrying in 5 seconds..." + sleep 5 + fi + done +fi + +sleep 600 +``` + +DEFAULT_SH:修改为训练所用配置 sh 文件路径。在此案例中修改为 [Qwen3-8B](https://github.com/verl-project/verl/blob/main/examples/gspo_trainer/run_qwen3_8b_fsdp.sh) 路径。 + +NNODES 和 NPUS_PER_NODE:修改为使用节点数量和每个节点 NPU 数量。在此案例中分别为4和16。 + +MASTER_ADDR:修改为对应主节点 IP。即所有节点的 MASTER_ADDR 应该相同。 + +SOCKET_IFNAME, HCCL_SOCKET_IFNAME, GLOO_SOCKET_IFNAME: 修改为对应通信网卡，通信网卡可以通过以下命令获取： + +``` +ifconfig |grep "$(hostname -I |awk '{print $1}'|awk -F '.' '{print $0}')" -B 1|awk -F ':' '{print$1}' | head -1 | tail -1 +``` + +## 性能调优 + +优化从训练、推理、调度和其他四个方面入手。 + +### 训练 + +#### 动态bsz + +```bash +actor_ppo_max_token_len=$(((max_prompt_length + max_response_length) / sp_size)) +infer_ppo_max_token_len=$(((max_prompt_length + max_response_length) / sp_size)) +``` + +**这个优化点主要调整上面这两个参数，不过需要注意这两个参数调整的太大会导致OOM** + +**主要调整** `actor_ppo_max_token_len`,调大了会降低训练的耗时，调整 `infer_ppo_max_token_len`没有明显的收益，可以不动 + +**这两个参数的作用介绍如下：** + +**这两个参数用于控制动态批处理(dynamic batch size)模式下每个GPU处理的最大token数量** + +- **`actor_ppo_max_token_len`**: Actor模型在PPO更新(前向+反向传播)时每个GPU能处理的最大token数 +- **`infer_ppo_max_token_len`**: 推理阶段(Reference policy和Rollout)计算log概率时每个GPU能处理的最大token数 + +### 推理 + +#### ACLgraph+FULL_DECODE_ONLY + +推理算子下发方面的优化，平均能有 `15%~20%`左右的性能收益。 + +先看单开**ACLgraph**，如下： + +```bash +# 开启ACLgraph+FULL_DECODE_ONLY（注意：当设置此参数为False时，TASK_QUEUE_ENABLE必须设置为1，不然会报错） +actor_rollout_ref.rollout.enforce_eager=False \ +actor_rollout_ref.rollout.engine_kwargs.vllm.compilation_config.cudagraph_capture_sizes='[8,16,32,64,128]' \ +actor_rollout_ref.rollout.engine_kwargs.vllm.compilation_config.cudagraph_mode='FULL_DECODE_ONLY' +``` + +`FULL_DECODE_ONLY`开启成功后有如下输出： + +![FULL_DECODE_ONLY result](https://github.com/wucong25/verl-data/blob/main/ascend_acl_graph.png) + +**`cudagraph_capture_sizes`参数设置指南** + +cudagraph_capture_sizes设置的值对应的是批大小，这里的批大小不是配置里的DP域对应的那个批次大小，这里是相较于vllm来说的批大小，单位为**token** + +默认生成的算法如下，可做参考 + +![cudagraph_capture_sizes](https://github.com/wucong25/verl-data/blob/main/ascend_set_cudagraph_sizes.png) + +##### 推理后端切换 + +使用方式：`export VLLM_ATTENTION_BACKEND=XFORMERS` + +![VLLM_ATTENTION_BACKEND](https://github.com/wucong25/verl-data/blob/main/ascend_vllm_attn_backend.png) + +注：需要注意某些后端在一些比较老的vllm-ascend版本内并不支持 + +##### 使能vllm v1版本 + +使用方式：`export VLLM_USE_V1=1` + +可以常开，一般都是正收益。 + +### 调度 + +#### AIV + +打开方式：设置 `export HCCL_OP_EXPANSION_MODE="AIV"` + +HCCL_OP_EXPANSION_MODE环境变量用于配置通信算法的编排展开位置，支持如下取值： + +- AI_CPU：代表通信算法的编排展开位置在Device侧的AI CPU计算单元。 +- AIV：代表通信算法的编排展开位置在Device侧的Vector Core计算单元。 +- HOST：代表通信算法的编排展开位置为Host侧CPU，Device侧根据硬件型号自动选择相应的调度器。 +- HOST_TS：代表通信算法的编排展开位置为Host侧CPU，Host向Device的Task Scheduler下发任务，Device的Task Scheduler进行任务调度执行。 + +下面介绍两种展开机制 + +##### HOST展开 + +image-20260113194257095 + +- 软件栈工作在hostcpu，通信算法展开一个个task +- 每个task调用runtime接口，下发到device的rtsqueue +- STARS从rstqueue上顺序拿取task +- 根据task类型分别调用掉SDMA和RDMA引擎。 + **单算子瓶颈**：hostbound 每个task提交是2~5us，一个通信算子有几百个task，单算子场景不会在device上缓存，下发一个执行一个 + +##### AICpu机制展开 + +image-20260113194333218 + +- host侧不下发一个个task，把通信算子作为一个个kernel，放在通信算子kernel的队列上去。 +- STARS调度kernel队列流上的kernel，把kernel放到AiCPU上去执行。 +- AICPU调用函数（kernel），用一个线程执行kernel 函数，在函数内把通信task展开，把task放到rstqueue上，STARS调用。 +- 降低host和aicpu交互，由几百次降低为一次。 +- task的提交在AICPU上提交，做了提交的部分合并。 + +#### TASK_QUEUE_ENABLE + +**使用方式：**`export TASK_QUEUE_ENABLE=2` + +TASK_QUEUE_ENABLE，下发优化，图模式设置为1（即开启图模式的时候这个要设置为1），非图模式设置为2 + +示意图： + +![ascend task queue](https://github.com/wucong25/verl-data/blob/main/ascend_task_queue2.png) + +##### 绑核优化 + +**使用方式：**`export CPU_AFFINITY_CONF=1` + +详细设置原理可看：https://www.hiascend.com/document/detail/zh/Pytorch/600/ptmoddevg/trainingmigrguide/performance_tuning_0059.html + +### 其他 + +以下内容汇总了若干全局环境变量的调优配置。由于这些参数在训练阶段与推理阶段往往都能带来正向收益，且目前尚缺乏足够精细的消融实验来严格区分它们各自对训练或推理的贡献占比，故统一归拢在此，供后续持续监控与进一步拆解分析。 + +#### 使能jemalloc + +使用方式（注意需要先安装jemalloc库）：`export LD_PRELOAD=/usr/local/lib/libjemalloc.so.2` + +**安装使用教程：**[MindSpeed-RL/docs/install_guide.md · Ascend/MindSpeed-RL - AtomGit | GitCode](https://gitcode.com/Ascend/MindSpeed-RL/blob/master/docs/install_guide.md#高性能内存库-jemalloc-安装) + +#### 多流复用 + +内存方面有优化 + +使能方式：`export MULTI_STREAM_MEMORY_REUSE=1` + +原理介绍：https://www.hiascend.com/document/detail/zh/Pytorch/600/ptmoddevg/trainingmigrguide/performance_tuning_0040.html + +#### VLLM_ASCEND_ENABLE_FLASHCOMM + +使用方式：`export VLLM_ASCEND_ENABLE_FLASHCOMM=1` + +启用昇腾 NPU 特有的FLASHCOMM高速通信优化技术 + +地址：https://vllm-ascend.readthedocs.io/zh-cn/latest/user_guide/release_notes.html + +#### VLLM_ASCEND_ENABLE_DENSE_OPTIMIZE + +使用方式：`export VLLM_ASCEND_ENABLE_DENSE_OPTIMIZE=1` + +启用昇腾 NPU针对大模型推理的稠密计算优化 + +地址：https://vllm-ascend.readthedocs.io/zh-cn/latest/user_guide/release_notes.html + +#### VLLM_ASCEND_ENABLE_PREFETCH_MLP + +使用方式：`export VLLM_ASCEND_ENABLE_PREFETCH_MLP=1` + +启用 MLP 层的权重预取机制 + +image-20251124173132677 + +### verl框架参数设置 + +以下是内存方面的一些设置开关（注意，这个里面的优化都或多或少会导致吞吐量有一定程度的劣化） + +```bash +# 梯度检查点 (Gradient Checkpointing) +# 作用: 通过重新计算激活值来节省显存,以计算换内存。在前向传播时不保存中间激活值,反向传播时重新计算,可以显著降低显存占用,允许使用更大的batch size。 +actor_rollout_ref.model.enable_gradient_checkpointing=True \ + +# 参数卸载 (Parameter Offload) +# 作用: 将模型参数卸载到CPU内存,训练时再加载回GPU。 +actor_rollout_ref.actor.fsdp_config.param_offload=True \ +actor_rollout_ref.ref.fsdp_config.param_offload=True \ + +# 优化器状态卸载 (Optimizer Offload) +# 作用: 将优化器状态(如Adam的动量)卸载到CPU。优化器状态通常占用大量显存(对于Adam,每个参数需要额外8字节),卸载可以节省显存。 +actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \ + +# 释放推理引擎缓存 (Free Cache Engine) +# 作用: 在训练阶段释放推理引擎的KV cache和权重。这是3D-HybridEngine的核心优化,允许在同一GPU上交替进行推理和训练,显著降低显存需求。 +actor_rollout_ref.rollout.free_cache_engine=True \ + +# 熵计算优化 +# entropy_checkpointing: 在训练时对熵计算启用重计算,降低显存峰值 +# entropy_from_logits_with_chunking: 分块处理logits张量(如2048 tokens一组),避免一次性加载整个[bsz*seq_len, vocab]张量 +actor_rollout_ref.actor.entropy_checkpointing=True \ +actor_rollout_ref.ref.entropy_checkpointing=True \ +actor_rollout_ref.actor.entropy_from_logits_with_chunking=True \ +actor_rollout_ref.ref.entropy_from_logits_with_chunking=True \ + +# 推理引擎显存配置 +# gpu_memory_utilization: 控制vLLM使用的GPU显存比例(0.90 = 90%) +# enforce_eager=False: 启用CUDA graphs加速推理,但会占用额外显存 +actor_rollout_ref.rollout.gpu_memory_utilization=0.90 \ +actor_rollout_ref.rollout.enforce_eager=False \ +``` + +## NPU调优参考文章 + +环境变量相关：[环境变量列表-Ascend Extension for PyTorch6.0.0-昇腾社区](https://www.hiascend.com/document/detail/zh/Pytorch/600/apiref/Envvariables/Envir_001.html) + +社区性能调优教程：[性能调优流程-Ascend Extension for PyTorch6.0.0-昇腾社区](https://www.hiascend.com/document/detail/zh/Pytorch/600/ptmoddevg/trainingmigrguide/performance_tuning_0001.html) diff --git a/verl/docs/ascend_tutorial/model_support/model_and_algorithm_support.md b/verl/docs/ascend_tutorial/model_support/model_and_algorithm_support.md new file mode 100644 index 0000000000000000000000000000000000000000..dcce74a7f3484e34f13a2910d069720c85597bae --- /dev/null +++ b/verl/docs/ascend_tutorial/model_support/model_and_algorithm_support.md @@ -0,0 +1,60 @@ +# NPU Model & Algorithms Support Status + +Last updated: 05/14/2026. + +### Table 1 RL Algorithms + +| No. | model | algorithm | download Link | actor.strategy | rollout.name | shell location | hardware | +|---|---------------------|------------|------|------------|-----------|------------------------------|-----------------------------------| +| 1 | Qwen2.5-7B-Instruct | DAPO | [7B](https://huggingface.co/Qwen/Qwen2.5-7B-Instruct) | FSDP | vllm | [`link`](https://github.com/verl-project/verl-recipe/blob/main/dapo/run_dapo_qwen2.5_7b_npu.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 2 | Qwen2.5-VL-3B-Instruct | DAPO | [3B](https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct) | FSDP2 | vllm | [`link`](https://github.com/verl-project/verl-recipe/blob/main/dapo/run_dapo_qwen2.5_vl_3b_fsdp2_npu.sh) | Atlas 900 A2 PODc | +| 3 | Qwen2.5-VL-7B-Instruct | GRPO | [7B](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct) | FSDP | vllm | [`link`](https://github.com/verl-project/verl/blob/main/examples/grpo_trainer/run_qwen2_5_vl_7b_fsdp.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 4 | Qwen2.5-VL-7B-Instruct | GRPO | [7B](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct) | FSDP | sglang | [`link`](https://github.com/verl-project/verl/blob/main/examples/grpo_trainer/run_qwen2_5_vl_7b_fsdp.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 5 | Qwen2.5-VL-7B-Instruct | DAPO | [7B](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct) | FSDP2 | vllm | [`link`](https://github.com/verl-project/verl-recipe/blob/main/dapo/run_dapo_qwen2.5_vl_7b_fsdp2_npu.sh) | Atlas 900 A2 PODc | +| 6 | Qwen2.5-32B | GRPO | [32B](https://huggingface.co/Qwen/Qwen2.5-32B) | FSDP | vllm | [`link`](https://github.com/verl-project/verl/blob/main/examples/grpo_trainer/run_qwen2_5_32b_fsdp.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 7 | Qwen2.5-32B | DAPO | [32B](https://huggingface.co/Qwen/Qwen2.5-32B) | FSDP | vllm | [`link`](https://github.com/verl-project/verl-recipe/blob/main/dapo/run_dapo_qwen2.5_32b_npu.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 8 | Qwen2.5-32B | DAPO | [32B](https://huggingface.co/Qwen/Qwen2.5-32B) | FSDP2 | vllm | [`link`](https://github.com/verl-project/verl-recipe/blob/main/dapo/run_dapo_qwen2.5_32b_fsdp2_4k_npu.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 9 | Qwen2.5-32B | DAPO | [32B](https://huggingface.co/Qwen/Qwen2.5-32B) | FSDP2 | vllm | [`link`](https://github.com/verl-project/verl-recipe/blob/main/dapo/run_dapo_qwen2.5_32b_fsdp2_20k_npu.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 10 | Qwen2.5-VL-32B-Instruct | DAPO | [32B](https://huggingface.co/Qwen/Qwen2.5-VL-32B-Instruct) | FSDP2 | vllm | [`link`](https://github.com/verl-project/verl-recipe/blob/main/dapo/run_dapo_qwen2.5_vl_32b_fsdp2_npu.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 11 | Qwen3-4B | GRPO | [4B](https://huggingface.co/Qwen/Qwen3-4B) | FSDP | vllm | [`link`](https://github.com/verl-project/verl/blob/main/examples/grpo_trainer/run_qwen3_4b_fsdp.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 12 | Qwen3-8B | GRPO | [8B](https://huggingface.co/Qwen/Qwen3-8B) | FSDP | vllm | [`link`](https://github.com/verl-project/verl/blob/main/examples/grpo_trainer/run_qwen3_8b_fsdp.sh) | Atlas 900 A2 PODc | +| 13 | Qwen3-8B | GRPO | [8B](https://huggingface.co/Qwen/Qwen3-8B) | FSDP | sglang | [`link`](https://github.com/verl-project/verl/blob/main/examples/grpo_trainer/run_qwen3_8b_fsdp.sh) | Atlas 900 A2 PODc | +| 14 | Qwen3-8B | PPO | [8B](https://huggingface.co/Qwen/Qwen3-8B) | FSDP | vllm | [`link`](https://github.com/verl-project/verl/blob/main/examples/ppo_trainer/run_qwen3_8b_fsdp.sh) | Atlas 900 A2 PODc | +| 15 | Qwen3-8B | SAPO | [8B](https://huggingface.co/Qwen/Qwen3-8B) | FSDP | vllm | [`link`](https://github.com/verl-project/verl/blob/main/examples/sapo_trainer/run_qwen3_8b_fsdp.sh) | Atlas 900 A2 PODc, Atlas 200T A2 Box16, Atlas 800T A3 | +| 16 | Qwen3-8B | GSPO | [8B](https://huggingface.co/Qwen/Qwen3-8B) | FSDP2 | vllm | [`link`](https://github.com/verl-project/verl/blob/main/examples/gspo_trainer/run_qwen3_8b_fsdp.sh) | Atlas 900 A2 PODc, Atlas 200T A2 Box16, Atlas 800T A3 | +| 17 | Qwen3-8B | GRPO | [8B](https://huggingface.co/Qwen/Qwen3-8B) | MindSpeed_LLM | sglang | [`link`](https://github.com/verl-project/verl/blob/main/examples/grpo_trainer/run_qwen3_8b_mindspeed.sh) | Atlas 900 A2 PODc, Atlas 200T A2 Box16, Atlas 800T A3 | +| 18 | Qwen3-8B | GRPO(One_Step_Off_Policy) | [8B](https://huggingface.co/Qwen/Qwen3-8B) | FSDP2 | vllm | [`link`](https://github.com/verl-project/verl/blob/main/verl/experimental/one_step_off_policy/shell/grpo_qwen3_8b_gsm8k_fsdp2_8_8_npu.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 19 | Qwen3-8B-Base | DAPO | [8B](https://huggingface.co/Qwen/Qwen3-8B-Base) | FSDP | vllm | [`link`](https://github.com/verl-project/verl-recipe/blob/main/dapo/run_dapo_qwen3_8b_base_npu.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 20 | Qwen3-14B-Base | DAPO | [14B](https://huggingface.co/Qwen/Qwen3-14B-Base) | FSDP | vllm | [`link`](https://github.com/verl-project/verl-recipe/blob/main/dapo/run_dapo_qwen3_14b_base_npu.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 21 | Qwen3-30B | DAPO | [30B](https://huggingface.co/Qwen/Qwen3-30B) | FSDP | vllm | [`link`](https://github.com/verl-project/verl-recipe/blob/main/dapo/run_dapo_qwen3_30b_fsdp_6k_npu.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 22 | Qwen3-30B-A3B-Base | DAPO | [30B](https://huggingface.co/Qwen/Qwen3-30B-A3B-Base) | FSDP | vllm | [`link`](https://github.com/verl-project/verl-recipe/blob/main/dapo/run_dapo_qwen3_moe_30b_base_fsdp_npu.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 23 | Qwen3-30B-A3B | GRPO | [30B](https://huggingface.co/Qwen/Qwen3-30B-A3B) | Megatron | sglang | [`link`](https://github.com/verl-project/verl/blob/main/examples/ascend_extras/grpo_trainer/run_qwen3_30b_a3b_megatron.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 24 | Qwen3-30B-A3B | GRPO | [30B](https://huggingface.co/Qwen/Qwen3-30B-A3B) | MindSpeed_LLM | sglang | [`link`](https://github.com/verl-project/verl/blob/main/examples/ascend_extras/grpo_trainer/run_qwen3_30b_a3b_mindspeed.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 25 | Qwen3-30B-A3B | DAPO | [30B](https://huggingface.co/Qwen/Qwen3-30B-A3B) | Megatron | vllm | [`link`](https://github.com/verl-project/verl-recipe/blob/main/dapo/run_dapo_qwen3_moe_30b_megatron_npu.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 26 | Qwen3-30B-A3B | GRPO (fully_async) | [30B](https://huggingface.co/Qwen/Qwen3-30B-A3B) | FSDP | vllm | [`link`](https://github.com/verl-project/verl/blob/main/verl/experimental/fully_async_policy/shell/dapo_30b_a3b_math_fsdp_npu.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 27 | Qwen3-32B | GRPO | [32B](https://huggingface.co/Qwen/Qwen3-32B) | FSDP | vllm | [`link`](https://github.com/verl-project/verl/blob/main/examples/ascend_extras/grpo_trainer/run_qwen3_32b_fsdp.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 28 | Qwen3-32B | GRPO | [32B](https://huggingface.co/Qwen/Qwen3-32B) | MindSpeed_LLM | sglang | [`link`](https://github.com/verl-project/verl/blob/main/examples/ascend_extras/grpo_trainer/run_qwen3_32b_mindspeed.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 29 | Qwen3-235B-A22B | GRPO | [235B](https://huggingface.co/Qwen/Qwen3-235B-A22B) | Megatron | vllm | [`link`](https://github.com/verl-project/verl/blob/main/examples/grpo_trainer/run_qwen3_235b_a22b_megatron.sh) | Atlas 900 A2 PODc, Atlas 200T A2 Box16, Atlas 800T A3 | +| 30 | Qwen3-235B-A22B | GRPO (fully_async) | [235B](https://huggingface.co/Qwen/Qwen3-235B-A22B) | Megatron | vllm | [`link`](https://github.com/verl-project/verl/blob/main/verl/experimental/fully_async_policy/shell/grpo_qwen3_235b_megatron_npu.sh%E2%80%8E) | Atlas 900 A2 PODc, Atlas 200T A2 Box16, Atlas 800T A3 | +| 31 | Qwen3-235B-A22B-Instruct | GRPO | [235B](https://huggingface.co/Qwen/Qwen3-235B-A22B-Instruct) | Megatron | vllm | [`link`](https://github.com/verl-project/verl/blob/main/examples/ascend_extras/grpo_trainer/run_qwen3_235b_256k_megatron.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 32 | Qwen3-VL-8B-Instruct | GRPO | [8B](https://huggingface.co/Qwen/Qwen3-VL-8B-Instruct) | FSDP2 | vllm | [`link`](https://github.com/verl-project/verl/blob/main/examples/grpo_trainer/run_qwen3_vl_8b_fsdp.sh) | Atlas 900 A2 PODc, Atlas 200T A2 Box16, Atlas 800T A3 | +| 33 | Qwen3-VL-8B-Instruct | GRPO (fully_async) | [8B](https://huggingface.co/Qwen/Qwen3-VL-8B-Instruct) | FSDP2 | vllm | [`link`](https://github.com/verl-project/verl/blob/main/verl/experimental/fully_async_policy/shell/geo3k_qwen3vl_8b_fsdp2_16_16_npu.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 34 | Qwen3-VL-30B-A3B-Instruct | GRPO | [30B](https://huggingface.co/Qwen/Qwen3-VL-30B-A3B-Instruct) | FSDP2 | vllm | [`link`](https://github.com/verl-project/verl/blob/main/examples/grpo_trainer/run_qwen3_vl_30b_a3b_fsdp.sh) | Atlas 900 A2 PODc, Atlas 200T A2 Box16, Atlas 800T A3 | +| 35 | Qwen3-VL-30B-A3B-Instruct | DAPO | [30B](https://huggingface.co/Qwen/Qwen3-VL-30B-A3B-Instruct) | FSDP2 | vllm | [`link`](https://github.com/verl-project/verl-recipe/blob/main/dapo/run_dapo_qwen3_vl_30b_fsdp2_npu.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 36 | Qwen3-VL-30B-A3B-Instruct | GRPO (fully_async) | [30B](https://huggingface.co/Qwen/Qwen3-VL-30B-A3B-Instruct) | Megatron | vllm | [`link`](https://github.com/verl-project/verl/blob/main/verl/experimental/fully_async_policy/shell/geo3k_qwen3vl_30b_megatron_6_2_npu_async.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 37 | Qwen3.5-27B | GRPO | [27B](https://huggingface.co/Qwen/Qwen3.5-27B) | FSDP2 | vllm | [`link`](https://github.com/verl-project/verl/blob/main/examples/grpo_trainer/run_qwen3_5_27b_fsdp.sh) | Atlas 900 A2 PODc, Atlas 200T A2 Box16, Atlas 800T A3 | +| 38 | Qwen3.5-27B | GRPO | [27B](https://huggingface.co/Qwen/Qwen3.5-27B) | FSDP(MindSpeed-MM) | vllm | [`link`](https://github.com/verl-project/verl-recipe/blob/main/grpo_mindspeed_mm/run_qwen3_5-27b_npu.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 39 | Qwen3.5-35B-A3B | GRPO | [35B](https://huggingface.co/Qwen/Qwen3.5-35B-A3B) | FSDP2 | vllm | [`link`](https://github.com/verl-project/verl/blob/main/examples/grpo_trainer/run_qwen3_5_35b_fsdp.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 40 | Qwen3.5-35B-A3B | GRPO | [35B](https://huggingface.co/Qwen/Qwen3.5-35B-A3B) | FSDP2(MindSpeed-MM) | vllm | [`link`](https://github.com/verl-project/verl-recipe/blob/main/grpo_mindspeed_mm/run_qwen3_5-35b_npu.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 41 | Qwen3-Next-80B-A3B-Instruct | GRPO | [80B](https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct) | FSDP2 | vllm | [`link`](https://github.com/verl-project/verl/blob/main/examples/grpo_trainer/run_qwen3_next_80b_a3b_fsdp.sh) | Atlas 900 A2 PODc, Atlas 200T A2 Box16, Atlas 800T A3 | +| 42 | DeepSeek-V3 | GRPO | [671B](https://huggingface.co/deepseek-ai/DeepSeek-V3) | Megatron | vllm | [`link`](https://github.com/verl-project/verl-recipe/blob/main/r1_ascend/run_deepseekv3_671b_grpo_megatron_npu.sh) | Atlas 200T A2 Box16, Atlas 800T A3 | +| 43 | Qwen3-1.7B | GRPO | [1.7B](https://huggingface.co/Qwen/Qwen3-1.7B) | VeOmni | vllm | [`link`](https://github.com/verl-project/verl-ascend-recipe/blob/main/verl_ascend_practice/run_qwen3_1_7b_npu.sh) | Atlas 800T A3 | +| 44 | Qwen3-30B-A3B | GRPO | [30B](https://huggingface.co/Qwen/Qwen3-30B-A3B) | VeOmni | vllm | [`link`](https://github.com/verl-project/verl/blob/main/examples/grpo_trainer/run_qwen3_30b_veomni.sh) | Atlas 800T A3 | +| 45 | Qwen3-VL-30B-A3B-Instruct | GRPO | [30B](https://huggingface.co/Qwen/Qwen3-VL-30B-A3B-Instruct) | VeOmni | vllm | [`link`](https://github.com/verl-project/verl/blob/main/examples/grpo_trainer/run_qwen3_vl_30b_moe_veomni.sh) | Atlas 800T A3 | + +### Table 2 SFT Algorithms + +| No. | model | algorithm | download link | actor.strategy | shell location | hardware | +|-----|-------|-----------|---------------|----------------|----------------|----------| +| 1 | Qwen3-8B | SFT-PEFT | [8B](https://huggingface.co/Qwen/Qwen3-8B) | FSDP | [`link`](https://github.com/verl-project/verl/blob/main/examples/sft/gsm8k/run_qwen3_8b_fsdp.sh) | Atlas 900 A2 PODc | +| 2 | Qwen2.5-7B-Instruct | ReTool-SFT | [7B](https://huggingface.co/Qwen/Qwen2.5-7B-Instruct) | FSDP | [`link`](https://github.com/verl-project/verl-recipe/blob/main/retool/run_qwen2_7b_sft_npu.sh) | Atlas 900 A2 PODc | diff --git a/verl/docs/blog/v0.7.md b/verl/docs/blog/v0.7.md new file mode 100644 index 0000000000000000000000000000000000000000..46be13b94ecc781c52dbf970cfa01af89a579a05 --- /dev/null +++ b/verl/docs/blog/v0.7.md @@ -0,0 +1,274 @@ +# verl 0.7 release blog + +**Author:** verl team + +Last updated: 01/03/2026. + +## Overview +verl adopts a Hybrid-Controller architecture (also known as HybridFlow). Sharing design principles with asynchronous sharded dataflow systems like Google Pathways, verl models Reinforcement Learning (RL) algorithms, such as PPO, GRPO, DAPO, and others, as a multi-stage, multi-model and parallelizable dataflow graph. + +To balance flexibility with performance, verl unifies two distinct programming models: + +**High-Level Single-Controller (MPMD)**: At the orchestration level, a single process `RLTrainer` manages the global computation graph. It handles macro-tasks such as scheduling rollout generation, triggering reward scoring, and dispatching distributed training jobs. + +**Internal Multi-Controller (SPMD)**: Internally, the Model Engine operates in standard distributed training mode. Workers execute identical programs, via trainer backends like FSDP, Megatron, or VeOmni, or rollout executors (not rollout server) like vLLM/SGLang/TensorRT-LLM, to perform heavy distributed computation, synchronizing via collective communication. + +
+ hybridflow.png +
+ +This hybrid approach offers significant advantages: + +**Flexible Orchestration**: The single-controller design allows verl to dynamically manage complex constraints within the computation graph, including flexible data dependencies, diverse resource allocation and model placement, and fine-grained asynchronous staleness control. + +**Abstraction of Complexity**: We encapsulate complex parallel strategies—such as 5D parallelism (DP, TP, CP, PP, and EP)—strictly within the Model Engine. This allows users to focus entirely on RL algorithm implementation without getting bogged down by the details of distributed training. + +Furthermore, leveraging Ray placement groups, verl provides `ResourcePool` and `WorkerGroup` abstractions. These enable flexible GPU sharing among the various roles in the RL process—such as actor, critic, reward, and rollout—allowing components to share resources efficiently while remaining isolated. + +As illustrated in the diagram below, the overall architecture of verl is divided into two layers: + +- **verl-core**: provides four components required for the RL pipeline: model engine, rollout engine, checkpoint engine, and transfer queue. Each component exposes abstract interfaces, making them both extensible and pluggable. +- **verl-trainer**: builds upon these components, construct various RL pipelines—such as on-policy, one-step-off-policy, and fully asynchronous—tailored to meet the demands of diverse scenarios. + +
+ verl-arch.png +
+ + +## verl-core +### Model Engine + +The Model Engine serves as verl's core training engine, defining a set of abstract interfaces that support pluggable backends. It operates in SPMD mode: +- SFT: Workers are launched via torchrun. +- RL: Workers are executed via the WorkerGroup API, invoked by the single-controller. + +The abstract interfaces include methods like `initialize`, `forward`, `optimizer_step`, and `load`/`offload`. Integrating a new training engine simply requires inheriting and implementing these interfaces. Crucially, because all backends adhere to this unified abstraction, adding a new Model Engine requires absolutely no code modification on the caller side. The RLTrainer remains completely agnostic to the backend's specific parallel strategy when calling these interfaces, while the WorkerGroup automatically handles data dispatch and collection based on the underlying parallelism. + +Currently, the Model Engine supports the following backends (more backend maybe supported in future, e.g torchtitan): +|Backend|Parallelism|Performance|Support Model|New Model Support Time +|-----|-----|----|----|----| +|FSDP| FSDP+SP|Dense medium/MoE low| all transformer models|Day 0 +|MCore| DP+TP+PP+EP+CP|High| see [Megatron-Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge) support model list|few weeks or month +|VeOmni| FSDP+SP+EP|Medium| see [VeOmni](https://github.com/ByteDance-Seed/VeOmni) support model list|~1 week + +```python +class BaseEngine: + def initialize(self): + """Instantiate or load the model, optimizer, and learning rate scheduler.""" + raise NotImplementedError + + def optimizer_zero_grad(self): + """Zero the gradients of the optimizer.""" + raise NotImplementedError + + def optimizer_step(self): + """Perform an optimization step using the optimizer.""" + raise NotImplementedError + + def lr_scheduler_step(self): + """Advance the learning rate scheduler by one step.""" + raise NotImplementedError + + def forward_backward_batch(self, data: TensorDict, loss_function: Callable, forward_only=False) -> Any: + """Perform a forward pass and optionally a backward pass on a batch of data.""" + raise NotImplementedError + + def get_per_tensor_param(self) -> tuple[Generator[tuple[str, torch.Tensor], None, None], Optional[dict]]: + """Get a generator that yields per-tensor parameters and optional peft config.""" + raise NotImplementedError + + def to(self, device: str, model: bool = True, optimizer: bool = True, grad: bool = True): + """Move model parameters, optimizer states, or both to the specified device.""" + raise NotImplementedError +``` + + +### Rollout Engine +As LLM reinforcement learning evolves from single-turn, static tasks to multi-turn, dynamic, and interactive agentic tasks, the legacy SPMD rollout mode previously used by verl has become insufficient. Consequently, in verl v0.7, we have removed the SPMD rollout mode and switched to rollout server mode by default. + +
+ rollout_engine.png +
+ +In the server mode, the LLM server operates as online serving rather than the traditional offline batch inference. Clients send per-sample requests to the server, enabling the engine to utilize dynamic batching. This significantly enhances throughput efficiency for multi-turn conversation. Furthermore, the server-based approach eliminates the need for intrusive modifications to the LLM inference engine, allowing for the seamless integration of modern inference backends such as vLLM, SGLang, and TensorRT-LLM. + +On the client side, verl introduces an extensible **AgentLoop** abstraction designed to define custom agentic task loops. This abstraction manages the cycle of requesting responses from the LLM server and interacting with external environments to obtain feedback. We provide two default implementations: +- **SingleTurnAgentLoop**: Designed for standard single-turn tasks. +- **ToolAgentLoop**: Designed for classic ReAct architectures involving multi-turn tool invocation. + +Users can implement custom AgentLoop logic tailored to their specific needs, such as [SWEAgentLoop](https://github.com/verl-project/verl/pull/4080) or GUIAgentLoop. + +```python +class AgentLoopBase(ABC): + @abstractmethod + async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentLoopOutput: + """Run agent loop to interact with LLM server and environment. + + Args: + sampling_params (Dict[str, Any]): LLM sampling params. + **kwargs: dataset fields from `verl.utils.dataset.RLHFDataset`. + + Returns: + AgentLoopOutput: Agent loop output. + """ + raise NotImplementedError +``` + +### TransferQueue +As mentioned, verl uses a global single-controller RLTrainer to orchestrate the computation graph. A major limitation in the current implementation is that the RLTrainer handles both control and data flow, creating a bottleneck when dispatching data between components. This issue is amplified by the massive data volumes in multimodal training (images, video, audio) and complex algorithms like router replay, which requires transmitting large tensors per sample. Our earlier attempt to solve this using the Ray object store yielded poor performance due to the lack of tensor optimization and fine-grained column access. + +
+ transfer_queue.png +
+ +In v0.7, we experimentally introduced **TransferQueue** to decouple control flow from data flow. The RLTrainer now only dispatch instructions and metadata, while TransferQueue handles data transmission via reference passing. TransferQueue is specifically optimized for PyTorch tensors (supporting zero-copy and RDMA) and allows for backend extensions like ZeroMQ, NIXL, and Ray RDT. We plan to make this the default transmission method in v0.8. + +```python +# In PPOTrainer +def fit(self): + batch = next(dataloader) + gen_batch: BatchMeta = self.rollout_manager.generate_sequences(batch) + output: BatchMeta = self.actor_rollout_wg.compute_log_prob(gen_batch) + gen_batch = gen_batch.union(output) + output = self.actor_rollout_wg.update_actor(gen_batch) + +# In Worker +def compute_log_prob(self, batch: BatchMeta) -> BatchMeta: + data = tq.get(batch) + output = self.actor.infer_batch(data=data) + return tq.put(output) +``` + +### Checkpoint Engine + +With the increase in LLM context lengths and the evolution of agentic tasks, the "long-tail" problem in rollout has become prominent, limiting the overall efficiency of RL training. + +To mitigate this, a viable strategy is moving from on-policy synchronous training to off-policy asynchronous training, e.g [Laminar](https://arxiv.org/abs/2510.12633), [Areal](https://arxiv.org/abs/2505.24298), [StreamRL](https://arxiv.org/abs/2504.15930), [LlamaRL](https://arxiv.org/pdf/2505.24034), [PipelineRL](https://arxiv.org/abs/2509.19128). This involves separating the rollout and model engines onto different nodes (a disaggregated architecture, as opposed to colocated), with data transmitted via queues. This separation alleviates the rollout long-tail issue and enables rollout elastic scaling, fault tolerance, and heterogeneous hardware. However, it introduces a new challenge: efficient cross-node parameter synchronization. + +
+ checkpoint_engine.png +
+ +To address this, we introduce the Checkpoint Engine: a unified abstraction layer designed to synchronize weights between various training and inference backends. +- It provides three unified APIs to implement the streaming transmission of parameters. +- Users can extend the Transport Layer implementation based on their specific infrastructure requirements (device, network, local cache, etc.). + +Currently, we provide two transport backends: NCCL (for broadcast collective communication) and NIXL (for P2P point-to-point communication). + +```python +class CheckpointEngine(ABC): + @abstractmethod + async def send_weights(self, weights: Generator[tuple[str, torch.Tensor], None, None]): + """Send the weights of the model. + + Args: + weights: A generator that yields the name of the weight tensor and the tensor itself. + """ + raise NotImplementedError + + @abstractmethod + async def receive_weights(self) -> Generator[tuple[str, torch.Tensor], None, None]: + """Receive the weights of the model. + + Yields: + A tuple of the name of the weight tensor and the tensor itself. + """ + raise NotImplementedError +``` + +## verl-trainer +Building upon the four core components provided by verl-core, verl-trainer constructs several RL training pipelines tailored to specific scenarios. These pipelines are designed to address training efficiency challenges across varying scales and requirements: + +**On-policy (Synchronous)** + - Main Features: Executes rollout and training serially, typically sharing GPU resources (Colocate). It strictly adheres to standard on-policy algorithm definitions, where training must wait for all samples to be generated. + - Scenarios: Best for baseline implementations, scenarios where strict algorithmic correctness is prioritized over training throughput. + +**One-step-off-policy (Async)** + - Main Features: Parallelizes generation and training by overlapping the current training step with the next batch's generation. It employs resource isolation and uses parameters from the previous step for rollout to minimize GPU idle time. + - Scenarios: Ideal for scenarios requiring moderate efficiency gains (20%–40%) while maintaining training stability very close to strict on-policy methods. + +**Fully async (Decoupled & Streaming)** + - Main Features: Completely decouples the Trainer and Rollouter onto separate nodes. It utilizes streaming data transfer, staleness control, and partial rollout mechanisms to maximize throughput and mitigate long-tail generation latency. + - Scenarios: Essential for large-scale training (e.g., 128+ GPUs) or complex reasoning tasks (e.g., long chain-of-thought) where generation latency significantly bottlenecks performance. + +
+ fully_async.png +
+ +## roadmap +### v0.7 release + +**Model Engine** +- Integrate Megatron-Bridge and support LoRA/PEFT, see blog post: [How We Build Trillion Parameter Reasoning RL with 10% GPUs](https://macaron.im/mindlab/research/building-trillion-parameter-reasoning-rl-with-10-gpus) +- Support experimental fp8 training for megatron backend +- Support new model for megatron backend: GPT-OSS, Qwen3-Next +- Comprehensive support for new mode engine, FSDP and Megatron engine are production ready. + - Dispatch tensordict with nested tensor instead of padded DataProto + - Add TrainingWorker that resembles Tinker-like API + - Add VLM support for model engine, SFT and RL trainer + - Add model engine based critic model + - Implement ActorRolloutRefWorker by TrainingWorker, support different backend in one worker +- New VeOmni engine added, still in alpha status. + +**Rollout Engine** +- Remove SPMD rollout mode +- Support blockwise fp8 rollout for vllm and sglang; support online quant for vllm with torchao +- Experimental router replay support for vllm +- Optimize multi-modal data fetch and preprocess, support video input +- Upgrade to vllm==0.12.0; sglang==0.5.6 + +**Reward** +- Support hybrid reward scenarios, including generative, discriminative, rule-based rewards, and their combinations. +- Refactor reward models into server mode, supporting both colocated and standalone deployments. +- Introduce new reward managers to handle more complex scenarios, limited mode for request rate control and remote mode for CPU-intensive tasks. + +**Algorithm** +- Add [CISPO](https://arxiv.org/pdf/2506.13585): Clipped IS-weight Policy Optimization +- Add [SAPO](https://arxiv.org/abs/2511.20347): Soft Adaptive Policy Optimization + +**Recipe** +- [NEW] VLA: add experimental support for VLA model +- [NEW] [rhymerl](https://arxiv.org/abs/2508.18588): History Rhymes: Accelerating LLM Reinforcement Learning with RhymeRL +- TransferQueue: support multiple data partition and optimize tensor zero-copy serialization +- One-step-off-policy/Fully async: optimize weight synchronization by checkpoint engine with bucket and pipeline support. + +### v0.8 + +**Model Engine** +- Deprecate DataProto by Tensordict for zero padding transmission +- Switch default to new model engine, mark legacy engine (fsdp_workers.py, megatron_workers.py) as deprecated +- Feature parity between new and legacy model engine: LoRA/PEFT, etc +- Polish VeOmni engine to production ready status +- Support MTP RL training +- Optimize GPU memory for long context: fine-grained activation recompuation/offload +- New model support: DeepSeek V3.2, etc + +**Rollout Engine** +- New rollout engine TensorRT-LLM +- Separate vllm worker from trainer process, update weights by cuda ipc + +**TransferQueue** +- Merge TransferQueue recipe into main +- Optimize e2e image/video vlm training pipeline by TransferQueue +- Optimize router replay transmission by TransferQueue + +**Checkpoint Engine** +- Add checkpoint engine abstract interface +- Add NCCL and NIXL transport backend +- Add more transport backend + +### v0.9 + +**Trainer** +- Merge Full async into main: refactor with verl-core component + +**Model Engine** +- Remove legacy model engine (fsdp_workers.py, megatron_workers.py) +- Support omni-model RL training: Qwen3-Omni, BAGEL, etc + +**Rollout Engine** +- New rollout engine vllm-omni + +**More agentic training recipe** +- SWEAgent +- GUIAgent diff --git a/verl/docs/conf.py b/verl/docs/conf.py new file mode 100644 index 0000000000000000000000000000000000000000..cbeabbd81b28e97fe0d0e8bcf436ab92f5833743 --- /dev/null +++ b/verl/docs/conf.py @@ -0,0 +1,113 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Configuration file for the Sphinx documentation builder. +# +# This file only contains a selection of the most common options. For a full +# list see the documentation: +# https://www.sphinx-doc.org/en/master/usage/configuration.html + +# -- Path setup -------------------------------------------------------------- + +# If extensions (or modules to document with autodoc) are in another directory, +# add these directories to sys.path here. If the directory is relative to the +# documentation root, use os.path.abspath to make it absolute, like shown here. +# +# import os +# import sys +# sys.path.insert(0, os.path.abspath('.')) + + +# -- Project information ----------------------------------------------------- + +project = "verl" +copyright = "2024 ByteDance Seed Foundation MLSys Team" +author = "Guangming Sheng, Chi Zhang, Yanghua Peng, Haibin Lin" + + +# -- General configuration --------------------------------------------------- +# The master toctree document. +master_doc = "index" + +# Add any Sphinx extension module names here, as strings. They can be +# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom +# ones. +extensions = [ + "myst_parser", + "sphinx.ext.autodoc", + "sphinx.ext.autosummary", + "sphinx.ext.autosectionlabel", + "sphinx.ext.napoleon", + "sphinx.ext.viewcode", +] + +# MyST-Parser settings +myst_enable_extensions = [ + "dollarmath", # Enables $...$ and $$...$$ syntax + "amsmath", # Enables amsmath environments +] + +# Use Google style docstrings instead of NumPy docstrings. +napoleon_google_docstring = True +napoleon_numpy_docstring = False + +# The suffix(es) of source filenames. +# You can specify multiple suffix as a list of string: +source_suffix = { + ".rst": "restructuredtext", + ".md": "markdown", +} + +# Add any paths that contain templates here, relative to this directory. +templates_path = ["_templates"] + +# The language for content autogenerated by Sphinx. Refer to documentation +# for a list of supported languages. +# +# This is also used if you do content translation via gettext catalogs. +# Usually you set "language" from the command line for these cases. +language = "en" + +# List of patterns, relative to source directory, that match files and +# directories to ignore when looking for source files. +# This pattern also affects html_static_path and html_extra_path. +exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] + + +# -- Options for HTML output ------------------------------------------------- + +# The theme to use for HTML and HTML Help pages. See the documentation for +# a list of builtin themes. +# +html_theme = "sphinx_rtd_theme" + +# Add any paths that contain custom static files (such as style sheets) here, +# relative to this directory. They are copied after the builtin static files, +# so a file named "default.css" will overwrite the builtin "default.css". +html_static_path = ["_static"] + +# Add the JavaScript file +html_js_files = [ + "js/runllm-widget.js", + "js/resizable-sidebar.js", +] + +# Add custom CSS file for full-width layout +html_css_files = [ + "custom.css", +] + +exclude_patterns += ["README.md", "README_vllm0.7.md"] + +suppress_warnings = ["ref.duplicate", "ref.myst"] diff --git a/verl/docs/contributing/editing-agent-instructions.md b/verl/docs/contributing/editing-agent-instructions.md new file mode 100644 index 0000000000000000000000000000000000000000..dc3657b0c347c52743277d6cf2993cdc03bdec8f --- /dev/null +++ b/verl/docs/contributing/editing-agent-instructions.md @@ -0,0 +1,95 @@ +# Editing Agent Instructions + +> Read this before modifying `AGENTS.md` or any guide it links to. + +## File Layout + +| Generic | Variants | Audience | Scope | +| ----------- | -------------------------- | --------------- | ------------------------------ | +| `AGENTS.md` | `CLAUDE.md`, ... | Agents | Project-wide instructions. | +| `.agent/` | `.claude/`, `.codex/`, ... | Agents | Agent specification directory. | +| `docs/` | N/A | Humans & Agents | Project documentation. | + +Generic files are framework-agnostic. Variants directly symlink generic files or refer to them and add framework-specific content. + +For skills, keep canonical content under `.agent/skills//SKILL.md`. +Variant trees such as `.codex/skills/` and `.claude/skills/` should be directory symlinks to the canonical `.agent` skill directory, not directories containing only a symlinked `SKILL.md`. +Codex documents support for symlinked skill folders and can skip file-level `SKILL.md` symlinks during discovery. Claude Code discovers the file-level symlink layout in current versions, but directory symlinks match the shared skill-package structure and keep supporting files in sync across frameworks. + +## Token Budget Mindset + +`AGENTS.md` loads on every agent request; domain guides load on entry to a relevant area. +Keep `AGENTS.md` under **200 lines** and each domain guide under **300 lines**. +When a file exceeds its budget, split or prune — do not compress prose to fit. + +## When NOT to Add Content + +Before writing a new rule, ask whether it is actually needed: + +- **Agents already do it.** Test with a prompt first. If the agent behaves correctly without the rule, don't add it. +- **One-off incident.** Prefer a code-level fix (lint rule, CI check, test assertion) over a new doc rule. +- **Hardcoded paths.** File paths change; use "search for X" patterns instead. +- **Upstream docs.** Don't reproduce pytest, ruff, or other tool docs — link to them. +- **Contradicts an existing rule.** Search all linked guides before adding. If two rules conflict, consolidate into one. +- **Already covered elsewhere.** Search `AGENTS.md` and every linked guide for overlapping guidance. + +If any of the above apply, **do not add the content**. + +## Where Content Belongs + +The goal is a lean `AGENTS.md` plus rich domain guides that teach agents what they can't learn from the code alone. + +| Scope | File | +| ------------------------------------------------------------------------------------------------ | ------------ | +| Project-wide invariants (contribution policy, env setup, test/lint commands, commit conventions) | `AGENTS.md` | +| Area-specific knowledge (model patterns, format details, deprecation timelines) | Domain guide | + +**Rules of thumb:** + +- If it only matters for one area, put it in a domain guide. +- If it matters for all areas, consider `AGENTS.md` — but first verify agents don't already do it. +- Create a new domain guide when you have 5 or more non-obvious instructions sharing a coherent scope. + +## What Makes a Good Domain Guide + +Add what agents can't infer from the code or public docs: project-specific +conventions that differ from standard patterns, correct approaches that require +cross-file context, and fixes for repeated mistakes. +Each entry should be short, specific, and actionable — e.g., which files to +touch, what order to change them in, and which tests to run. + +## Keeping Docs Lean + +- Every addition should trigger review of surrounding content for stale or redundant items. +- Refer to existing files (e.g., "follow the PR template") instead of restating their content — keep a single source of truth. +- Prefer examples over explanations — a 3-line snippet beats a paragraph of prose. +- Merge related bullets into one principle instead of listing variants. +- Use `search for X` instead of hardcoded file paths. +- PR references are fine in domain guides for traceability, but avoid them in `AGENTS.md`. + +## Anti-Patterns + +| Pattern | Problem | +| ------------------------- | ----------------------------------------------------------------------- | +| Reactive accumulation | Adding a rule per incident without pruning leads to bloat | +| Copy-paste between guides | Duplicated content drifts apart; keep in one place, link from the other | +| Imperative walls | Long DO NOT lists that agents skim past; consolidate into principles | +| Config snapshots | Show the command to get the value, not the value itself | + +## Change Checklist + +Before submitting changes to any agent instruction file: + +- [ ] **Non-obvious?** Would an agent do the wrong thing without this rule? +- [ ] **No conflicts?** Searched all linked guides for contradictions? +- [ ] **Right file?** Project-wide goes in `AGENTS.md`, area-specific in a domain guide? +- [ ] **Offset the addition?** Removed or consolidated something to compensate? +- [ ] **Under budget?** `AGENTS.md` < 200 lines, domain guides < 300 lines? +- [ ] **No hardcoded paths?** Uses "search for X" where paths may change? +- [ ] **Tested?** Verified that an agent actually follows the new instruction? + +## Acknowledgements + +This guide is adapted from the [vLLM project](https://github.com/vllm-project/vllm)'s [`editing-agent-instructions.md`](https://github.com/vllm-project/vllm/blob/main/docs/contributing/editing-agent-instructions.md). + +Last updated: 05/13/2026 diff --git a/verl/docs/data/transfer_queue.md b/verl/docs/data/transfer_queue.md new file mode 100644 index 0000000000000000000000000000000000000000..aeb31e0d5f93bd3c64ca33ec028ede7f4713ceeb --- /dev/null +++ b/verl/docs/data/transfer_queue.md @@ -0,0 +1,341 @@ +# TransferQueue Data System + +Last updated: 05/15/2026. + +This doc introduce [TransferQueue](https://github.com/Ascend/TransferQueue), an asynchronous streaming data management system for efficient post-training. + +🔥 **Now TransferQueue is open-sourced at both [Github](https://github.com/Ascend/TransferQueue) and [GitCode](https://gitcode.com/Ascend/TransferQueue). You are welcome to submit contributions or propose new ideas on either platform!** + + +> At the mean time, the early development history remains accessible at: https://github.com/TransferQueue/TransferQueue. + +

Overview

+ +TransferQueue is a high-performance data storage and transfer module with panoramic data visibility and streaming scheduling capabilities, optimized for efficient dataflow in post-training workflows. + +

+ +

+ +TransferQueue offers **fine-grained, sub-sample-level** data management and **load-balancing** capabilities. It serves as a data gateway that decouples explicit data dependencies across computational tasks, enabling a divide-and-conquer approach that significantly simplifies algorithm controller design. + +

+ +

+ +

Updates

+ + - **April 15, 2026**: 🔥 TransferQueue has been adopted in [Relax](https://github.com/redai-infra/Relax)! By leveraging the `StreamingDataLoader` abstraction, it schedules training data across the cluster at micro-batch granularity, reducing synchronization barriers in a single-controller setup. + - **April 10, 2026**: 🔥 TransferQueue is now officially integrated into [verl](https://github.com/verl-project/verl/pull/5401)! **We achieved an end-to-end performance gain of 49.1% for multi-modal post-training on a 128 × H100 GPU cluster!** Refer to [our blog](https://www.yuque.com/haomingzi-lfse7/lhp4el/gm8mkpfu83luuhxg?singleDoc#) for more details. + - **Feb 8, 2026**: 🔥 Initialization and usage are greatly simplified by high-level APIs [PR#26](https://github.com/Ascend/TransferQueue/pull/26), [PR#28](https://github.com/Ascend/TransferQueue/pull/28). You can now use a Redis-style API to take advantage of most of the advanced features provided by TransferQueue! + - **Jan 28, 2026**: We experimentally introduce the `StreamingDataLoader` interface for a fully-streamed production-consumption pipeline. Refer to our [tutorials/06_streaming_dataloader.py](https://github.com/Ascend/TransferQueue/blob/main/tutorial/06_streaming_dataloader.py) for details. + - **Dec 30, 2025**: **TransferQueue x verl** integration has been tested with the DAPO algorithm at scale **(64 nodes, 1024 cards)**. It significantly optimizes host memory utilization and accelerates data transfers. Stay tuned for more details! + - **Dec 20, 2025**: 🔥 The official [tutorial](https://github.com/Ascend/TransferQueue/tree/main/tutorial) is released! Feel free to check it out. + - **Nov 10, 2025**: We disentangled the data retrieval logic from TransferQueueController [PR#101](https://github.com/TransferQueue/TransferQueue/pull/101). Now you can implement your own `Sampler` to customize data consumption. + - **Nov 5, 2025**: We provide a `KVStorageManager` that simplifies the integration with KV-based storage backends [PR#96](https://github.com/TransferQueue/TransferQueue/pull/96). The first available KV-based backend is [openYuanrong](https://gitcode.com/openeuler/yuanrong-datasystem). + - **Nov 4, 2025**: Data partitioning capability is available in [PR#98](https://github.com/TransferQueue/TransferQueue/pull/98). Now you can define logical data partitions to manage your train/val/test datasets. + - **Oct 25, 2025**: Storage backends are now pluggable in [PR#66](https://github.com/TransferQueue/TransferQueue/pull/66). You can try to integrate your own storage backend with TransferQueue now! + - **Oct 21, 2025**: Early integration with verl is ready [verl/pull/3649](https://github.com/volcengine/verl/pull/3649). Following PRs will optimize the single controller architecture by fully decoupling data & control flows. + - **July 22, 2025**: We published a series of Chinese blog posts on Zhihu 1, 2. + - **July 21, 2025**: We initiated an RFC in the verl community [verl/RFC#2662](https://github.com/volcengine/verl/discussions/2662). + - **July 2, 2025**: We published the paper [AsyncFlow](https://arxiv.org/abs/2507.01663). + +

Components

+ +### Control Plane: Panoramic Data Management + +In the control plane, `TransferQueueController` tracks the **production status** and **consumption status** of each training sample as metadata. Once all required data fields are ready (i.e., written to the `StorageManager`), the data sample can be consumed by downstream tasks. + +We also track the consumption history for each computational task (e.g., `generate_sequences`, `compute_log_prob`, etc.). Therefore, even when different computational tasks require the same data field, they can consume the data independently without interfering with each other. + +

+ +

+ +To make the data retrieval process more customizable, we provide a `Sampler` class that allows users to define their own data retrieval and consumption logic. Refer to the [Customize](#customize) section for details. + +> **load-balancing** capabilities are experimentally supported in the control plane. This design enables us to offload some data management capabilities from single controller. Refer to [#PR70](https://github.com/Ascend/TransferQueue/pull/70) for details. + +### Data Plane: Distributed Data Storage + +In the data plane, we utilize a pluggable design, enabling TransferQueue to integrate with different storage backends based on user requirements. + +Specifically, we provide a `StorageManager` abstraction class that defines the core APIs as follows: + +- `async def put_data(self, data: TensorDict, metadata: BatchMeta) -> None` +- `async def get_data(self, metadata: BatchMeta) -> TensorDict` +- `async def clear_data(self, metadata: BatchMeta) -> None` + +This class encapsulates the core interaction logic within the TransferQueue system. You only need to write a simple subclass to integrate your custom storage backend. Refer to the [Customize](#customize) section for details. + +Currently, we support the following storage backends: + +- SimpleStorage: A basic CPU memory storage with minimal data format constraints and ease of use. +- [Yuanrong](https://gitee.com/openeuler/yuanrong-datasystem) (beta, [#PR107](https://github.com/TransferQueue/TransferQueue/pull/107), [#PR96](https://github.com/TransferQueue/TransferQueue/pull/96)): An Ascend native data system that provides hierarchical storage interfaces including HBM/DRAM/SSD. +- [MooncakeStore](https://github.com/kvcache-ai/Mooncake) (beta, [#PR162](https://github.com/TransferQueue/TransferQueue/pull/162)): A high-performance, KV-based hierarchical storage that supports RDMA transport between GPU and DRAM. +- [RayRDT](https://docs.ray.io/en/master/ray-core/direct-transport.html) (alpha, [#PR167](https://github.com/TransferQueue/TransferQueue/pull/167)): Ray's new feature that allows Ray to store and pass objects directly between Ray actors. + +Among them, `SimpleStorageUnit` serves as our default storage backend, coordinated by the `AsyncSimpleStorageManager` class. Each storage unit can be deployed on a separate node, allowing for distributed data management. + +`SimpleStorageUnit` employs a 2D data structure as follows: + +- Each row corresponds to a training sample, assigned a unique index within the corresponding global batch. +- Each column represents the input/output data fields for computational tasks. + +This data structure design is motivated by the computational characteristics of the post-training process, where each training sample is generated in a relayed manner across task pipelines. It provides precise addressing capabilities, enabling fine-grained, concurrent data read/write operations in a streaming manner. + +

+ +

+ +### User Interface: High-Level & Low-Level APIs + +| Level | Tier | Style | Fine-Grained Access | Streaming | Sampler | Multiple-Backends | +|---|---|---|---|------------------|---|---| +| High | **KV Interface** ([PR#28](https://github.com/Ascend/TransferQueue/pull/28))| Put/Get/List/Clear | ✓ | ○ | ✗ | ✓ | +| High | **StreamingDataLoader** ([PR#23](https://github.com/Ascend/TransferQueue/pull/23)) | PyTorch DataLoader | ✓ | ✓ | ✓ | ✓ | +| Low | **TransferQueueClient** | Metadata-based | ✓ | ✓ | ✓ | ✓ | + + +#### Key-Value based API + +To simplify the usage of TransferQueue, we provide a Redis-style high-level API that exposes most of its advanced features ([PR#28](https://github.com/Ascend/TransferQueue/pull/28)). + +**Methods** + +- **(async_)kv_put**: Insert/Update a multi-column sample by key, with an optional metadata tag. +- **(async_)kv_batch_put**: Put multiple key-value pairs efficiently in batches. +- **(async_)kv_batch_get**: Retrieve samples (by keys), supporting column selection (by fields). +- **(async_)kv_list**: List keys and tags (metadata) in a partition. +- **(async_)kv_clear**: Remove key-value pairs from storage. + +**Key Features** + +- **Redis-style Semantics**: Familiar KV interface (Put/Get/List) for a zero learning curve. +- **Fine-grained Access**: Update or retrieve specific fields (columns) within a key (row) without requiring a full-row operation. +- **Partition Isolation**: Logical separation of storage namespaces. +- **Metadata Tags**: Lightweight metadata for status tracking. +- **Pluggable Backends**: Supports multiple backends. + +Refer to [tutorials/basic.ipynb](https://github.com/Ascend/TransferQueue/blob/main/tutorial/basic.ipynb) and [tutorials/02_kv_interface.py](https://github.com/Ascend/TransferQueue/blob/main/tutorial/02_kv_interface.py) for detailed usage examples. + +#### StreamingDataLoader API + +Designed as a drop-in replacement for the standard PyTorch `DataLoader`, this API allows each rank to automatically consume data without single-controller intervention. + +In this scenario, `TransferQueueController` serves as a side-controller for data dispatching, with a user-defined `Sampler` class to organize the dataflow. +It encapsulates the complex scheduling and data transfer logic required for various parallelism strategies, seamlessly integrating TransferQueue into existing training workflows and simplifying the development of disaggregated frameworks. + +See the [Roadmap](https://github.com/Ascend/TransferQueue/issues/1) and [tutorials/06_streaming_dataloader.py](https://github.com/Ascend/TransferQueue/blob/main/tutorial/06_streaming_dataloader.py) for more details. + +#### Low-Level Native API + +The native interfaces of TransferQueue are implemented in `TransferQueueClient`. It offers maximum flexibility through native, atomic operations. + +Developers can leverage `TransferQueueClient` directly to implement advanced features that require fine-grained control and fully streamed data scheduling, as illustrated in the following tutorials: +- [tutorial/03_metadata_concepts.py](https://github.com/Ascend/TransferQueue/blob/main/tutorial/03_metadata_concepts.py) +- [tutorial/04_understanding_controller.py](https://github.com/Ascend/TransferQueue/blob/main/tutorial/04_understanding_controller.py) +- [tutorial/05_custom_sampler.py](https://github.com/Ascend/TransferQueue/blob/main/tutorial/05_custom_sampler.py) + + +

🔥 Showcases

+ +### Collocated Example + +#### verl +The primary motivation for integrating TransferQueue into verl is to **alleviate the data transfer bottleneck of the single controller `RayPPOTrainer`**. Currently, all `DataProto` objects must be routed through `RayPPOTrainer`, resulting in a single-point bottleneck for the entire post-training system. + +

+ +

+ +Official integration with verl is available at [verl/pull/5401](https://github.com/verl-project/verl/pull/5401), with the design doc at [[RFC] PPOTrainer with TransferQueue Integration](https://github.com/verl-project/verl/issues/5400). You may also refer to our [recipe](https://github.com/Ascend/TransferQueue/blob/main/recipe/simple_use_case/single_controller_demo.py), where we mimic verl usage in a high-level manner. + + +### Disaggregated Example + +We have experimentally implemented a **standardized, fully-streamed distributed** workflow via TransferQueue. + +By leveraging the `RankAwareSampler` and `StreamingDataLoader` interfaces, we achieve a **streamlined micro-batch-level producer-consumer pipeline**. This design eliminates the need to manually determine data dispatching logic across varying parallelism strategies—a typical complexity in the single-controller paradigm—thereby greatly simplifying framework design. + +Please refer to our [Roadmap](https://github.com/Ascend/TransferQueue/issues/1) and [tutorials/05_streaming_dataloader.py](https://github.com/Ascend/TransferQueue/blob/main/tutorial/05_streaming_dataloader.py) for more details. + +

+ +

+ +

🚀 Quick Start

+ +### Use Python package +```bash +pip install TransferQueue +``` + +### Install from source code + +1. Clone the source code from the GitHub repository + ```bash + git clone https://github.com/Ascend/TransferQueue/ + cd TransferQueue + ``` + +2. Install from source code + ```bash + pip install . + ``` + +### Build wheel package from source code + +1. Clone the source code from the GitHub repository + ```bash + git clone https://github.com/Ascend/TransferQueue/ + cd TransferQueue + ``` + +2. Install dependencies + ```bash + pip install build + ``` + +3. Build and install + ```bash + python -m build --wheel + pip install dist/*.whl + ``` + +

📊 Performance

+ +### Simple Case: Regular Tensor Only +

+ +

+ +### Complex Case: Regular Tensor + NestedTensor + NonTensor +

+ +

+ +> Note: Optimization for MooncakeStore and other backends are still in process. Warmly welcome contributions from the community! + +For detailed performance benchmarks, please refer to [this blog](https://www.yuque.com/haomingzi-lfse7/lhp4el/tml8ke0zkgn6roey?singleDoc#). + +We also provide a [stress test report](https://www.yuque.com/haomingzi-lfse7/lhp4el/mt0vedqy7c337pgg?singleDoc#) that demonstrates more than **8192 concurrent clients writing 2 TB of data** into TransferQueue across 4 nodes. The system remains stable without any crashes or data loss. + +

🛠️ Customize TransferQueue

+ +### Define your own data retrieval logic +We provide a `BaseSampler` abstraction class, which defines the following interface: + +```python3 +@abstractmethod +def sample( + self, + ready_indexes: list[int], + batch_size: int, + *args: Any, + **kwargs: Any, +) -> tuple[list[int], list[int]]: + """Sample a batch of indices from the ready indices. + + Args: + ready_indexes: List of global indices for which all required fields of the + corresponding samples have been produced, and the samples are not labeled as + consumed in the corresponding task. + batch_size: Number of samples to select + *args: Additional positional arguments for specific sampler implementations + **kwargs: Additional keyword arguments for specific sampler implementations + + Returns: + List of sampled global indices of length batch_size + List of global indices of length batch_size that should be labeled as consumed + (will never be retrieved in the future) + + Raises: + ValueError: If batch_size is invalid or ready_indexes is insufficient + """ + raise NotImplementedError("Subclasses must implement sample") +``` + +In this design, we separate data retrieval and data consumption through the two return values, which enables us to easily control sample replacement. We have implemented two reference designs: `SequentialSampler` and `GRPOGroupNSampler`. + +The `Sampler` class or instance should be passed to the `TransferQueueController` during initialization. During each `get_meta` call, you can provide dynamic sampling parameters to the `Sampler`. + +```python3 +from transfer_queue import TransferQueueController, TransferQueueClient, GRPOGroupNSampler, process_zmq_server_info + +# Option 1: Pass the sampler class to the TransferQueueController +controller = TransferQueueController.remote(GRPOGroupNSampler) + +# Option 2: Pass the sampler instance to the TransferQueueController (if you need custom configuration) +your_own_sampler = YourOwnSampler(config) +controller = TransferQueueController.remote(your_own_sampler) + +# Use the sampler +batch_meta = client.get_meta( + data_fields=["input_ids", "attention_mask"], + batch_size=8, + partition_id="train_0", + task_name="generate_sequences", + sampling_config={"n_samples_per_prompt": 4} # Put the required sampling parameters here +) +``` + +**Refer to [tutorial/05_custom_sampler.py](https://github.com/Ascend/TransferQueue/blob/main/tutorial/05_custom_sampler.py) for more details.** + + +### How to integrate a new storage backend + +The data plane is organized as follows: +```text + transfer_queue/ + ├── storage/ + │ ├── __init__.py + │ │── simple_backend.py # Default distributed storage backend (SimpleStorageUnit) by TQ + │ ├── managers/ # Managers are upper level interfaces that encapsulate the interaction logic with TQ system. + │ │ ├── __init__.py + │ │ ├──base.py # StorageManager, KVStorageManager, StorageManagerFactory + │ │ ├──simple_storage_manager.py # AsyncSimpleStorageManager + │ │ ├──yuanrong_manager.py # YuanrongStorageManager + │ │ └──mooncake_manager.py # MooncakeStorageManager + │ └── clients/ # Clients are lower level interfaces that directly manipulate the target storage backend. + │ │ ├── __init__.py + │ │ ├── base.py # StorageKVClient, StorageClientFactory + │ │ ├── yuanrong_client.py # YuanrongStorageClient + │ │ ├── mooncake_client.py # MooncakeStorageClient + │ │ └── ray_storage_client.py # RayStorageClient +``` + +To integrate TransferQueue with a custom storage backend, start by implementing a subclass that inherits from `StorageManager`. This subclass acts as an adapter between the TransferQueue system and the target storage backend. For KV-based storage backends, you can simply inherit from `KVStorageManager`, which can serve as the general manager for all KV-based backends. + +Distributed storage backends often come with their own native clients serving as the interface of the storage system. In such cases, a low-level adapter for this client can be written, following the examples provided in the `storage/clients` directory. + +Factory classes are provided for both `StorageManager` and `StorageClient` to facilitate easy integration. Adding necessary descriptions of required parameters in the factory class helps enhance the overall user experience. + +

✏️ Contribution Guide

+ +**Contributions are warmly welcome!** + +New ideas, feature suggestions, and user experience feedback are all encouraged—feel free to submit issues or PRs. We will respond as soon as possible. + +We recommend using pre-commit for better code format. + +```bash +# install pre-commit +pip install pre-commit + +# run the following command in your repo folder, then fix the check before committing your code +pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always +``` + + +

Citation

+Please kindly cite our paper if you find this repo is useful: + +```bibtex +@article{han2025asyncflow, + title={AsyncFlow: An Asynchronous Streaming RL Framework for Efficient LLM Post-Training}, + author={Han, Zhenyu and You, Ansheng and Wang, Haibo and Luo, Kui and Yang, Guang and Shi, Wenqi and Chen, Menglong and Zhang, Sicheng and Lan, Zeshun and Deng, Chunshi and others}, + journal={arXiv preprint arXiv:2507.01663}, + year={2025} +} +``` \ No newline at end of file diff --git a/verl/docs/examples/config.rst b/verl/docs/examples/config.rst new file mode 100644 index 0000000000000000000000000000000000000000..1c4821dcc2b1464826b92acc62c98d7cd7df9e42 --- /dev/null +++ b/verl/docs/examples/config.rst @@ -0,0 +1,729 @@ +.. _config-explain-page: + +Config Explanation +=================== + +Last updated: 06/18/2025. + +ppo_trainer.yaml for RL FSDP Backend +------------------------------------- + +Data +~~~~ + +.. code:: yaml + + data: + tokenizer: null + train_files: ~/data/rlhf/gsm8k/train.parquet + val_files: ~/data/rlhf/gsm8k/test.parquet + train_max_samples: -1 # set to -1 to use full dataset + val_max_samples: -1 # set to -1 to use full dataset + prompt_key: prompt + max_prompt_length: 512 + max_response_length: 512 + train_batch_size: 1024 + return_raw_input_ids: False # This should be set to true when the tokenizer between policy and rm differs + return_raw_chat: False + return_full_prompt: False + shuffle: True + seed: 42 + filter_overlong_prompts: False + filter_overlong_prompts_workers: 1 + truncation: error + image_key: images + trust_remote_code: True + custom_cls: + path: null + name: null + +- ``data.train_files``: Training set parquet. Can be a list or a single + file. The program will read all files into memory, so it can't be too + large (< 100GB). The path can be either local path or HDFS path. For + HDFS path, we provide utils to download it to DRAM and convert the + HDFS path to local path. +- ``data.val_files``: Validation parquet. Can be a list or a single + file. +- ``data.train_max_samples``: Maximum number of samples to use from the + training dataset. Set to -1 to use the full dataset. +- ``data.val_max_samples``: Maximum number of samples to use from the + validation dataset. Set to -1 to use the full dataset. +- ``data.prompt_key``: The field in the dataset where the prompt is + located. Default is 'prompt'. +- ``data.max_prompt_length``: Maximum prompt length. All prompts will be + left-padded to this length. An error will be reported if the length is + too long +- ``data.max_response_length``: Maximum response length. Rollout in RL + algorithms (e.g. PPO) generates up to this length +- ``data.train_batch_size``: Batch size sampled for one training + iteration of different RL algorithms. +- ``data.return_raw_input_ids``: Whether to return the original + input_ids without adding chat template. This is mainly used to + accommodate situations where the reward model's chat template differs + from the policy. It needs to be decoded first, then apply the RM's + chat template. If using a model-based RM, and the policy and RM + chat_templates are different, this flag needs to be set +- ``data.return_raw_chat``: Whether to return the original chat (prompt) + without applying chat template. +- ``data.return_full_prompt``: Whether to return the full prompt with chat template +- ``data.shuffle``: Whether to shuffle the data in the dataloader. +- ``data.seed``: An integer seed to use when shuffling the data. If not set or set to + `null`, the data shuffling will not be seeded, resulting in a different data order on each run. +- ``data.filter_overlong_prompts``: Default don't filter. +- ``data.filter_overlong_prompts_workers``: For large-scale dataset, filtering + overlong prompts could be timeconsuming. You cat set the ``filter_overlong_prompts_workers`` + to use multiprocessing for speed up. Default to 1. +- ``data.truncation``: Truncate the input_ids or prompt length if they + exceed max_prompt_length. Default is 'error', not allow exceed the + max_prompt_length. The users should increase the max_prompt_length if + throwing the error. You can also set ``left``, ``right`` and ``middle``. + When ``middle`` is selected, the logic splits the allowed max length roughly in half + and keeps the head and tail of the sequence, effectively discarding the middle section. +- ``data.image_key``: The field in the multi-modal dataset where the image is + located. Default is 'images'. +- ``data.trust_remote_code``: If the remote tokenizer has python file, we can use this field to allow + using remote tokenizer. For example: moonshotai/Moonlight-16B-A3B-Instruct + +Customized Dataset +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Customized dataset extension is implemented for the SFT trainer and can be extended to other trainers with similar changes. + +.. code:: yaml + + custom_cls: + path: null + name: null + +- ``data.custom_cls.path``: The path to the file containing your customized dataset class. If not specified, pre-implemented dataset will be used. +- ``data.custom_cls.name``: The name of the dataset class within the specified file. + +Actor/Rollout/Reference Policy +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. code:: yaml + + actor_rollout_ref: + hybrid_engine: True + model: + path: ~/models/deepseek-llm-7b-chat + external_lib: null + override_config: + attn_implementation: flash_attention_2 # or eager, sdpa - attention implementation override + model_config: {} + moe_config: # Megatron only, can adjust moe configuration + freeze_moe_router: False # Megatron only, can freeze moe router (no grad) + enable_gradient_checkpointing: False + enable_activation_offload: False + trust_remote_code: False + use_remove_padding: False + actor: + strategy: fsdp # This is for backward-compatibility + ppo_mini_batch_size: 256 + ppo_micro_batch_size: null # will be deprecated, use ppo_micro_batch_size_per_gpu + ppo_micro_batch_size_per_gpu: 8 + use_dynamic_bsz: False + ppo_max_token_len_per_gpu: 16384 # n * ${data.max_prompt_length} + ${data.max_response_length} + grad_clip: 1.0 + clip_ratio: 0.2 + entropy_coeff: 0.0 + use_kl_loss: False # True for GRPO + # Rollout Correction (corrects distribution mismatch between rollout and training) + rollout_correction: + rollout_is: token # IS weights + rollout_is_threshold: 2.0 # TIS upper bound, or "0.5_5.0" for IcePop + rollout_rs: null # Rejection sampling + rollout_rs_threshold: null # RS upper threshold + use_torch_compile: True # False to disable torch compile + kl_loss_coef: 0.001 # for grpo + kl_loss_type: low_var_kl # for grpo + ppo_epochs: 1 + data_loader_seed: null + shuffle: False + ulysses_sequence_parallel_size: 1 # sp size + optim: + lr: 1e-6 + lr_warmup_steps: -1 # Prioritized. Negative values mean delegating to lr_warmup_steps_ratio. + lr_warmup_steps_ratio: 0. # the total steps will be injected during runtime + min_lr_ratio: 0.0 # only used with cosine lr scheduler, default to 0.0 + num_cycles: 0.5 # only used with cosine lr scheduler, default to 0.5 + lr_scheduler_type: constant # select from constant/cosine + total_training_steps: -1 # must be override by program + fsdp_config: + wrap_policy: + # transformer_layer_cls_to_wrap: None + min_num_params: 0 + param_offload: False + optimizer_offload: False + fsdp_size: -1 + checkpoint: + # What to include in saved checkpoints + # with 'hf_model' you can save whole model as hf format, now only use sharded model checkpoint to save space + save_contents: ['model', 'optimizer', 'extra'] + # For more flexibility, you can specify the contents to load from the checkpoint. + load_contents: ${actor_rollout_ref.actor.checkpoint.save_contents} + ref: + fsdp_config: + param_offload: False + wrap_policy: + # transformer_layer_cls_to_wrap: None + min_num_params: 0 + log_prob_micro_batch_size: null # will be deprecated, use log_prob_micro_batch_size_per_gpu + log_prob_micro_batch_size_per_gpu: 16 + log_prob_use_dynamic_bsz: ${actor_rollout_ref.actor.use_dynamic_bsz} + log_prob_max_token_len_per_gpu: ${actor_rollout_ref.actor.ppo_max_token_len_per_gpu} + ulysses_sequence_parallel_size: ${actor_rollout_ref.actor.ulysses_sequence_parallel_size} # sp size + rollout: + name: vllm + temperature: 1.0 + top_k: -1 # 0 for hf rollout, -1 for vllm rollout + top_p: 1 + prompt_length: ${data.max_prompt_length} # not use for opensource + response_length: ${data.max_response_length} + # for vllm rollout + dtype: bfloat16 # should align with FSDP + gpu_memory_utilization: 0.5 + ignore_eos: False + enforce_eager: True + free_cache_engine: True + load_format: dummy_dtensor + tensor_model_parallel_size: 2 + max_num_batched_tokens: 8192 + max_num_seqs: 1024 + log_prob_micro_batch_size: null # will be deprecated, use log_prob_micro_batch_size_per_gpu + log_prob_micro_batch_size_per_gpu: 16 + log_prob_use_dynamic_bsz: ${actor_rollout_ref.actor.use_dynamic_bsz} + log_prob_max_token_len_per_gpu: ${actor_rollout_ref.actor.ppo_max_token_len_per_gpu} + # for hf rollout + do_sample: True + engine_kwargs: # inference engine parameters, please refer vllm/sglang official doc for detail + vllm: {} + sglang: {} + + n: 1 # for each prompt, sample n responses (i.e. num sample times). set it to values > 1 for grpo, rloo + calculate_log_probs: False # set to True for computing log probs via rollouts + val_kwargs: + # sampling parameters for validation + top_k: -1 # 0 for hf rollout, -1 for vllm rollout + top_p: 1.0 + temperature: 0 + n: 1 + do_sample: False # default eager for validation + + agent: + custom_async_server: # Use custom async server implementation for rollout + path: null + name: null + +**Common config for actor, rollout and reference model** + +- ``actor_rollout_ref.hybrid_engine``: Whether it's a hybrid engine, + currently only supports hybrid engine +- ``actor_rollout_ref.model.path``: Huggingface model path. This can be + either local path or HDFS path. For HDFS path, we provide utils to + download it to DRAM and convert the HDFS path to local path. +- ``actor_rollout_ref.model.external_libs``: Additional Python packages + that need to be imported. Used to register models or tokenizers into + the Huggingface system. +- ``actor_rollout_ref.model.override_config``: Used to override some of + the model's original configurations. Common overrides include: + + - ``attn_implementation``: Override the attention implementation. Default is ``flash_attention_2``. + Supported values: ``flash_attention_2``, ``eager``, ``sdpa``. Use ``eager`` for debugging or + compatibility issues. See :ref:`attention-implementation-override` for detailed usage. + +- ``actor_rollout_ref.model.enable_gradient_checkpointing``: FSDP only, decide + Whether to enable gradient checkpointing for the actor, + Megatron uses recompute options in ``override_transformer_config`` to set this +- ``actor_rollout_ref.model.enable_activation_offload``: Whether to enable + activation offloading for the actor +- ``actor_rollout_ref.model.trust_remote_code``: Whether to enable loading + a remote code model +- ``actor_rollout_ref.model.use_fused_kernels``: Whether to use fused + kernels in the model. If set to True, the following parameters will be + used. + + - ``actor_rollout_ref.model.fused_kernel_options.impl_backend``: The + implementation backend for fused kernels. Options: "triton" or + "torch". Default is "torch". + While in megatron, we only support "triton" as the + implementation backend, so there is no need for this option. + +- ``actor_rollout_ref.model.use_remove_padding``: Whether to use remove + padding in the model. If set to True, the model will remove padding + tokens in the input_ids and response_ids. This helps a lot in improving model running efficiency. + +- ``actor_rollout_ref.model.tiled_mlp``: TiledMLP configuration for memory-efficient + MLP computation. Reduces peak memory by processing MLP forward/backward in tiles. + Only compatible with FSDP2 (requires ``actor_rollout_ref.actor.strategy=fsdp2``). + + - ``actor_rollout_ref.model.tiled_mlp.enabled``: Whether to enable TiledMLP. + Default is False. + - ``actor_rollout_ref.model.tiled_mlp.num_shards``: Number of shards to split + the input. Higher values reduce peak memory but may slightly impact performance. + Default is 4. + +**Actor model** + +- ``actor_rollout_ref.actor.strategy``: fsdp or megatron. In this + example, we use fsdp backend. + +- ``actor_rollout_ref.actor.ppo_mini_batch_size``: One sample is split + into multiple sub-batches with batch_size=ppo_mini_batch_size for PPO + updates. The ppo_mini_batch_size is a global num across all workers/gpus + +- ``actor_rollout_ref.actor.ppo_micro_batch_size``: [Will be deprecated, use ppo_micro_batch_size_per_gpu] + Similar to gradient accumulation, the micro_batch_size_per_gpu for one forward pass, + trading speed for GPU memory. The value represent the global view. + +- ``actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu``: Similar to gradient + accumulation, the micro_batch_size_per_gpu for one forward pass, trading speed + for GPU memory. The value represent the local num per gpu. + +- ``actor_rollout_ref.actor.grad_clip``: Gradient clipping for actor + updates +- ``actor_rollout_ref.actor.use_kl_loss``: to use kl loss in actor. When used, we are not applying KL in the reward function. + +- ``actor_rollout_ref.actor.clip_ratio``: PPO clip ratio + +- ``actor_rollout_ref.actor.use_torch_compile``: Whether to use torch compile in actor + +- ``actor_rollout_ref.actor.entropy_coeff``: The weight of entropy when + calculating PPO loss. The default value is changed to 0.0 since v0.3.x + +- ``actor_rollout_ref.actor.ppo_epochs``: Number of epochs for PPO + updates on one set of sampled data + +- ``actor_rollout_ref.actor.data_loader_seed``: From torch 2.6.0 Megatron backend can get wrong seed generated by pytorch + between cp ranks and cause misalignment between data on these ranks, so we shall manually set the seed to avoid hanging + issue. if ``actor_rollout_ref.actor.shuffle`` is not null, this must be set. + +- ``actor_rollout_ref.actor.shuffle``: Whether to shuffle data when + there are multiple epochs + +- ``actor_rollout_ref.actor.optim``: Actor's optimizer parameters + +- ``actor_rollout_ref.actor.fsdp_config``: FSDP config for actor + training + + - ``wrap_policy``: FSDP wrap policy. By default, it uses Huggingface's + wrap policy, i.e., wrapping by DecoderLayer + + - No need to set transformer_layer_cls_to_wrap, so we comment it. + + - ``*_offload``: Whether to enable parameter, gradient and optimizer + offload + + - Trading speed for GPU memory. + +- ``actor_rollout_ref.actor.use_kl_loss``: Whether to enable kl loss. Default is False. + +- ``actor_rollout_ref.actor.kl_loss_coef``: The coefficient of kl loss. Default is 0.001. + +- ``actor_rollout_ref.actor.kl_loss_type``: Support ``kl`` (``k1``), ``abs``, ``mse`` (``k2``), ``low_var_kl`` (``k3``) and ``full``. Appending ``+`` in the end (e.g., ``k1+`` and ``k3+``) would use straight-through to employ ``k2`` for unbiased gradient estimation, regardless of the kl value estimation (see https://github.com/verl-project/verl/pull/2953#issuecomment-3162113848 for more details). How to calculate the kl divergence between actor and reference policy. For specific options, refer to `kl_penalty()` in `core_algos.py `_ . See this blog post for detailed analysis: http://joschu.net/blog/kl-approx.html + +- ``actor_rollout_ref.actor.checkpoint``: The configurations of checkpoint function in actor + + - ``save_contents``: The contents to save in the checkpoint. By default, we save model, optimizer and extra information in the checkpoint. + The extra information includes Rng states currently, FSDP supported lr_scheduler, and Megatron opt_param_scheduler will coming soon. + We do not store hf_model in checkpoint by default, but we provide a tool in ``scripts/model_merge.py`` to convert checkpoint format to hf format. + + - ``load_contents``: The contents to load in the checkpoint, you can specify different checkpoint loading contents. By default, it is the same with ``save_checkpoint``. + +**Reference Model** + +Reference model will be enabled when ``actor.use_kl_loss`` or/and ``algorithm.use_kl_in_reward`` is/are True. + +- ``actor_rollout_ref.ref``: FSDP config same as actor. **For models + larger than 7B, it's recommended to turn on offload for ref by + default** + +- ``actor_rollout_ref.ref.log_prob_micro_batch_size``: [Will be deprecate, use log_prob_micro_batch_size_per_gpu] + The batch size for one forward pass in the computation of ``ref_log_prob``. The value represent the global num. + +- ``actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu``: The batch size + for one forward pass in the computation of ``ref_log_prob``. The value represent the local num per gpu. + +**Rollout Model** + +- ``actor_rollout_ref.rollout.name``: hf/vllm/sglang. + +- Rollout (Auto-regressive) parameters. The key should be equal to the + property name in vLLM's ``SamplingParams``. + + - ``temperature``, ``top_k``, ``top_p`` and others: Sampling + parameters in ``SamplingParams``. + +- ``actor_rollout_ref.rollout.dtype``: Rollout model parameters type. This should be align with + the actor model parameter type in FSDP/Megatron backend. + +- ``actor_rollout_ref.rollout.gpu_memory_utilization``: + + - For vLLM v0.7.0 and later: The fraction of **total** GPU memory to be used for the vLLM instance. + - For SGLang: Corresponding to ``mem_fraction_static``, the fraction of the free GPU memory used for **static** memory like model weights and KV cache. + +- ``actor_rollout_ref.rollout.tensor_model_parallel_size``: TP size for rollout. Only effective + for vllm. + +- ``actor_rollout_ref.rollout.log_prob_micro_batch_size``: [Will be deprecate, use log_prob_micro_batch_size_per_gpu] + The batch size for one forward pass in the computation of ``log_prob``. The value represent the global num. + +- ``actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu``: Micro batch size per gpu (The batch size for + one forward pass) for recalculating ``log_prob``. The value represent the local num per gpu. + +- ``actor_rollout_ref.rollout.do_sample``: Whether to sample during training rollout. If set to False, the rollout model + will perform greedy sampling. + +- ``actor_rollout_ref.rollout.val_kwargs```: Sampling parameters used specifically during validation. + + - ``top_k``: Top-k sampling parameter. Default to -1 for vLLM rollout or 0 for HF rollout. + - ``top_p``: Top-p sampling parameter. Default is 1.0 (disabled). + - ``temperature``: Sampling temperature. Default is 0 (deterministic greedy). + - ``n``: Number of responses to generate during validation. Default is 1. + - ``do_sample``: Whether to use sampling during validation. Default is False for + deterministic outputs. When set to True, the rollout will use the ``actor_rollout_ref.rollout.val_kwargs`` parameters + (top_k, top_p, temperature) to control the sampling behavior. + +- ``actor_rollout_ref.rollout.engine_kwargs.vllm``: extra vllm engine args, please refer vllm official doc for detail + +- ``actor_rollout_ref.rollout.engine_kwargs.sglang``: extra sglang engine args, please refer sglang official doc for detail + +- ``actor_rollout_ref.rollout.ignore_eos``: Whether to ignore the EOS + token and continue generating tokens after the EOS token is generated. + +- ``actor_rollout_ref.rollout.free_cache_engine``: Offload the KVCache + after rollout generation stage. Default is True. When set to True, + for vllm v0.5.4 and v0.6.3, we need to disable the usage of CUDAGraph + (set ``enforce_eager`` to True.) + +- ``actor_rollout_ref.rollout.enforce_eager``: Whether to use CUDAGraph + in vLLM generation. Default set to True to disable CUDAGraph. + +- ``actor_rollout_ref.rollout.load_format``: Which weight loader to use + to load the actor model weights to the rollout model. + + - ``auto``: Use Megatron weight loader. + - ``megatron``: Use Megatron weight loader. Deployed with Megatron + backend. The input model ``state_dict()`` is already partitioned + along TP dimension and already gathered along PP dimension. This + weight loader requires that the Rollout model and Actor model's + parameters shape and name should be identical. + - ``dtensor``: Default solution when using Huggingface weight loader. + Deployed with FSDP backend and the state_dict_type is + ``StateDictType.SHARDED_STATE_DICT``. Recommend to use this weight + loader + - ``hf``: Use Huggingface weight loader. Deployed with FSDP backend + and the state_dict_type is ``StateDictType.FULL_STATE_DICT``. This + solution doesn't need to rewrite the weight loader for each model + implemented in vLLM but it results in larger peak memory usage. + - ``dummy_hf``, ``dummy_megatron``, ``dummy_dtensor``: Random + initialization. + +.. note:: **NOTED**: In this config field, users only need to select from ``dummy_megatron``, ``dummy_dtensor``, ``dummy_hf`` for rollout initialization and our hybrid engine will select the corresponding weight loader (i.e., ``megatron``, ``dtensor``, ``hf``) during actor/rollout weight synchronization. + + +Megatron Optimizer and Optimizer Parameter Scheduler +____________________________________________________ + +.. code:: yaml + + optim: + optimizer: adam + lr: 1e-6 + clip_grad: 1.0 + total_training_steps: -1 # must be override by program + lr_warmup_init: 0.0 # initial learning rate for warmup, default to 0.0 + lr_warmup_steps: -1 # Prioritized. Negative values mean delegating to lr_warmup_steps_ratio. + lr_warmup_steps_ratio: 0. # the total steps will be injected during runtime + lr_decay_steps: null + lr_decay_style: constant # select from constant/linear/cosine/inverse_square_root + min_lr: 0.0 # minimum learning rate, default to 0.0 + weight_decay: 0.01 + weight_decay_incr_style: constant # select from constant/linear/cosine + lr_wsd_decay_style: exponential # select from constant/exponential/cosine + lr_wsd_decay_steps: null + use_checkpoint_opt_param_scheduler: False # use checkpoint optimizer parameter scheduler + + +Notice that there are some differences in APIs between Megatron optimizer and FSDP optimizer. + +- Megatron optimizer scheduler names the period after lr_warmup as lr_decay_steps, so the ``lr_scheduler_type`` actually means the style of lr decay after warmup. +- Megatron optimizer also support weight decay decay mechanism +- ``use_checkpoint_opt_param_scheduler`` determines whether to use the checkpoint optimizer parameter scheduler. If set to True, the optimizer parameter scheduler will be saved in the checkpoint and loaded from the checkpoint during resuming training. + +For learning rate decay, original Megatron pretrain default option of ``lr_decay_style`` is ``linear``, +meaning that the learning rate will be linearly decayed from the initial learning rate to ``min_lr`` within the +``lr_decay_steps``. However, in verl, to align with FSDP's default behavior, we set the default +``lr_decay_style`` to ``constant``, meaning that the learning rate will be kept constant after the warmup stage. + + +Critic Model +~~~~~~~~~~~~ + +Most parameters for Critic are similar to Actor Model. + +Reward Model +~~~~~~~~~~~~ + +.. code:: yaml + + reward_model: + enable: False + model: + input_tokenizer: ${actor_rollout_ref.model.path} # set this to null if the chat template is identical + path: ~/models/Anomy-RM-v0.1 + external_lib: ${actor_rollout_ref.model.external_lib} + trust_remote_code: False + fsdp_config: + min_num_params: 0 + param_offload: False + micro_batch_size_per_gpu: 16 + max_length: null + reward_manager: naive + +- ``reward_model.enable``: Whether to enable reward model. If False, we + compute the reward only with the user-defined reward functions. In + GSM8K and Math examples, we disable reward model. For RLHF alignment + example using full_hh_rlhf, we utilize reward model to assess the + responses. If False, the following parameters are not effective. +- ``reward_model.model`` + + - ``input_tokenizer``: Input tokenizer. If the reward model's chat + template is inconsistent with the policy, we need to first decode to + plaintext, then apply the rm's chat_template. Then score with RM. If + chat_templates are consistent, it can be set to null. + - ``path``: RM's HDFS path or local path. Note that RM only supports + AutoModelForSequenceClassification. Other model types need to define + their own RewardModelWorker and pass it from the code. + - ``trust_remote_code``: Whether to enable loading a remote code model, + default to False. +- ``reward_model.reward_manager``: Reward Manager. This defines the mechanism + of computing rule-based reward and handling different reward sources. Default + is ``naive``. If all verification functions are multiprocessing-safe, the reward + manager can be set to ``prime`` for parallel verification. + +Customized Reward Function +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. code:: yaml + + custom_reward_function: + path: null + name: compute_score + +- ``custom_reward_function.path``: The path to the file containing your customized reward function. If not specified, pre-implemented reward functions will be used. +- ``custom_reward_function.name`` (Optional) : The name of the reward function within the specified file. Default is 'compute_score'. + +Algorithm +~~~~~~~~~ + +.. code:: yaml + + algorithm: + gamma: 1.0 + lam: 1.0 + adv_estimator: gae + use_kl_in_reward: False + kl_penalty: kl # how to estimate kl divergence + kl_ctrl: + type: fixed + kl_coef: 0.005 + horizon: 10000 + target_kl: 0.1 + # Rollout Correction + rollout_correction: + rollout_is: null # IS weights + rollout_is_threshold: 2.0 # Upper threshold for IS weights + rollout_rs: null # Rejection sampling + rollout_rs_threshold: null # RS upper threshold + +- ``gamma``: discount factor +- ``lam``: Trade-off between bias and variance in the GAE estimator +- ``adv_estimator``: Support ``gae``, ``grpo``, ``reinforce_plus_plus``, ``reinforce_plus_plus_baseline``, ``rloo``, ``rloo_vectorized``, ``grpo_vectorized`` +- ``use_kl_in_reward``: Whether to enable in-reward kl penalty. Default is False. +- ``kl_penalty``: Support ``kl``, ``abs``, ``mse``, ``low_var_kl`` and ``full``. How to + calculate the kl divergence between actor and reference policy. For + specific options, refer to `kl_penalty()` in `core_algos.py `_ . +- ``kl_ctrl``: Config for in-reward kl_penalty controller + + - ``kl_coef``: The (initial) coefficient of in-reward kl_penalty. Default is 0.001. + - ``type``: 'fixed' for FixedKLController and 'adaptive' for AdaptiveKLController. + - ``horizon`` and ``target_kl``: See source code of AdaptiveKLController for details. + +- ``rollout_correction``: Rollout Correction configuration (nested dict). Set to ``null`` to disable. + When enabled, contains: + + - ``rollout_is``: IS weights aggregation level, ``null`` to disable IS weights. + - ``rollout_is_threshold``: Upper threshold for IS weights (e.g., 2.0). + - ``rollout_rs``: Rejection sampling mode, ``null`` to disable RS. + - ``rollout_rs_threshold``: RS upper threshold. + + Note: Rollout Correction requires setting ``actor_rollout_ref.rollout.calculate_log_probs=True``. + +Trainer +~~~~~~~ + +.. code:: yaml + + trainer: + total_epochs: 30 + project_name: verl_examples + experiment_name: gsm8k + logger: ['console', 'wandb'] + log_val_generations: 0 + nnodes: 1 + n_gpus_per_node: 8 + save_freq: -1 + val_before_train: True + test_freq: 2 + critic_warmup: 0 + default_hdfs_dir: null # hdfs checkpoint path + default_local_dir: checkpoints/${trainer.project_name}/${trainer.experiment_name} # local checkpoint path + resume_mode: auto # or disable or resume_path if resume_from_path is set + resume_from_path: null + remove_previous_ckpt_in_save: False + del_local_ckpt_after_load: False + ray_wait_register_center_timeout: 300 + +- ``trainer.total_epochs``: Number of epochs in training. +- ``trainer.project_name``: For wandb, swanlab, mlflow +- ``trainer.experiment_name``: For wandb, swanlab, mlflow +- ``trainer.logger``: Support console and wandb, swanlab, mlflow, tensorboard, trackio +- ``trainer.log_val_generations``: The number of logged generation during validation (default ``0``) +- ``trainer.nnodes``: Number of nodes used in the training. +- ``trainer.n_gpus_per_node``: Number of GPUs per node. +- ``trainer.save_freq``: The frequency (by iteration) to save checkpoint + of the actor and critic model. +- ``trainer.val_before_train``: Whether to run validation before training. +- ``trainer.test_freq``: The validation frequency (by iteration). +- ``trainer.critic_warmup``: The number of iteration to train the critic + model before actual policy learning. +- ``trainer.resume_mode``: The mode of resuming training. Support + ``disable``, ``auto`` and ``resume_path``. If set to ``auto`` as default, the + program will automatically resume from the latest checkpoint in the + ``default_local_dir``. If set to ``resume_path``, the program will resume + from the path specified in ``resume_from_path``. +- ``trainer.resume_from_path``: The path to resume training from. Only + effective when ``resume_mode`` is set to ``resume_path``. +- ``trainer.remove_previous_ckpt_in_save``: Whether to remove previous + checkpoints in the save directory. Default is False. +- ``trainer.del_local_ckpt_after_load``: Whether to delete local + checkpoints after loading them. Default is False. +- ``trainer.ray_wait_register_center_timeout``: The timeout for waiting + for the ray register center to be ready. Default is 300 seconds. + + +This figure illustrates how the configurations affect the training. + +https://excalidraw.com/#json=pfhkRmiLm1jnnRli9VFhb,Ut4E8peALlgAUpr7E5pPCA + +.. image:: https://github.com/user-attachments/assets/16aebad1-0da6-4eb3-806d-54a74e712c2d + + +evaluation.yaml +--------------- + +Data +~~~~ + +.. code:: yaml + + data: + path: /tmp/math_Qwen2-7B-Instruct.parquet + prompt_key: prompt + response_key: responses + data_source_key: data_source + reward_model_key: reward_model + +- ``data.path``: Path to the dataset file (Parquet format). +- ``data.prompt_key``: The field in the dataset where the prompt is located. Default is 'prompt'. +- ``data.response_key``: The key holds the generated responses. This should be a list of strings representing the responses. Default is 'responses'. +- ``data.data_source_key``: This is used to separate metric calculations for different data sources, ensuring that metrics are calculated independently for each source. +- ``data.reward_model_key``: The key holds the reference answers. These reference answers typically serve as the ground truth or test cases for the task. + +Customized Reward Function +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. code:: yaml + + custom_reward_function: + path: null + name: compute_score + +- ``custom_reward_function.path``: The path to the file containing your customized reward function. If not specified, pre-implemented reward functions will be used. +- ``custom_reward_function.name`` (Optional) : The name of the reward function within the specified file. Default is 'compute_score'. + +sft_trainer.yaml for SFT FSDP Backend +-------------------------------------- + + +Optim +~~~~~~~ + +.. code:: yaml + + optim: + optimizer: AdamW + optimizer_impl: torch.optim + lr: 1e-5 + weight_decay: 0.01 + lr_warmup_steps_ratio: 0.1 + clip_grad: 1.0 + lr_scheduler: cosine + override_optimizer_config: null + +- ``optimizer``: Optimizer class name (e.g., ``"AdamW"``, ``"AdamW8bit"``, ``"_AdamW"``). The class name as it appears in the module. +- ``optimizer_impl``: Module path to import optimizer from (e.g., ``"torch.optim"``, ``"torchao.optim"``, ``"bitsandbytes.optim"``). +- ``optim.lr``: Learning rate for the optimizer. +- ``optim.weight_decay``: Weight decay for the optimizer. +- ``optim.lr_warmup_steps_ratio``: Ratio of warmup steps to total training steps. +- ``optim.clip_grad``: Gradient clipping value. +- ``optim.lr_scheduler``: Learning rate scheduler type. Options: + + - ``cosine``: Cosine learning rate scheduler with warmup (default). + - ``wsd``: Warmup-Stable-Decay scheduler that provides a stable learning rate phase between warmup and decay phases. + +- ``override_optimizer_config``: Dictionary of additional optimizer-specific keyword arguments. For example, to use ``torchao.optim``'s ``_AdamW`` with BF16 stochastic rounding: ``{"bf16_stochastic_round": true}`` + +Model +~~~~~~~~~~~~ + +Most parameters for Model are similar to Reward Model. + +.. code:: yaml + + model: + partial_pretrain: ~/models/gemma-1.1-7b-it + fsdp_config: + model_dtype: fp32 + wrap_policy: + min_num_params: 0 + cpu_offload: False + offload_params: False + external_lib: null + enable_gradient_checkpointing: False + trust_remote_code: False + lora_rank: 0 + lora_alpha: 16 + target_modules: all-linear + use_liger: False + +- ``partial_pretrain``: HDFS path or local path for the pretrained model. +- ``fsdp_config`` + + - ``model_dtype``: Model parameters type, default to ``fp32``. + Support: ``bf16``, ``fp16``, ``fp32``. + - ``cpu_offload``: Whether to enable CPU offloading for FSDP. If True, + the offload_params will be used as argument. + - ``offload_params``: Whether to offload parameters to CPU + when not involved in computation. If True, then this offloads gradients + to CPU as well, meaning that the optimizer step runs on CPU. + +- ``lora_rank``: The rank of the LoRA model, default to 0. If ``lora_rank``>0, + we will train LoRA modules instead of tuning the full model. +- ``lora_alpha``: The alpha parameter for LoRA scaling, default to 16. +- ``target_modules``: The names of the modules to apply the adapter to, + default to ``all-linear``. See `peft docs `_ for detail. + +- ``use_liger``: Whether to enable Liger kernel, default to False. If True, + we apply Liger kernel to the model (depends on `liger-kernel`). diff --git a/verl/docs/examples/gsm8k_example.rst b/verl/docs/examples/gsm8k_example.rst new file mode 100644 index 0000000000000000000000000000000000000000..e653fe065e23d8d38736f4a2c29f705615d70aa9 --- /dev/null +++ b/verl/docs/examples/gsm8k_example.rst @@ -0,0 +1,188 @@ +GSM8K Example +============= + +Last updated: 03/25/2025. + +Introduction +------------ + +In this example, we train an LLM to tackle the GSM8k task. + +Paper: https://arxiv.org/pdf/2110.14168 + +Dataset: https://huggingface.co/datasets/openai/gsm8k + +Note that the original paper mainly focuses on training a verifier (a +reward model) to solve math problems via Best-of-N sampling. In this +example, we train an RLHF agent using a rule-based reward model. + +Dataset Introduction +-------------------- + +GSM8k is a math problem dataset. The prompt is an elementary school +problem. The LLM model is required to answer the math problem. + +The training set contains 7473 samples and the test set contains 1319 +samples. + +**An example** + +Prompt + + Katy makes coffee using teaspoons of sugar and cups of water in the + ratio of 7:13. If she used a total of 120 teaspoons of sugar and cups + of water, calculate the number of teaspoonfuls of sugar she used. + +Solution + + The total ratio representing the ingredients she used to make the + coffee is 7+13 = <<7+13=20>>20 Since the fraction representing the + number of teaspoons she used is 7/20, she used 7/20\ *120 = + <<7/20*\ 120=42>>42 #### 42 + +Step 1: Prepare dataset +----------------------- + +.. code:: bash + + cd examples/data_preprocess + python3 gsm8k.py --local_save_dir ~/data/gsm8k + +Step 2: Download Model +---------------------- + +There're three ways to prepare the model checkpoints for post-training: + +- Download the required models from huggingface or modelscope + +.. code:: bash + + hf download deepseek-ai/deepseek-math-7b-instruct --local-dir ~/models/deepseek-math-7b-instruct --local-dir-use-symlinks False + # or + modelscope download --model deepseek-ai/deepseek-math-7b-instruct --local_dir ~/models/deepseek-math-7b-instruct + +- Already store your store model in the local directory or HDFS path. +- Also, you can directly use the model name in huggingface (e.g., + deepseek-ai/deepseek-math-7b-instruct) in + ``actor_rollout_ref.model.path`` and ``critic.model.path`` field in + the run script. You can also download models from modelscope by setting environmental variable ``VERL_USE_MODELSCOPE=True``. + +Noted that users should prepare checkpoints for actor, critic and reward +model. + +[Optional] Step 3: SFT your Model +--------------------------------- + +We provide a SFT Trainer using PyTorch FSDP in +`sft_trainer.py `_. +Users can customize their own SFT +script using our FSDP SFT Trainer. + +We also provide various training scripts for SFT on GSM8K dataset in `gsm8k sft directory `_. + +.. code:: shell + + set -x + + torchrun -m verl.trainer.sft_trainer \ + data.train_files=$HOME/data/gsm8k/train.parquet \ + data.val_files=$HOME/data/gsm8k/test.parquet \ + data.messages_key=messages \ + data.micro_batch_size_per_gpu=8 \ + model.path=deepseek-ai/deepseek-coder-6.7b-instruct \ + trainer.project_name=gsm8k-sft \ + trainer.experiment_name=gsm8k-sft-deepseek-coder-6.7b-instruct \ + trainer.total_epochs=4 \ + trainer.logger='["console","wandb"]' + + +If you use AMD GPUs (ROCm kernel), you need to add the following environment variables into the run script: + + .. code-block:: bash + + export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 + export ROCR_VISIBLE_DEVICES=$HIP_VISIBLE_DEVICES + export CUDA_VISIBLE_DEVICES=$HIP_VISIBLE_DEVICES + + +Step 4: Perform PPO training with your model on GSM8K Dataset +------------------------------------------------------------- + +- Prepare your own run.sh script. Here's an example for GSM8k dataset + and deepseek-llm-7b-chat model. +- Users could replace the ``data.train_files`` ,\ ``data.val_files``, + ``actor_rollout_ref.model.path`` and ``critic.model.path`` based on + their environment. +- See :doc:`config` for detailed explanation of each config field. + +**Reward Model/Function** + +We use a rule-based reward model. We force the model to produce a final +answer following 4 “#” as shown in the solution. We extract the final +answer from both the solution and model's output using regular +expression matching. We compare them and assign a reward of 1 to correct +answer, 0.1 to incorrect answer and 0 to no answer. + +**Training Script** + +The training script example for FSDP and Megatron-LM backend are stored in examples/ppo_trainer directory. + +.. code:: bash + + cd ../ppo_trainer + bash run_deepseek_llm_7b_fsdp.sh + +The script of run_deepseek_llm_7b_fsdp.sh + +.. code:: bash + + set -x + + python3 -m verl.trainer.main_ppo \ + data.train_files=$HOME/data/gsm8k/train.parquet \ + data.val_files=$HOME/data/gsm8k/test.parquet \ + data.train_batch_size=1024 \ + data.max_prompt_length=512 \ + data.max_response_length=512 \ + actor_rollout_ref.model.path=deepseek-ai/deepseek-llm-7b-chat \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.model.use_remove_padding=True \ + actor_rollout_ref.actor.ppo_mini_batch_size=256 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \ + actor_rollout_ref.actor.fsdp_config.param_offload=False \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=4 \ + actor_rollout_ref.rollout.name=vllm \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 \ + actor_rollout_ref.ref.fsdp_config.param_offload=True \ + critic.optim.lr=1e-5 \ + critic.model.use_remove_padding=True \ + critic.model.path=deepseek-ai/deepseek-llm-7b-chat \ + critic.model.enable_gradient_checkpointing=True \ + critic.ppo_micro_batch_size_per_gpu=32 \ + critic.fsdp.param_offload=False \ + critic.fsdp.optimizer_offload=False \ + algorithm.kl_ctrl.kl_coef=0.001 \ + trainer.critic_warmup=0 \ + trainer.logger='["console","wandb"]' \ + trainer.project_name='verl_example_gsm8k' \ + trainer.experiment_name='deepseek_llm_7b_function_rm' \ + trainer.n_gpus_per_node=8 \ + trainer.nnodes=1 \ + trainer.save_freq=-1 \ + trainer.test_freq=1 \ + trainer.total_epochs=15 $@ + + +If you use AMD GPUs (ROCm kernel), you need to add the following environment variables into the run script: + + .. code-block:: bash + + export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 + export ROCR_VISIBLE_DEVICES=$HIP_VISIBLE_DEVICES + export CUDA_VISIBLE_DEVICES=$HIP_VISIBLE_DEVICES + +If you encounter any issues in using AMD GPUs running VeRL, feel free to contact me - `Yusheng Su `_. \ No newline at end of file diff --git a/verl/docs/examples/megatron_fsdp_example.rst b/verl/docs/examples/megatron_fsdp_example.rst new file mode 100644 index 0000000000000000000000000000000000000000..6b2cad4b5841bb5ad133aa00c2e70c287bb00fe5 --- /dev/null +++ b/verl/docs/examples/megatron_fsdp_example.rst @@ -0,0 +1,73 @@ +Megatron-FSDP Example +======================== + +Last updated: 04/29/2026. + +Introduction +------------ + +In this example, we run SFT and RL training with Megatron-FSDP: + +- Runtime image: ``verlai/verl:vllm011.dev7`` + +Step 1: Prepare +-------------------- + +Download ``Megatron-LM`` and ``Megatron-Bridge``. The required Megatron-FSDP support has already been merged into + ``Megatron-LM`` main + (``) and + ``Megatron-Bridge`` main + (``). + +.. code:: bash + + git clone https://github.com/NVIDIA/Megatron-LM.git + git clone https://github.com/NVIDIA-NeMo/Megatron-Bridge.git + +Step 2: Run Megatron-FSDP SFT +---------------------------- + +Before launch, check and update key fields ``MODEL_PATH`` and ``SAVE_PATH`` in the script. + +.. code:: bash + + bash examples/sft/gsm8k/run_qwen_megatron_fsdp.sh + +Step 3: Run Megatron-FSDP RL +---------------------------- + +Before launch, check and update key fields in +``examples/grpo_trainer/run_qwen2-7b_math_megatron_fsdp.sh``: + +- ``actor_rollout_ref.model.path``: model name or local model path. +- ``train_files`` / ``test_files``: parquet paths for GSM8K and MATH. +- ``trainer.n_gpus_per_node`` and ``trainer.nnodes``: hardware topology. +- ``trainer.project_name`` and ``trainer.experiment_name``: experiment identifiers. + +Then run: + +.. code:: bash + + bash examples/grpo_trainer/run_qwen2-7b_math_megatron_fsdp.sh + +The script launches RL training and enables Megatron-FSDP with: + +- ``actor_rollout_ref.actor.megatron.use_mbridge=True`` +- ``actor_rollout_ref.actor.megatron.vanilla_mbridge=False`` +- ``actor_rollout_ref.actor.megatron.use_megatron_fsdp=True`` + +Checkpoint Notes +---------------- + +Megatron-FSDP checkpoints are saved as DTensor checkpoints under ``dist_ckpt``. +When ``checkpoint.save_contents`` includes ``model``, verl also saves the HuggingFace config and +tokenizer under ``huggingface``; HF weights can also be exported through Megatron-Bridge. + +Current Megatron-FSDP checkpoint examples assume: + +- ``use_distributed_optimizer=True``. +- ``CUDA_DEVICE_MAX_CONNECTIONS`` is unset or greater than ``1``. +- PEFT + Megatron-FSDP checkpoint save/load is not covered by this example yet. +- ``checkpoint.async_save=True`` is not covered for Megatron-FSDP DTensor checkpoints yet. +- Megatron-FSDP checkpoints do not support saving optimizer state by itself; include ``model`` whenever + ``optimizer`` is listed in ``checkpoint.save_contents``. diff --git a/verl/docs/examples/multi_modal_example.rst b/verl/docs/examples/multi_modal_example.rst new file mode 100644 index 0000000000000000000000000000000000000000..c71b562b4614dae40ea2e204d998307b7db649c8 --- /dev/null +++ b/verl/docs/examples/multi_modal_example.rst @@ -0,0 +1,45 @@ +Multi-Modal Example Architecture +================================= + +Last updated: 04/28/2025. + +Introduction +------------ + +Now, verl has supported multi-modal training. You can use fsdp and +vllm/sglang to start a multi-modal RL task. Megatron supports is also +on the way. + +Follow the steps below to quickly start a multi-modal RL task. + +Step 1: Prepare dataset +----------------------- + +.. code:: python + + # it will be saved in the $HOME/data/geo3k folder + python examples/data_preprocess/geo3k.py + +Step 2: Download Model +---------------------- + +.. code:: bash + + # download the model from huggingface + python3 -c "import transformers; transformers.pipeline(model='Qwen/Qwen2.5-VL-7B-Instruct')" + +Step 3: Perform GRPO training with multi-modal model on Geo3K Dataset +--------------------------------------------------------------------- + +.. code:: bash + + # run the task + bash examples/grpo_trainer/run_qwen2_5_vl_7b_fsdp.sh + + + + + + + + diff --git a/verl/docs/examples/ppo_code_architecture.rst b/verl/docs/examples/ppo_code_architecture.rst new file mode 100644 index 0000000000000000000000000000000000000000..aaa91a7e7cf834fe30243f302b8759cafd363aa8 --- /dev/null +++ b/verl/docs/examples/ppo_code_architecture.rst @@ -0,0 +1,206 @@ +PPO Example Architecture +======================== + +Last updated: 02/17/2025. + +Let's start with the Proximal Policy Optimization algorithm, which is +most widely used algorithm in LLM post-training. + +The main entry point of the PPO algorithm example is: +`main_ppo.py `_. +In this tutorial, we will go through the code architecture in `main_ppo.py `_. + +Define the data +--------------- + +Users need to preprocess and store the dataset in parquet files. +And we implement `RLHFDataset` to load and tokenize the parquet files. + +For ``RLHFDataset`` (Default), at least 1 fields are required: + +- ``prompt``: Contains the string prompt + +We already provide some examples of processing the datasets to parquet +files in `data_preprocess directory `_. Currently, we support +preprocess of GSM8k, MATH, HellaSwag, Full_hh_rlhf datasets. See :doc:`../preparation/prepare_data` for +more information. + +Define the reward functions for different datasets +-------------------------------------------------- + +In this main entry point, the users only need to define their own reward +function based on the datasets (or applications) utilized in PPO +training. + +For example, we already provide reward functions for `GSM8k `_ +and `MATH `_ +datasets in the ``_select_rm_score_fn``. In the ``RewardManager``, we +will compute the reward score based on the data_source to select +corresponding reward functions. For some RLHF datasets (e.g., +full_hh_rlhf), the reward model is utilized to assess the responses +without any reward functions. In this case, the ``RewardManager`` will +return the ``rm_score`` computed by the reward model directly. + +See `reward functions `_ for detailed implementation. + +Define worker classes +--------------------- + +verl ships a single, unified model-engine worker implementation. The actor/rollout/ref +policy live in :class:`verl.workers.engine_workers.ActorRolloutRefWorker`, and the +critic/reward-model live in :class:`verl.workers.engine_workers.TrainingWorker`. +The underlying backend (FSDP, FSDP2, Megatron-LM, torchtitan, veomni, ...) is selected +at runtime from ``config.actor_rollout_ref.actor.strategy`` / ``config.critic.strategy``. + +.. code:: python + + from verl.single_controller.ray import RayWorkerGroup + from verl.trainer.ppo.ray_trainer import ResourcePoolManager, Role + from verl.workers.engine_workers import ActorRolloutRefWorker, TrainingWorker + + ray_worker_group_cls = RayWorkerGroup + + role_worker_mapping = { + Role.ActorRollout: ActorRolloutRefWorker, + Role.Critic: TrainingWorker, + Role.RefPolicy: ActorRolloutRefWorker + } + + global_pool_id = 'global_pool' + resource_pool_spec = { + global_pool_id: [config.trainer.n_gpus_per_node] * config.trainer.nnodes, + } + mapping = { + Role.ActorRollout: global_pool_id, + Role.Critic: global_pool_id, + Role.RefPolicy: global_pool_id, + } + +Step 1: Construct the mapping between roles and workers +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +A role represents a group of workers in the same process. We have +pre-defined several roles in `ray_trainer.py `_. + +.. code:: python + + class Role(Enum): + """ + To create more roles dynamically, you can subclass Role and add new members + """ + Actor = 0 # This worker only has Actor + Rollout = 1 # This worker only has Rollout + ActorRollout = 2 # This worker has both actor and rollout, it's a HybridEngine + Critic = 3 # This worker only has critic + RefPolicy = 4 # This worker only has reference policy + RewardModel = 5 # This worker only has reward model + ActorRolloutRef = 6 # This worker contains actor, rollout and reference policy simultaneously + +Step 2: Define the worker class corresponding to this role +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +- We have pre-implemented the ``ActorRolloutRefWorker``. Through + different configs, it can be a standalone actor, a standalone rollout, + an ActorRollout HybridEngine, or an ActorRolloutRef HybridEngine. +- The ``TrainingWorker`` is the generic training worker used for + ``Critic`` and ``Reward Model`` roles. +- Backend selection (PyTorch FSDP/FSDP2, Megatron-LM, torchtitan, veomni, ...) + is driven by ``config.actor_rollout_ref.actor.strategy`` and + ``config.critic.strategy`` and handled internally by the model engine. + See `engine workers `_ + and the `model engine package `_ + for more information. + +Step 3: Define resource pool id and resource pool spec +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +- Resource pool is a division of global GPU resources, + ``resource_pool_spec`` is a dict, mapping from id to # of GPUs + + - In the above example, we defined a global resource pool: + global_pool_id, and then put all roles on this one resource pool + with all the GPUs in this post-training task. This refers to + *co-locate* placement where all the models share the same set of + GPUs. + +- See resource pool and placement for advance usage. + +Defining reward model/function +------------------------------ + +.. code:: python + + # we should adopt a multi-source reward function here + # - for rule-based rm, we directly call a reward score + # - for model-based rm, we call a model + # - for code related prompt, we send to a sandbox if there are test cases + # - finally, we combine all the rewards together + # - The reward type depends on the tag of the data + if config.reward_model.enable: + from verl.workers.engine_workers import TrainingWorker + role_worker_mapping[Role.RewardModel] = TrainingWorker + mapping[Role.RewardModel] = global_pool_id + + reward_fn = RewardManager(tokenizer=tokenizer, num_examine=0) + + # Note that we always use function-based RM for validation + val_reward_fn = RewardManager(tokenizer=tokenizer, num_examine=1) + + resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping) + +Since not all tasks use model-based RM, users need to define here +whether it's a model-based RM or a function-based RM + +- If it's a model-based RM, directly add the ``RewardModel`` role in the + resource mapping and add it to the resource pool mapping. + + - The default ``TrainingWorker`` handles reward models through the unified + model engine and supports the typical huggingface + ``AutoModelForSequenceClassification`` layout. For custom reward models + you can subclass :class:`verl.workers.engine_workers.TrainingWorker` + or build a dedicated worker on top of the `model engine package + `_. + +- If it's a function-based RM, the users are required to classified the + reward function for each datasets. + +.. code:: python + + def _select_rm_score_fn(data_source): + if data_source == 'openai/gsm8k': + return gsm8k.compute_score + elif data_source == 'lighteval/MATH': + return math.compute_score + else: + raise NotImplementedError + +See reward functions implemented in `directory `_ +for more information. + +Define, init and run the PPO Trainer +------------------------------------ + +.. code:: python + + trainer = RayPPOTrainer(config=config, + tokenizer=tokenizer, + role_worker_mapping=role_worker_mapping, + resource_pool_manager=resource_pool_manager, + ray_worker_group_cls=ray_worker_group_cls, + reward_fn=reward_fn, + val_reward_fn=val_reward_fn) + trainer.init_workers() + trainer.fit() + +- We first initialize the ``RayPPOTrainer`` with user config, tokenizer + and all the above worker mapping, resource pool, worker group and + reward functions +- We first call the ``trainer.init_workers()`` to initialize the models + on the allocated GPUs (in the resource pool) +- The actual PPO training will be executed in ``trainer.fit()`` + +verl can be easily extended to other RL algorithms by reusing the Ray +model workers, resource pool and reward functions. See :doc:`extension<../advance/dpo_extension>` for +more information. + +Details of the ``RayPPOTrainer`` is discussed in :doc:`Ray Trainer<../workers/ray_trainer>`. diff --git a/verl/docs/examples/sandbox_fusion_example.rst b/verl/docs/examples/sandbox_fusion_example.rst new file mode 100644 index 0000000000000000000000000000000000000000..a141efdfd1e90346de2f8a3f8d8a90e784a4e445 --- /dev/null +++ b/verl/docs/examples/sandbox_fusion_example.rst @@ -0,0 +1,52 @@ +Sandbox Fusion Example +============================ + +Last updated: 05/17/2026. + +Introduction +------------ + +Sandbox Fusion is a remote code sandbox service that provides a secure environment for running and evaluating code generated by Large Language Models (LLMs). This example demonstrates how to train an LLM and use Sandbox Fusion to verify generated code, enhancing both security and performance. + +By leveraging a remote code sandbox service with greater CPU resources for concurrent code verification, you can reduce the reward stage time by 10-30%, depending on the quality of the generated code. + +Step 1: Prepare the Dataset +--------------------------- + +We use the Eurus-2-RL-Data dataset for training. This dataset combines math and code questions, making it suitable for LLM training tasks. You can download it from HuggingFace: `Eurus-2-RL-Data Dataset `_. + +Step 2: Set Up the Sandbox Fusion Service +----------------------------------------- + +Sandbox Fusion is a remote code sandbox service designed to securely run and evaluate LLM-generated code. To use it: + +1. **Access Full Documentation**: For detailed setup instructions, refer to the `Sandbox Fusion Documentation `_. +2. **Deploy the Service**: Choose one of the following deployment methods: + + - **Local Deployment**: Follow the guide `here `_. + - **FaaS Instance (Volcengine)**: Create an instance using the `Volcengine Documentation `_. + +After deployment, you will receive an API endpoint in the format: ``https:///run_code``. + +Step 3: Configure the Training Script +------------------------------------- + +To integrate Sandbox Fusion into your training script, configure the following parameters: + +**Key Settings for Sandbox Fusion** + +- ``reward_model.sandbox_fusion.url=''``: Enable Sandbox Fusion by specifying the API endpoint (must end with ``/run_code``). +- ``reward_model.sandbox_fusion.max_concurrent=256``: Set the maximum number of concurrent API requests to the Sandbox Fusion service. +- ``reward_model.sandbox_fusion.memory_limit_mb=1024``: Set the memory limit (in MB) for each sandbox instance. Defaults to 1024MB if not specified. + +**Additional Optimization** + +To further reduce code verification time, enable parallel processing with: + +- ``reward_model.reward_manager=prime``: The Prime reward manager verifies code across multiple subprocesses concurrently. + +**Example Notebook** + +For a practical implementation, refer to the example notebook: + +``examples/tutorial/agent_loop_get_started/agent_loop_tutorial.ipynb`` diff --git a/verl/docs/examples/skypilot_examples.rst b/verl/docs/examples/skypilot_examples.rst new file mode 100644 index 0000000000000000000000000000000000000000..3834f38842f4f0ff3f766160027c967de603133c --- /dev/null +++ b/verl/docs/examples/skypilot_examples.rst @@ -0,0 +1,146 @@ +SkyPilot Examples +================= + +Last updated: 09/04/2025. + +This guide provides examples of running VERL reinforcement learning training on Kubernetes clusters or cloud platforms with GPU nodes using `SkyPilot `_. + +Installation and Configuration +------------------------------- + +Step 1: Install SkyPilot +~~~~~~~~~~~~~~~~~~~~~~~~~ + +Choose the installation based on your target platform: + +.. code-block:: bash + + # For Kubernetes only + pip install "skypilot[kubernetes]" + + # For AWS + pip install "skypilot[aws]" + + # For Google Cloud Platform + pip install "skypilot[gcp]" + + # For Azure + pip install "skypilot[azure]" + + # For multiple platforms + pip install "skypilot[kubernetes,aws,gcp,azure]" + +Step 2: Configure Your Platform +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +See https://docs.skypilot.co/en/latest/getting-started/installation.html + +Step 3: Set Up Environment Variables +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Export necessary API keys for experiment tracking: + +.. code-block:: bash + + # For Weights & Biases tracking + export WANDB_API_KEY="your-wandb-api-key" + + # For HuggingFace gated models (if needed) + export HF_TOKEN="your-huggingface-token" + +Examples +-------- + +All example configurations are available in the `examples/tutorial/skypilot/ `_ directory on GitHub. See the `README `_ for additional details. + +PPO Training +~~~~~~~~~~~~ + +.. code-block:: bash + + sky launch -c verl-ppo verl-ppo.yaml --secret WANDB_API_KEY -y + +Runs PPO training on GSM8K dataset using Qwen2.5-0.5B-Instruct model across 2 nodes with H100 GPUs. Based on examples in ``examples/ppo_trainer/``. + +`View verl-ppo.yaml on GitHub `_ + +GRPO Training +~~~~~~~~~~~~~ + +.. code-block:: bash + + sky launch -c verl-grpo verl-grpo.yaml --secret WANDB_API_KEY -y + +Runs GRPO (Group Relative Policy Optimization) training on MATH dataset using Qwen2.5-7B-Instruct model. Memory-optimized configuration for 2 nodes. Based on examples in ``examples/grpo_trainer/``. + +`View verl-grpo.yaml on GitHub `_ + +Multi-turn Tool Usage Training +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. code-block:: bash + + sky launch -c verl-multiturn verl-multiturn-tools.yaml \ + --secret WANDB_API_KEY --secret HF_TOKEN -y + +Single-node training with 8xH100 GPUs for multi-turn tool usage with Qwen2.5-3B-Instruct. Includes tool and interaction configurations for GSM8K. Based on examples in ``examples/sglang_multiturn/`` but uses vLLM instead of sglang. + +`View verl-multiturn-tools.yaml on GitHub `_ + +Configuration +------------- + +The example YAML files are pre-configured with: + +- **Infrastructure**: Kubernetes clusters (``infra: k8s``) - can be changed to ``infra: aws`` or ``infra: gcp``, etc. +- **Docker Image**: VERL's official Docker image with CUDA 12.6 support +- **Setup**: Automatically clones and installs VERL from source +- **Datasets**: Downloads required datasets during setup phase +- **Ray Cluster**: Configures distributed training across nodes +- **Logging**: Supports Weights & Biases via ``--secret WANDB_API_KEY`` +- **Models**: Supports gated HuggingFace models via ``--secret HF_TOKEN`` + +Launch Command Options +---------------------- + +- ``-c ``: Cluster name for managing the job +- ``--secret KEY``: Pass secrets for API keys (can be used multiple times) +- ``-y``: Skip confirmation prompt + +Monitoring Your Jobs +-------------------- + +Check Cluster Status +~~~~~~~~~~~~~~~~~~~~ + +.. code-block:: bash + + sky status + +View Logs +~~~~~~~~~ + +.. code-block:: bash + + sky logs verl-ppo # View logs for the PPO job + +SSH into Head Node +~~~~~~~~~~~~~~~~~~ + +.. code-block:: bash + + ssh verl-ppo + +Access Ray Dashboard +~~~~~~~~~~~~~~~~~~~~ + +.. code-block:: bash + + sky status --endpoint 8265 verl-ppo # Get dashboard URL + +Stop a Cluster +~~~~~~~~~~~~~~ + +.. code-block:: bash + + sky down verl-ppo diff --git a/verl/docs/faq/faq.rst b/verl/docs/faq/faq.rst new file mode 100644 index 0000000000000000000000000000000000000000..901d8077045371551b1b5ea01fdf2211047f9e4b --- /dev/null +++ b/verl/docs/faq/faq.rst @@ -0,0 +1,209 @@ +Frequently Asked Questions +==================================== + +Last updated: 09/24/2025. + +Ray related +------------ + +How to add breakpoint for debugging with distributed Ray? +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Please checkout the official debugging guide from Ray: https://docs.ray.io/en/latest/ray-observability/ray-distributed-debugger.html + + +"Unable to register worker with raylet" +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +The cause of this issue is due to some system setting, e.g., SLURM added some constraints on how the CPUs are shared on a node. +While `ray.init()` tries to launch as many worker processes as the number of CPU cores of the machine, +some constraints of SLURM restricts the `core-workers` seeing the `raylet` process, leading to the problem. + +To fix this issue, you can set the config term ``ray_init.num_cpus`` to a number allowed by your system. + +Distributed training +------------------------ + +How to run multi-node post-training with Ray? +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +You can start a ray cluster and submit a ray job, following the official guide from Ray: https://docs.ray.io/en/latest/ray-core/starting-ray.html + +Then in the configuration, set the ``trainer.nnode`` config to the number of machines for your job. + +How to use verl on a Slurm-managed cluster? +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Ray provides users with `this `_ official +tutorial to start a Ray cluster on top of Slurm. We have verified the :doc:`GSM8K example<../examples/gsm8k_example>` +on a Slurm cluster under a multi-node setting with the following steps. + +1. [Optional] If your cluster support `Apptainer or Singularity `_ and you wish +to use it, convert verl's Docker image to an Apptainer image. Alternatively, set up the environment with the package +manager available on your cluster or use other container runtimes (e.g. through `Slurm's OCI support `_) available to you. + +.. code:: bash + + apptainer pull /your/dest/dir/vemlp-th2.4.0-cu124-vllm0.6.3-ray2.10-te1.7-v0.0.3.sif docker://verlai/verl:vemlp-th2.4.0-cu124-vllm0.6.3-ray2.10-te1.7-v0.0.3 + +2. Follow :doc:`GSM8K example<../examples/gsm8k_example>` to prepare the dataset and model checkpoints. + +3. Modify `examples/tutorial/slurm/ray_on_slurm.slurm `_ with your cluster's own information. + +4. Submit the job script to the Slurm cluster with `sbatch`. + +Please note that Slurm cluster setup may vary. If you encounter any issues, please refer to Ray's +`Slurm user guide `_ for common caveats. + +If you changed Slurm resource specifications, please make sure to update the environment variables in the job script if necessary. + + +Install related +------------------------ + +NotImplementedError: TensorDict does not support membership checks with the `in` keyword. +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Detail error information: + +.. code:: bash + + NotImplementedError: TensorDict does not support membership checks with the `in` keyword. If you want to check if a particular key is in your TensorDict, please use `key in tensordict.keys()` instead. + +Cause of the problem: There is no suitable version of tensordict package for the linux-arm64 platform. The confirmation method is as follows: + +.. code:: bash + + pip install tensordict==0.6.2 + +Output example: + +.. code:: bash + + ERROR: Could not find a version that satisfies the requirement tensordict==0.6.2 (from versions: 0.0.1a0, 0.0.1b0, 0.0.1rc0, 0.0.2a0, 0.0.2b0, 0.0.3, 0.1.0, 0.1.1, 0.1.2, 0.8.0, 0.8.1, 0.8.2, 0.8.3) + ERROR: No matching distribution found for tensordict==0.6.2 + +Solution 1st: + Install tensordict from source code: + +.. code:: bash + + pip uninstall tensordict + git clone https://github.com/pytorch/tensordict.git + cd tensordict/ + git checkout v0.6.2 + python setup.py develop + pip install -v -e . + +Solution 2nd: + Temperally modify the error takeplace codes: tensordict_var -> tensordict_var.keys() + + +Illegal memory access +--------------------------------- + +If you encounter the error message like ``CUDA error: an illegal memory access was encountered`` during rollout, please check the vLLM documentation for troubleshooting steps specific to your vLLM version. + +Checkpoints +------------------------ + +If you want to convert the model checkpoint into huggingface safetensor format, please refer to ``verl/model_merger``. + + +Triton ``compile_module_from_src`` error +------------------------------------------------ + +If you encounter triton compilation error similar to the stacktrace below, please set the ``use_torch_compile`` flag according to +https://verl.readthedocs.io/en/latest/examples/config.html to disable just-in-time compilation for fused kernels. + +.. code:: bash + + File "/data/lbh/conda_envs/verl/lib/python3.10/site-packages/triton/runtime/jit.py", line 345, in + return lambda *args, **kwargs: self.run(grid=grid, warmup=False, *args, **kwargs) + File "/data/lbh/conda_envs/verl/lib/python3.10/site-packages/triton/runtime/autotuner.py", line 338, in run + return self.fn.run(*args, **kwargs) + File "/data/lbh/conda_envs/verl/lib/python3.10/site-packages/triton/runtime/jit.py", line 607, in run + device = driver.active.get_current_device() + File "/data/lbh/conda_envs/verl/lib/python3.10/site-packages/triton/runtime/driver.py", line 23, in __getattr__ + self._initialize_obj() + File "/data/lbh/conda_envs/verl/lib/python3.10/site-packages/triton/runtime/driver.py", line 20, in _initialize_obj + self._obj = self._init_fn() + File "/data/lbh/conda_envs/verl/lib/python3.10/site-packages/triton/runtime/driver.py", line 9, in _create_driver + return actives[0]() + File "/data/lbh/conda_envs/verl/lib/python3.10/site-packages/triton/backends/nvidia/driver.py", line 371, in __init__ + self.utils = CudaUtils() # TODO: make static + File "/data/lbh/conda_envs/verl/lib/python3.10/site-packages/triton/backends/nvidia/driver.py", line 80, in __init__ + mod = compile_module_from_src(Path(os.path.join(dirname, "driver.c")).read_text(), "cuda_utils") + File "/data/lbh/conda_envs/verl/lib/python3.10/site-packages/triton/backends/nvidia/driver.py", line 57, in compile_module_from_src + so = _build(name, src_path, tmpdir, library_dirs(), include_dir, libraries) + File "/data/lbh/conda_envs/verl/lib/python3.10/site-packages/triton/runtime/build.py", line 48, in _build + ret = subprocess.check_call(cc_cmd) + File "/data/lbh/conda_envs/verl/lib/python3.10/subprocess.py", line 369, in check_call + raise CalledProcessError(retcode, cmd) + +What is the meaning of train batch size, mini batch size, and micro batch size? +------------------------------------------------------------------------------------------ + +This figure illustrates the relationship between different batch size configurations. + +https://excalidraw.com/#json=pfhkRmiLm1jnnRli9VFhb,Ut4E8peALlgAUpr7E5pPCA + +.. image:: https://github.com/user-attachments/assets/16aebad1-0da6-4eb3-806d-54a74e712c2d + +How to generate ray timeline to analyse performance of a training job? +------------------------------------------------------------------------------------------ + +To generate the ray timeline file, you can set the config term ``ray_init.timeline_json_file`` to a json file path. +For example: + +.. code:: bash + + ray_init.timeline_json_file=/tmp/ray_timeline.json + +The file will be generated in the specified path at the end of a training job. +You can use tools like chrome://tracing or the Perfetto UI and view the ray timeline file. + +This figure shows the ray timeline file generated by from a training job on 1 node with 4 GPUs + +.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray_timeline.png?raw=true + +How to set proxy only for wandb? +------------------------------------------------------------------------------------------ + +If you need a proxy to access wandb, you can add below config in your training job script. +Comparing to using global https_proxy env variable, this approach won't mess up other http requests, such as ChatCompletionScheduler. + +.. code:: bash + + +trainer.wandb_proxy=http:// + +Missmatch between inference and training sequence (high actor/grad_norm) +------------------------------------------------------------------------------------------ + +If you encounter the issue of actor/grad_norm metric continuously increasing during training, it might be caused by a significant precision mismatching between the inference engine and training. You can use the following parameter to confirm this: + +.. code:: bash + + actor_rollout_ref.rollout.calculate_log_probs=True + +This parameter will add metrics like training/rollout_probs_diff_mean , which can be used to verify if there is a precision difference between inference and training. + +Under normal circumstances, the value of training/rollout_probs_diff_mean should be below 0.005. If you observe this value to be higher than 0.01, it indicates a precision issue from the inference engine. +The precision issue is known to occur under the following conditions: + +1. Using non-Hopper architecture GPUs, such as A100, L20, B200, etc. + +2. Using vLLM `with issue 22103 `_ as the inference engine. + +3. The input and output texts are long, for example, in multi-turn scenarios using reasioning models like Qwen3 for RL training. + +If all three conditions above are met and you observe that rollout_probs_diff_mean is too high, it is recommended to add the following parameter to resolve the precision issue: + +.. code:: bash + + +actor_rollout_ref.rollout.engine_kwargs.vllm.disable_cascade_attn=True + +The root cause of this issue is a bug in the flash attention used by vLLM. Although it has been fixed, the fix has not yet been released in the latest version of vLLM (v0.10.2). +For a more detailed explanation of this issue, please refer to `Fix LSE output error in FA2 kv-split `_. + +Until vLLM releases a new version with this fix, it is recommended to use the configuration above to disable cascade attention as a workaround. diff --git a/verl/docs/hybrid_flow.rst b/verl/docs/hybrid_flow.rst new file mode 100644 index 0000000000000000000000000000000000000000..1bd1918eba9a3aafd069acdb21c64bc400199b1b --- /dev/null +++ b/verl/docs/hybrid_flow.rst @@ -0,0 +1,259 @@ +========================================================= +HybridFlow Programming Guide +========================================================= + +Last updated: 06/02/2025. + +.. _vermouth: https://github.com/vermouth1992 + +Author: `Chi Zhang `_ + +verl is an open source implementation of the paper `HybridFlow `_ [1]_. In this section, we will introduce the basic concepts of HybridFlow, the motivation and how to program with verl APIs. + +Motivation and Design +------------------------ +We use dataflow to represent RL systems. [4]_. + +DataFlow +~~~~~~~~~~~~~~~~~~~~ + +Dataflow is an abstraction of computations. Neural Network training is a typical dataflow. It can be represented by computational graph. + +.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/dataflow.jpeg?raw=true + :alt: The dataflow graph from CS231n 2024 lecture 4 + +This figure [2]_ represents the computation graph of a polynomial function followed by a sigmoid function. In the data flow of neural network computation, each node represents an operator, and each edge represents the direction of forward/backward propagation. The computation graph determines the architecture of the neural network. + +RL as a dataflow problem +++++++++++++++++++++++++++++++++++++++++++++++ + +Reinforcement learning (RL) training can also be represented as a dataflow. Below is the dataflow graph that represents the PPO algorithm used in RLHF [3]_: + +.. image:: https://picx.zhimg.com/70/v2-cb8ab5ee946a105aab6a563e92682ffa_1440w.avis?source=172ae18b&biz_tag=Post + :alt: PPO dataflow graph, credit to Zhihu 低级炼丹师 + +However, the dataflow of RL has fundamental differences compared with dataflow of neural network training as follows: + ++--------------------------+--------------------------------------------------+---------------------+ +| Workload | Node | Edge | ++--------------------------+--------------------------------------------------+---------------------+ +| Neural Network Training | Operator (+/-/matmul/softmax) | Tensor movement | ++--------------------------+--------------------------------------------------+---------------------+ +| Reinforcement Learning | High-level operators (rollout/model forward) | Data Movement | ++--------------------------+--------------------------------------------------+---------------------+ + +In the case of tabular reinforcement learning, each operator is a simple scalar math operation (e.g., bellman update). In deep reinforcement learning(DRL), each operator is a high-level neural network computation such as model inference/update. This makes RL a two-level dataflow problem: + +- Control flow: defines how the high-level operators are executed (e.g., In PPO, we first perform rollout. Then, we perform advantage computation. Finally, we perform training). It expresses the **core logics of RL algorithms**. +- Computation flow: defines the dataflow of **neural network computation** (e.g., model forward/backward/optimizer). + + +Design Choices +~~~~~~~~~~~~~~~~~~~~ +The model size used in DRL before the LLM era is typically small. Thus, the high-level neural network computation can be done in a single process. This enables embedding the computation flow inside the control flow as a single process. + +However, in the LLM era, the computation flow (e.g., training neural network) becomes a multi-process program. This naturally leads to two design choices: + +1. Convert the control flow into a multi-process program as well. Then colocate with computation flow (unified multi-controller) + +- Advantages: + + - Achieves the **optimal performance** under fixed computation flow and control flow as the communication overhead in both training and data transfer is minimized. + +- Disadvantages: + + - The computation and/or control flow is **hard to reuse** from software perspective as computation code is coupled with specific controller code. For example, the training loop of PPO is generic. Say we have an PPO training flow implemented with a specific computation flow such as FSDP. Neither the control flow or computation flow can be reused if we want to switch the computation flow from FSDP to Megatron, due to the coupling of control and computation flows. + - Requires more efforts from the user under flexible and dynamic control flows, due to the multi-process nature of the program. + +2. Separate the flows: single process for the control flow and multi-process for computation flow + +- Advantages: + + - The computation flow defined elsewhere can be **easily reused** after the decoupling. + - The controller runs on a single process. Implementing a new RL algorithm with a **different control flow is simple and easy**. + +- Disadvantages: + + - Additional **data communication overhead** each time the controller process and computatation processes interact. The data has to be sent back and forth. + +In verl, the latter strategy with separate control flow and computation flow is adopted. verl is designed to decouple the control flow of RL algorithms, and the implementation of computation engines. + +Overall Execution Diagram +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Below is a simplified diagram denoting the execution of a reinforcement learning job. In the diagram, the controller runs on a single process, while the generator/actor workers, critic workers run on multiple processes, placed with specific resource groups. For rollout, the controller passes the data to the generator to perform sample generation. When the rollout is done, the data is passed back to controller for the next step of the algorithm. Similar execution is done for other workers. With the hybrid controller design, the data flow and computation is decoupled to provide both efficiency in computation and flexibility in defining algorithm training loops. + +.. figure:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/driver_worker.png?raw=true + :alt: The execution diagram + +Codebase walkthrough (PPO) +------------------------------------------------ + +Entry function +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +Code: https://github.com/verl-project/verl/blob/main/verl/trainer/main_ppo.py + +In this file, we define a remote function `main_task` that serves as the controller (driver) process as shown in the above figure. We also define a ``RewardManager``, where users can customize their reward function based on the data source in the dataset. Note that `RewardManager` should return the final token-level reward that is optimized by RL algorithms. Note that users can combine model-based rewards and rule-based rewards. +The ``main_task`` constructs a RayPPOTrainer instance and launch the fit. Note that ``main_task`` **runs as a single process**. + +We highly recommend that the ``main_task`` is NOT scheduled on the head of the ray cluster because ``main_task`` will consume a lot of memory but the head usually contains very few resources. + +Ray trainer +~~~~~~~~~~~~~~~~~~~~ +Code: https://github.com/verl-project/verl/blob/main/verl/trainer/ppo/ray_trainer.py + +The RayPPOTrainer manages + +- Worker and WorkerGroup construction +- Runs the main loop of PPO algorithm + +Note that, the fit function of RayPPOTrainer **runs as a single process**. + +Worker and WorkerGroup construction +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Each workerGroup manages a list of workers that runs remotely. Note that the worker group runs in the process of its constructor. +Each worker inside the WorkerGroup runs on a GPU. The worker group serves as a proxy for the controller process to interact with a list of workers, in order to perform certain computations. **In order to do so, we have to bind the methods of the worker into the method of the WorkerGroup and define the data dispatch and data collection**. This is done via simple decoration that will be introduced in the Worker definition section. + +For example, in PPO, we define 3 worker groups: + +- ActorRolloutRef: manages actor, rollout and reference policy. ActorRolloutRefWorker can be instantiated as a single actor, a single rollout, a single reference policy, a combined actor/rollout or a combined actor/rollout/ref. This design is aimed for the maximum code reuse in various scenarios. The reason for colocating actor and rollout is for fast weight transfer using nccl. The reason for coloating actor and reference is to implement an efficient lora PPO as the reference policy is simply the base model of PPO in lora. The colocation is done via ``verl.single_controller.ray.base.create_colocated_worker_cls``, where it creates a single ray remote class exposing all class methods from these roles. +- Critic: manages the critic model +- Reward: manages the reward model + +The worker group will be constructed on the resource pool it designates. The resource pool is a set of GPUs in the ray cluster. + +Worker definition +~~~~~~~~~~~~~~~~~~~~ + +.. _ActorRolloutRefWorker: https://github.com/verl-project/verl/blob/main/verl/workers/engine_workers.py + +We take `ActorRolloutRefWorker `_ for an example. +The APIs it should expose to the controller process are: + +- init_model: build the underlying model +- generate_sequences: given prompts, generate responses +- compute_log_prob: compute the log-probability of a generated sequence using actor +- compute_ref_log_prob: compute the log-probability of a generated sequence using reference policy +- save_checkpoint: save the checkpoint + +Note that these methods are defined in the worker that can only be invoked via remote calls. For example, if the controller process wants to initialize the model, it has to call + +.. code-block:: python + + for worker in actor_rollout_ref_wg: + worker.init_model.remote() + +If the controller process wants to generate sequences, it has to call + +.. code-block:: python + + data = xxx + # split the data into dp chunks + data_dp_lst = data.split(dp_size) + output_dp_lst = [] + for i, worker in enumerate(actor_rollout_ref_wg): + output_future = worker.generate_sequences.remote(data_dp_lst[i]) + output_dp_lst.append(output_future) + output = torch.cat(ray.get(output_dp_lst), dim=0) + +We observe that controller process calling worker group methods in general can be divided into 3 parts: + +- Split the data into data parallel sizes +- Dispatch the corresponding data into each worker +- Collect and concatenate the data when the computation finishes + +In verl, we design a syntax sugar to encapsulate the 3 processes into a single call from the controller process. + +.. code-block:: python + + @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO) + def generate_sequences(data): + ... + + # on the driver + output = actor_rollout_ref_wg.generate_sequences(data) + +We decorate the method of the worker with a ``register`` that explicitly defines how the input data should be split and dispatched to each worker, and how the output data should be collected and concatenated by the controller. For example, ``Dispatch.DP_COMPUTE_PROTO`` splits the input data into dp chunks, dispatch each data to each worker, collect the output and concatenate the results. Note that this function requires the input and output to be a DataProto defined here (https://github.com/verl-project/verl/blob/main/verl/protocol.py). + + +PPO main loop +~~~~~~~~~~~~~~~~~~~~ +With the aforementioned APIs, we can implement the main loop of PPO as if it is a single process program + +.. code-block:: python + + for prompt in dataloader: + output = actor_rollout_ref_wg.generate_sequences(prompt) + old_log_prob = actor_rollout_ref_wg.compute_log_prob(output) + ref_log_prob = actor_rollout_ref_wg.compute_ref_log_prob(output) + values = critic_wg.compute_values(output) + rewards = reward_wg.compute_scores(output) + # compute_advantages is running directly on the control process + advantages = compute_advantages(values, rewards) + output = output.union(old_log_prob) + output = output.union(ref_log_prob) + output = output.union(values) + output = output.union(rewards) + output = output.union(advantages) + # update actor + actor_rollout_ref_wg.update_actor(output) + critic.update_critic(output) + +Takeaways +~~~~~~~~~~~~~~~~~~~~ +- This programming paradigm enables users to use different computation backend without modification of the control process. +- This programming paradigm enables flexible placement (by changing the mapping of WorkerGroup and ResourcePool) without modification of the control process. + +Repository organization +------------------------------------------------ + +Important code files in the repository are organized as below: + +.. code-block:: bash + + verl # the verl package + trainer + main_ppo.py # the entrypoint for RL training + ppo + ray_trainer.py # the training loop for RL algorithms such as PPO + sft_trainer.py # the SFT trainer with FSDP backend + config + generation.yaml # configuration template for rollout + ppo_trainer.yaml # configuration template for the RL trainer + workers + protocol.py # the interface of DataProto + engine_workers.py # unified ActorRolloutRefWorker / TrainingWorker implementations + engine + fsdp # FSDP / FSDP2 actor+critic engine implementations + megatron # Megatron-LM actor+critic engine implementations + torchtitan # torchtitan engine implementation + veomni # veomni engine implementation + mcore # shared mcore helpers + rollout + vllm + vllm_rollout.py # rollout with vllm backend + hf_rollout.py # rollout with huggingface TGI backend + sglang_rollout # rollout with SGLang backend + utils + dataset # datasets for SFT/RM/RL + reward_score # function based reward + gsm8k.py # reward function for gsm8k dataset + math.py # reward function for math dataset + seqlen_balancing.py # the sequence balance optimization + models + llama # Megatron implementation for llama, deepseek, mistral, etc + transformers # ulysses integration with transformer models such as llama, qwen, etc + weight_loader_registery.py # registry of weight loaders for loading hf ckpt into Megatron + third_party + vllm # adaptor for vllm's usage in RL + vllm_spmd # vllm >= v0.7 adaptor + examples # example scripts + tests # integration and unit tests + .github # the configuration of continuous integration tests + + +.. [1] HybridFlow: A Flexible and Efficient RLHF Framework: https://arxiv.org/abs/2409.19256v2 +.. [2] Data flow graph credit to CS231n 2024 lecture 4: https://cs231n.stanford.edu/slides/2024/lecture_4.pdf +.. [3] PPO dataflow graph credit to 低级炼丹师 from Zhihu​: https://zhuanlan.zhihu.com/p/635757674 +.. [4] RLFlow diff --git a/verl/docs/index.rst b/verl/docs/index.rst new file mode 100644 index 0000000000000000000000000000000000000000..85a5da221554620b083f39d3533759f2dd321178 --- /dev/null +++ b/verl/docs/index.rst @@ -0,0 +1,250 @@ +Welcome to verl's documentation! +================================================ + +verl is a flexible, efficient and production-ready RL training framework designed for large language models (LLMs) post-training. It is an open source implementation of the `HybridFlow `_ paper. + +verl is flexible and easy to use with: + +- **Easy extension of diverse RL algorithms**: The hybrid programming model combines the strengths of single-controller and multi-controller paradigms to enable flexible representation and efficient execution of complex Post-Training dataflows. Allowing users to build RL dataflows in a few lines of code. + +- **Seamless integration of existing LLM infra with modular APIs**: Decouples computation and data dependencies, enabling seamless integration with existing LLM frameworks, such as PyTorch FSDP, Megatron-LM, vLLM and SGLang. Moreover, users can easily extend to other LLM training and inference frameworks. + +- **Flexible device mapping and parallelism**: Supports various placement of models onto different sets of GPUs for efficient resource utilization and scalability across different cluster sizes. + +- Ready integration with popular HuggingFace models + + +verl is fast with: + +- **State-of-the-art throughput**: By seamlessly integrating existing SOTA LLM training and inference frameworks, verl achieves high generation and training throughput. + +- **Efficient actor model resharding with 3D-HybridEngine**: Eliminates memory redundancy and significantly reduces communication overhead during transitions between training and generation phases. + +-------------------------------------------- + +.. _Contents: + +.. toctree:: + :maxdepth: 2 + :caption: Quickstart + + start/install + start/quickstart + start/multinode + start/ray_debug_tutorial + start/more_resources + start/agentic_rl + +.. toctree:: + :maxdepth: 2 + :caption: Programming guide + + hybrid_flow + single_controller + +.. toctree:: + :maxdepth: 1 + :caption: Data Preparation + + preparation/prepare_data + preparation/reward_function + +.. toctree:: + :maxdepth: 2 + :caption: Configurations + + examples/config + +.. toctree:: + :maxdepth: 1 + :caption: PPO Example + + examples/ppo_code_architecture + examples/gsm8k_example + examples/megatron_fsdp_example + examples/multi_modal_example + examples/skypilot_examples + +.. toctree:: + :maxdepth: 1 + :caption: Algorithms + + algo/ppo.md + algo/grpo.md + algo/dapo.md + algo/spin.md + algo/sppo.md + algo/entropy.md + algo/opo.md + algo/baseline.md + algo/gpg.md + algo/rollout_corr.md + algo/rollout_corr_math.md + algo/otb.md + algo/dppo.md + algo/opd.md + +.. toctree:: + :maxdepth: 1 + :caption: PPO Trainer and Workers + + workers/ray_trainer + workers/model_engine + workers/engine_workers + workers/automodel_workers + workers/sglang_worker + workers/trtllm_worker + +.. toctree:: + :maxdepth: 1 + :caption: Performance Tuning Guide + + perf/dpsk.md + perf/best_practices + perf/perf_tuning + README_vllm0.8.md + perf/device_tuning + perf/verl_profiler_system.md + perf/nsight_profiling.md + perf/torch_profiling.md + +.. toctree:: + :maxdepth: 1 + :caption: Adding new models + + advance/fsdp_extension + advance/megatron_extension + +.. toctree:: + :maxdepth: 1 + :caption: Async Training + + advance/one_step_off + advance/fully_async + advance/async-on-policy-distill + +.. toctree:: + :maxdepth: 1 + :caption: Low Precision + + low_precision/fp8.md + low_precision/nvfp4_qat.md + +.. toctree:: + :maxdepth: 1 + :caption: Advanced Features + + advance/checkpoint + advance/rope + advance/attention_implementation + advance/ppo_lora.rst + sglang_multiturn/multiturn.rst + advance/placement + advance/dpo_extension + examples/sandbox_fusion_example + advance/rollout_trace.rst + advance/rollout_skip.rst + advance/agent_loop + advance/reward_loop + data/transfer_queue.md + advance/grafana_prometheus.md + advance/mtp.md + +.. toctree:: + :maxdepth: 2 + :caption: Hardware Support + + amd_tutorial/amd_build_dockerfile_page.rst + amd_tutorial/amd_vllm_page.rst + amd_tutorial/amd_quick_start.rst + ascend_tutorial/README.md + ascend_tutorial/get_start/dockerfile_build_guidance.rst + ascend_tutorial/get_start/install_guidance.rst + ascend_tutorial/get_start/quick_start.rst + ascend_tutorial/feature_support/ascend_backend_features.md + ascend_tutorial/feature_support/npu_advance_features.md + ascend_tutorial/model_support/model_and_algorithm_support.md + ascend_tutorial/model_support/examples/ascend_retool_best_pratice.rst + ascend_tutorial/model_support/examples/ascend_sglang_best_practices.rst + ascend_tutorial/model_support/examples/dapo_multi_model_optimization_practice.md + ascend_tutorial/model_support/examples/gspo_optimization_practice.md + ascend_tutorial/dev_guide/model_dev/evaluation.md + ascend_tutorial/dev_guide/model_dev/parameter_and_metrics.md + ascend_tutorial/dev_guide/model_dev/transfer_to_npu_guide.md + ascend_tutorial/dev_guide/precision_analysis/precision_alignment_zh.md + ascend_tutorial/dev_guide/precision_analysis/precision_debugger_zh.md + ascend_tutorial/dev_guide/performance/ascend_performance_analysis_guide.md + ascend_tutorial/dev_guide/performance/perf_tuning_on_ascend.rst + ascend_tutorial/dev_guide/performance/ascend_profiling_zh.rst + ascend_tutorial/dev_guide/performance/ascend_profiling_en.rst + ascend_tutorial/faq/faq.rst + ascend_tutorial/contribution_guide/ascend_ci_guide_zh.rst + +.. toctree:: + :maxdepth: 1 + :caption: API References + + api/data + api/single_controller.rst + api/trainer.rst + api/utils.rst + +.. toctree:: + :maxdepth: 1 + :caption: Blog + + blog/v0.7.md + +.. toctree:: + :maxdepth: 2 + :caption: FAQ + + faq/faq + +.. toctree:: + :maxdepth: 1 + :caption: Contributing + + contributing/editing-agent-instructions.md + +.. toctree:: + :maxdepth: 1 + :caption: Development Notes + + sglang_multiturn/sandbox_fusion.rst + +Contribution +------------- + +verl is free software; you can redistribute it and/or modify it under the terms +of the Apache License 2.0. We welcome contributions. +Join us on `GitHub `_, `Slack `_ and `Wechat `_ for discussions. + +Contributions from the community are welcome! Please check out our `project roadmap `_ and `good first issues `_ to see where you can contribute. + +Code Linting and Formatting +^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +We use pre-commit to help improve code quality. To initialize pre-commit, run: + +.. code-block:: bash + + pip install pre-commit + pre-commit install + +To resolve CI errors locally, you can also manually run pre-commit by: + +.. code-block:: bash + + pre-commit run + +Adding CI tests +^^^^^^^^^^^^^^^^^^^^^^^^ + +If possible, please add CI test(s) for your new feature: + +1. Find the most relevant workflow yml file, which usually corresponds to a ``hydra`` default config (e.g. ``ppo_trainer``, ``ppo_megatron_trainer``, ``sft_trainer``, etc). +2. Add related path patterns to the ``paths`` section if not already included. +3. Minimize the workload of the test script(s) (see existing scripts for examples). + +We are HIRING! Send us an `email `_ if you are interested in internship/FTE opportunities in MLSys/LLM reasoning/multimodal alignment. diff --git a/verl/docs/low_precision/fp8.md b/verl/docs/low_precision/fp8.md new file mode 100644 index 0000000000000000000000000000000000000000..0b4290be044fca65af57405dacb7e9211f7a77c6 --- /dev/null +++ b/verl/docs/low_precision/fp8.md @@ -0,0 +1,192 @@ +# FP8 RL in verl + +Last updated: 03/05/2026 + +verl supports two FP8 modes for accelerating RL training: + +| Mode | Training Precision | Rollout Precision | +|------|-------------------|-------------------| +| **FP8 Rollout Only** | BF16 | FP8 | +| **FP8 End-to-End** | FP8 (Megatron) | FP8 (vLLM) | + +> [!TIP] +> For ready-to-run scripts, see the [low-precision recipe directory](https://github.com/verl-project/verl-recipe/low_precision). + +--- + +## FP8 Rollout Only + +FP8 rollout-only mode keeps training in BF16 and quantizes rollout inference to FP8. This reduces GPU memory during generation and speeds up rollout without affecting training precision. + +### Implementation + +We monkey patch several vLLM functions to enable FP8 rollout for reinforcement learning: + +1. **Quantize weights**: Quantize model weights on-the-fly from higher-precision formats to FP8. +2. **Process weights after loading**: For vLLM, we replace the `vllm.model_executor.layers.quantization.fp8.Fp8LinearMethod.process_weights_after_loading` function to handle weight processing after quantization. For SGLang, this patch is not needed as it natively supports loading quantized weights. + +### Support Matrix + +- FP8 blockwise quantization for rollout + - Used in Deepseek, which is 1x128 quantization for activations and 128x128 quantization for model weights +- Dense models and MoE models +- Async rollout interfaces +- vLLM 0.10.x & vLLM 0.11 & vLLM 0.12 & SGLang 0.5.5 +- FSDP and Megatron training backends + +### Usage + +Enable in config file: + +```yaml +rollout: + quantization: "fp8" +``` + +Or via command line: + +```bash +actor_rollout_ref.rollout.quantization=fp8 +``` + +### Experiments and Outcomes + +#### Qwen3-8B-Base Dense Model + +**Configuration** +- DAPO recipe. AIME24 online validation. +- vLLM(FP8 spmd rollout) + FSDP + - Note that SPMD rollout has been deprecated, so we removed the FP8 SPMD rollout. +- Prompt batch size 32, n=16. +- Rollout batch size: 32\*3*16 +- Train_batch_size & ppo_mini_batch_size 32 +- Max response length 20K +- Token-level TIS, C=2 +- 8*H100 +- vLLM 0.10.0+CUDA 12.6 vs vLLM 0.11.0+CUDA 12.9 + +**Accuracy** +![Qwen3-8b-base_fp8_acc]( +https://github.com/Agoniii/verl/blob/xueh/fp8_pr_images/docs/advance/images/Qwen3-8b-base_fp8_acc.png?raw=true) +*dark green: BF16, orange: FP8 rollout + token-level TIS, light green: FP8 rollout without TIS* + +Results and observations: +- With TIS, FP8 rollout aligns with BF16 +- Obvious accuracy drop when TIS is not enabled +- Higher mismatch kl but within acceptable range throughout the training + + +**Performance** + +![Qwen3-8b-base_fp8_rollout_perf]( +https://github.com/Agoniii/verl/blob/xueh/fp8_pr_images/docs/advance/images/Qwen3-8b-base_fp8_rollout_perf.png?raw=true) +*green: BF16, orange: FP8 rollout + CUDA12.6 + DeepGemm, purple: FP8 rollout + CUDA 12.9 + DeepGemm* + +Results and observations: +- FP8 rollout leads to around ~12% rollout speedup with CUDA 12.6 + DeepGemm +- When upgrading to CUDA 12.9, speedup can be up to ~18% + +#### Qwen3-30B-A3B-Base MoE Model + +**Configuration** +- DAPO recipe. AIME24 online validation. +- FP8 async rollout, vLLM+FSDP +- Prompt batch size 32 +- Rollout batch size: 32\*3*16 +- Train_batch_size & ppo_mini_batch_size 32 +- Max response length 20K +- Token-level TIS, C=2 +- 2\*8*H100 +- vLLM 0.10.0+CUDA 12.6 + +**Accuracy** +![Qwen3-30b-a3b_fp8_acc]( +https://github.com/Agoniii/verl/blob/xueh/fp8_pr_images/docs/advance/images/Qwen3-30b-a3b_fp8_acc.png?raw=true) +*grey: BF16 + token-level TIS, red: FP8 rollout + token-level TIS* + +Results and observations: +- Rollout & training distribution mismatch is in general higher for MoE +- Rollout correction required even for BF16 +- FP8 rollout with token-level TIS aligns with BF16 + + +**Performance** + +![Qwen3-30b-a3b_fp8_perf]( +https://github.com/Agoniii/verl/blob/xueh/fp8_pr_images/docs/advance/images/Qwen3-30b-a3b_fp8_perf.png?raw=true) +*grey: BF16 + token-level TIS, red: FP8 rollout + token-level TIS​* + +Results and observations: +- FP8 rollout : over 35% rollout speedup +- Expecting more perf gain with CUDA 12.9 + +--- + +## FP8 End-to-End (Training + Rollout) + +FP8 E2E applies FP8 to the entire RL pipeline: forward/backward passes via Transformer Engine, FP8 optimizer states, and FP8 rollout inference via vLLM. This maximizes memory savings and throughput. + +### Requirements + +- **CUDA 12.9+** (required for block-wise FP8 scaling) +- **Transformer Engine** with block-wise FP8 support +- Environment variable: `NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1` + +### Key Configuration + +```yaml +# FP8 training via Transformer Engine +actor_rollout_ref.actor.megatron.override_transformer_config: + fp8: "hybrid" # FP8 forward + backward; also supports "e4m3" + fp8_recipe: "blockwise" # block-wise scaling + +# FP8 optimizer +actor_rollout_ref.actor.optim.override_optimizer_config: + fp8_recipe: "blockwise" + +# FP8 rollout inference (vLLM) +actor_rollout_ref.rollout: + quantization: fp8 +``` + +### Support Matrix + +- Megatron training backend (via Megatron-Bridge) +- Verified on Qwen3-30B-A3B and Qwen3-8B +- Block-wise FP8 scaling (`fp8_recipe: "blockwise"`) + +### Experiments and Results + +#### Qwen3-30B-A3B MoE Model + +**Configuration** +- DAPO recipe. AIME24 online validation. +- Megatron + Megatron-Bridge, FP8 async rollout with vLLM +- MoE router in BF16 for both vLLM and Megatron-Core +- Prompt batch size 128, n=16 +- Max response length 20K +- Token-level TIS, C=2 +- 2\*8*H100, CUDA 12.9 + +![Qwen3-30b-a3b_fp8_e2e](https://github.com/user-attachments/assets/70fb1396-ec73-40d7-9a43-1d48553c0ad9) +*Orange: BF16, Green: FP8 E2E, Red: FP8 rollout + BF16 training* + +Results and observations: +- FP8 E2E achieves comparable accuracy to the BF16 baseline, with the two curves closely aligned throughout training. +- The training/inference precision mismatch (measured by KL divergence) follows the ordering: FP8 rollout-only > FP8 E2E > BF16 E2E. This is expected, as FP8 E2E maintains consistent precision across both training and inference, resulting in lower distribution mismatch than the FP8 rollout-only setting where training remains in BF16. + +--- + +## Citation + +For more extensive experiments, ablation studies, and analysis on FP8 reinforcement learning, please refer to our technical report: + +```bibtex +@article{qiu2026fp8rl, + title={FP8-RL: A Practical and Stable Low-Precision Stack for LLM Reinforcement Learning}, + author={Qiu, Zhaopeng and Yu, Shuang and Zhang, Jingqi and Zhang, Shuai and Huang, Xue and Yang, Jingyi and Lai, Junjie}, + journal={arXiv preprint arXiv:2601.18150}, + year={2026}, + url={https://arxiv.org/abs/2601.18150} +} +``` diff --git a/verl/docs/low_precision/nvfp4_qat.md b/verl/docs/low_precision/nvfp4_qat.md new file mode 100644 index 0000000000000000000000000000000000000000..f7c2f19a23e934d9471f2b2a44c013eb1e065539 --- /dev/null +++ b/verl/docs/low_precision/nvfp4_qat.md @@ -0,0 +1,87 @@ +# NVFP4 QAT (Quantization-Aware Training) in verl + +Last updated: 04/02/2026 + +verl supports NVFP4 Quantization-Aware Training (QAT), which applies fake quantization during training so the model learns to tolerate NVFP4 quantization error. At rollout time, weights are packed into real NVFP4 format for vLLM inference. This closes the precision gap between training and inference, preventing KL divergence explosion. + +| Training Backend | Training Precision | Rollout Precision | vLLM Quant Method | +|---|---|---|---| +| **FSDP** | BF16 + fake quantization | NVFP4 W4A16 | `compressed-tensors` | +| **Megatron** | BF16 + fake quantization | NVFP4 W4A16 | `modelopt` | + +> [!TIP] +> For ready-to-run scripts, environment setup, and experimental results, see the [QAT recipe](https://github.com/verl-project/verl-recipe/tree/main/qat). + +--- + +## Key Configuration + +### FSDP Backend + +Configured under `actor_rollout_ref.actor.fsdp_config.qat`: + +```yaml +actor_rollout_ref: + actor: + fsdp_config: + qat: + enable: true + mode: "w4a16" + group_size: 16 + ignore_patterns: + - "lm_head" + - "embed_tokens" + - "re:.*mlp.gate$" + quantization_config_path: "recipe/qat/config/nvfp4_w4a16.json" +``` + +| Parameter | Description | Default | +|-----------|-------------|---------| +| `fsdp_config.qat.enable` | Enable QAT | `False` | +| `fsdp_config.qat.mode` | Quantization mode | `"w4a16"` | +| `fsdp_config.qat.group_size` | Quantization group size | `16` | +| `fsdp_config.qat.ignore_patterns` | Layers to skip. Supports `re:` prefix for regex, otherwise substring match | `["lm_head", "embed_tokens", "re:.*mlp.gate$"]` | +| `fsdp_config.qat.quantization_config_path` | vLLM quantization config JSON path | Required | + +### Megatron Backend + +Configured under `actor_rollout_ref.actor.megatron.qat`: + +```yaml +actor_rollout_ref: + actor: + megatron: + qat: + enable: true + mode: "w4a16" + group_size: 16 + ignore_patterns: + - "lm_head" + - "*mlp.gate" + quantization_config_path: "recipe/qat/config/nvfp4_w4a16_megatron.json" +``` + +| Parameter | Description | Default | +|-----------|-------------|---------| +| `megatron.qat.enable` | Enable QAT | `False` | +| `megatron.qat.mode` | Quantization mode | `"w4a16"` | +| `megatron.qat.group_size` | Quantization group size | `16` | +| `megatron.qat.ignore_patterns` | Layers to skip. Uses `fnmatch` glob syntax | `["lm_head", "*mlp.gate"]` | +| `megatron.qat.quantization_config_path` | vLLM quantization config JSON path | Required | + +--- + +## Support Matrix + +- NVFP4 W4A16 (weight-only FP4 quantization) +- Dense models and MoE models +- FSDP and Megatron training backends +- Full quantization and FFN-only quantization strategies +- Verified on Qwen3-8B-Base and Qwen3-30B-A3B-Base + +--- + +## Notes + +- FSDP backend has scalability limitations for very large models. For large-scale training, use the Megatron backend. +- FSDP uses `re:` prefix regex for `ignore_patterns`, while Megatron uses `fnmatch` glob syntax. The two are not interchangeable. diff --git a/verl/docs/perf/best_practices.rst b/verl/docs/perf/best_practices.rst new file mode 100644 index 0000000000000000000000000000000000000000..66f3e3122f5b258db07cdb2146fb982324999678 --- /dev/null +++ b/verl/docs/perf/best_practices.rst @@ -0,0 +1,242 @@ +Verl LLM Best Practices (DAPO + Qwen3-235B) +=========================================== + +Last updated: 11/03/2025. + +Purpose +------- + +This guide uses DAPO training on Qwen3-235B as a concrete example. We unpack every parameter that appears in the optimization objective, map it to Verl configuration entries, and share field-tested recommendations so you can derive sensible settings for your own workloads. + +.. note:: + + 1. The guide only covers the subset of parameters required to reproduce the DAPO experiments discussed here. For the full list, refer to the ``config`` components in the Verl source tree: https://github.com/verl-project/verl/tree/main/verl/trainer/config + 2. PPO and GRPO introduce KL-constrained policies. We therefore include that setup in the explanations below. You can treat all configurations mentioned here as a DAPO pipeline augmented with a KL penalty. + +Optimization Objectives +----------------------- + +DAPO objective +~~~~~~~~~~~~~~ + +.. math:: + + \begin{aligned} + \mathcal{J}_{\mathrm{DAPO}}(\theta)= & \mathbb{E}_{(q, a) \sim \mathcal{D},\left\{o_i\right\}_{i=1}^G \sim \pi_{\theta_{\text {old }}}(\cdot \mid q)} \ + {\left[\frac{1}{\sum_{i=1}^G\left|o_i\right|} \sum_{i=1}^G \sum_{t=1}^{\left|o_i\right|} \min \left(r_{i, t}(\theta) \hat{A}_{i, t}, \operatorname{clip}\left(r_{i, t}(\theta), 1-\varepsilon_{\text {low }}, 1+\varepsilon_{\text {high }}\right) \hat{A}_{i, t}\right)\right] } \\ + \end{aligned} + +.. math:: + \text { s.t. } \quad 0<\mid\left\{o_i \mid \text { is_equivalent }\left(a, o_i\right)\right\} \mid 2 * model_parameters`` (bf16/fp16). Increase TP gradually to expand KV cache capacity while watching communication cost—especially once TP > 8. + - ``actor_rollout_ref.rollout.temperature`` / ``top_p`` / ``top_k``: + Sampling knobs for rollout. Keep enough randomness; ``temperature=1.0``, ``top_p=1.0``, ``top_k=-1`` are good defaults. + - ``actor_rollout_ref.rollout.val_kwargs.temperature`` / ``top_p`` / ``top_k`` / ``do_sample`` / ``n``: + Sampling options for validation. Set ``temperature > 0`` to prevent repetitive thinking chains. For small test sets (e.g., AIME24) raise ``n`` (64 is a common choice) to reduce variance. A practical starting point is ``temperature=1.0``, ``top_p=0.7``, ``top_k=-1``, ``do_sample=True``, ``n=1`` and then increase ``n`` as needed. + - ``+actor_rollout_ref.rollout.engine_kwargs.vllm.*`` / ``+actor_rollout_ref.rollout.engine_kwargs.sglang.*``: + Extra backend options injected via the ``+`` syntax. Consult backend docs for exact semantics. Some switches (for example ``pipeline_parallel_size``) may not be supported yet; when TP=32, ``enable_expert_parallel=True`` can even slow down DeepSeek-V3 rollout, so benchmark carefully. + +:math:`\pi_\theta` + - ``data.train_batch_size``: + Total batch size per training iteration. Each rollout produces ``train_batch_size * n`` samples. Larger values reduce the number of rollouts but increase off-policy drift. + - ``actor_rollout_ref.actor.ppo_mini_batch_size``: + Mini-batch size per optimization step. Tune it the same way you would for standard deep learning workloads. + - ``actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu``: + Samples processed per forward pass on one GPU group (a Megatron group contains TP * PP * CP GPUs). Keep it ≤ ``ppo_mini_batch_size`` and as large as memory allows. + - ``actor_rollout_ref.actor.use_dynamic_bsz``: + Enable dynamic batch sizing to adapt to sequence length and improve throughput. + - ``actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu``: + Maximum tokens per GPU when computing log probabilities under dynamic batching. Set it to at least a multiple of ``max_prompt_length + max_response_length`` to prevent truncation. + - Megatron parallelism parameters (``pipeline_model_parallel_size`` / ``tensor_model_parallel_size`` / ``expert_model_parallel_size`` / ``expert_tensor_parallel_size`` / ``context_parallel_size``): + Balance PP/TP/EP/ETP/CP to match memory and network constraints. In bf16/fp16, each parameter consumes roughly ``2 / TP`` bytes; if you keep FP32 master weights or skip optimizer offload, reserve another 4–8 bytes for Adam. Activations scale with ``micro_batch_size × sequence_length × hidden_size`` and can be mitigated with gradient checkpointing, dynamic batches, or offload. Prefer increasing TP first, add PP when necessary, extend sequence capacity with CP, align EP/ETP with TP for MoE models, and keep DP minimal on constrained clusters while combining with offload. Always align the setup with hardware topology and communication cost. + - ``actor_rollout_ref.model.use_fused_kernels``: + Enable Verl’s fused kernels for supported models to squeeze out additional performance. + +:math:`\hat{A}_{i,t}` + - ``algorithm.adv_estimator``: + Advantage estimator. Set to ``grpo`` for DAPO/GRPO. + +:math:`R_i` + - ``reward_model.reward_manager``: + Reward aggregation strategy. Use ``dapo`` for DAPO and ``naive`` for GRPO. + +:math:`D_{KL}` + - ``algorithm.use_kl_in_reward``: + Whether to add a KL term to the reward. ``True`` for PPO, ``False`` for GRPO and DAPO. + - ``actor_rollout_ref.actor.use_kl_loss``: + Whether to include a KL loss term. ``False`` for PPO, ``True`` for GRPO, ``False`` for DAPO. + +:math:`\beta` + - ``actor_rollout_ref.actor.kl_loss_coef``: + Weight of the KL loss. Start around 0.001. Larger values curb reward hacking but reduce exploration. + - ``algorithm.kl_ctrl.kl_coef``: + KL coefficient applied within the reward. Adjust to match your tolerance for divergence. + +:math:`\pi_{old}` + - ``actor_rollout_ref.rollout.log_prob_use_dynamic_bsz``: + Enable dynamic batching when the old policy computes log-probabilities. Recommended. + +:math:`\pi_{ref}` + - ``actor_rollout_ref.ref.log_prob_use_dynamic_bsz``: + Enable dynamic batching for the reference policy. Recommended. + - Reference Megatron parallelism: + Keep ``pipeline_model_parallel_size``, ``tensor_model_parallel_size``, ``expert_model_parallel_size``, ``expert_tensor_parallel_size``, and ``context_parallel_size`` in sync with the actor. + - ``actor_rollout_ref.ref.megatron.param_offload``: + Offload reference parameters to CPU when the actor does so. Even without gradients or optimizer states, parity helps with capacity planning. + +:math:`o_i` / :math:`|o_i|` + - ``actor_rollout_ref.actor.loss_agg_mode``: + Loss aggregation mode. Token-level ``token-mean`` matches the recommendations from Dr.GRPO and DAPO; use ``seq-mean-token-mean`` to reproduce the original GRPO behavior. + +:math:`\pi_\theta(o_{i,t} \mid q_i,o_{i,`_ + - `SimonHuang `_ + +1.5B +~~~ + +.. list-table:: + :widths: auto + :header-rows: 1 + + * - Tag + - Model + - Task + - Resource + - MaxBatch + - Train + - Infer + - Link + - Contributor + * - MIN + - Qwen2.5-1.5B + - GRPO-LoRA + - 1*H100 + - 128 + - fsdp + - vllm0.8.3 + - `qwen2-1.5b_grpo-lora_1_h100_fsdp_vllm.sh `_ + - `SimonHuang `_ + +3B +~~~ + +.. list-table:: + :widths: auto + :header-rows: 1 + + * - Tag + - Model + - Task + - Resource + - MaxBatch + - Train + - Infer + - Link + - Contributor + * - MIN + - Qwen2.5-3B + - GRPO-LoRA + - 1*H100 + - 62 + - fsdp + - vllm0.8.3 + - `qwen2-3b_grpo-lora_1_h100_fsdp_vllm.sh `_ + - `SimonHuang `_ + +7B +~~~ + +.. list-table:: + :widths: auto + :header-rows: 1 + + * - Tag + - Model + - Task + - Resource + - MaxBatch + - Train + - Infer + - Link + - Contributor + * - MIN + - Qwen2-7B + - GRPO + - 2*H800 + - \ + - fsdp + - vllm0.8.2 + - `qwen2-7b_grpo_2_h800_fsdp_vllm `_ + - `Xiangyongan `_ + * - MIN + - Qwen2.5-7B + - GRPO-LoRA + - 1*H100 + - 16 + - fsdp + - vllm0.8.3 + - `qwen2-7b_grpo-lora_1_h100_fsdp_vllm.sh `_ + - `SimonHuang `_ + +14B +~~~ + +.. list-table:: + :widths: auto + :header-rows: 1 + + * - Tag + - Model + - Task + - Resource + - MaxBatch + - Train + - Infer + - Link + - Contributor + * - MIN + - Qwen2-14B + - GRPO + - 4*H800 + - \ + - fsdp + - vllm0.8.2 + - `qwen2-14b_grpo_4_h800_fsdp_vllm `_ + - `Xiangyongan `_ + * - MIN + - Qwen2.5-14B + - GRPO-LoRA + - 2*H100 + - 116 + - fsdp + - vllm0.8.3 + - `qwen2-14b_grpo-lora_2_h100_fsdp_vllm.sh `_ + - `SimonHuang `_ + +32B +~~~ + +.. list-table:: + :widths: auto + :header-rows: 1 + + * - Tag + - Model + - Task + - Resource + - MaxBatch + - Train + - Infer + - Link + - Contributor + * - MIN + - Qwen2-32B + - GRPO + - 8*H20 + - \ + - megatron + - vllm0.8.2 + - `qwen2-32b_grpo_8_h20_megatron_vllm `_ + - `Xiangyongan `_ + * - MIN + - Qwen2.5-32B + - GRPO-LoRA + - 4*H100 + - 180 + - fsdp + - vllm0.8.3 + - `qwen2-32b_grpo-lora_4_h100_fsdp_vllm.sh `_ + - `SimonHuang `_ + +70B +~~~ + +.. list-table:: + :widths: auto + :header-rows: 1 + + * - Tag + - Model + - Task + - Resource + - MaxBatch + - Train + - Infer + - Link + - Contributor + * - MIN + - Qwen2-70B + - GRPO + - 32*H20 + - \ + - fsdp + - vllm0.8.2 + - `qwen2-70b_grpo_32_h20_fsdp_vllm `_ + - `Xiangyongan `_ + * - MIN + - Qwen2-70B + - GRPO + - 32*H800 + - \ + - fsdp + - vllm0.8.3 + - `qwen2-70b_grpo_32_h800_fsdp_vllm `_ + - `Xiangyongan `_ + * - MIN + - Qwen2.5-72B + - GRPO-LoRA + - 8*H100 + - 176 + - fsdp + - vllm0.8.3 + - `qwen2-72b_grpo-lora_8_h100_fsdp_vllm.sh `_ + - `SimonHuang `_ + +405B +~~~~ + +.. table:: + :widths: auto + + ====== ====== ====== ======== ======== ====== ====== ====== + tag model task resource MaxBatch train infer link + ====== ====== ====== ======== ======== ====== ====== ====== + \ \ \ \ \ \ \ + ====== ====== ====== ======== ======== ====== ====== ====== + +671B +~~~~ + +.. table:: + :widths: auto + + ====== ====== ====== ======== ======== ====== ====== ====== + tag model task resource MaxBatch train infer link + ====== ====== ====== ======== ======== ====== ====== ====== + \ \ \ \ \ \ \ + ====== ====== ====== ======== ======== ====== ====== ====== diff --git a/verl/docs/perf/dpsk.md b/verl/docs/perf/dpsk.md new file mode 100644 index 0000000000000000000000000000000000000000..400066cfc0b67e39c34acb0713384f24677ea165 --- /dev/null +++ b/verl/docs/perf/dpsk.md @@ -0,0 +1,88 @@ +# Training DeepSeek 671b + +Last updated: 08/20/2025. + +verl integrates Megatron to support large MoE models such as `Qwen3-235B-A22B` and `deepseek-ai/DeepSeek-V3`. This is an ongoing community effort. + +In the journey the community added the following features and optimizations that enable verl with larger models: +- per tensor weight resharding between rollout and training +- context parallelism and expert parallelism enabled via megatron +- dynamic batch size (sequence balance) for megatron +- reduced ray-related serialization overhead +- optimizer offloading, recomputation, and efficient kernels +- various debugging metrics and utils +- hybrid optimizer + +and the megatron backend now has a wider list of models supported: +- DeepSeek-V3 +- Moonlight +- Qwen3 +- Qwen2.5-VL (to be merged soon) +- Qwen2 +- Mixtral + +## Getting Started + +### preparation +The recommended image with pre-built Megatron dependency is `verlai/verl:app-verl0.4-vllm0.8.5-mcore0.13.0-preview`, which is built using the Dockerfile at [docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.vllm.mcore0.13.preview](https://github.com/verl-project/verl/blob/main/docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.vllm.mcore0.13.preview). + +The image is build in Hopper GPUs with DeepEP. It does not support None-Hopper GPUs, such as A100. You may need to reinstall DeepEP to work with A100. + +With `OFFLOAD_FRACTION=1`, the system's minimum requirements are lowered. It can run on as few as 96 H20 (96GB) GPUs for DeepSeek-V3, and on as few as 32 H20 (96GB) GPUs for Qwen3-235B-A22B. However, this configuration will use 1.6TB CPU memory per node. If you run out of CPU memory or require faster training speed, you can add more nodes. + +### DeepSeek 671b + +For DeepSeek-V3 671b, please refer to [examples/grpo_trainer/run_deepseek_v3_671b_megatron.sh](https://github.com/verl-project/verl/blob/main/examples/grpo_trainer/run_deepseek_v3_671b_megatron.sh). + +MTP and quantilization is disabled during RL training. + +To train your project, configure the following environment variables based on the number of available GPUs. These are recommended settings and can be adjusted based on your specific hardware. +| num gpus | NNODES | TP | PP | EP | OFFLOAD_FRACTION | OFFLOAD_OPTIM | LAST_LAYER | +| -- | -- | -- | -- | -- | -- | -- | -- | +| 96 | 12 | 8 | 12 | 8 | 1. | False | 6 | +| 128 | 16 | 8 | 16 | 8 | 0.5 | True | 1 | +| 256 | 32 | 8 | 16 | 8 | 0. | True | 1 | +| 512 | 64 | 1 | 16 | 32 | 0 | True | 1 | + +### Qwen3 235b + +For Qwen3-235b, please refer to [examples/grpo_trainer/run_qwen3_235b_a22b_megatron.sh](https://github.com/verl-project/verl/blob/main/examples/grpo_trainer/run_qwen3_235b_a22b_megatron.sh). + +To train your project, configure the following environment variables based on the number of available GPUs. These are recommended settings and can be adjusted based on your specific hardware. +| num gpus | NNODES | TP | PP | EP | OFFLOAD_FRACTION | OFFLOAD_OPTIM | LAST_LAYER | +| -- | -- | -- | -- | -- | -- | -- | -- | +| 32 | 4 | 4 | 8 | 4 | 1. | False | 6 | +| 64 | 8 | 4 | 8 | 4 | 0.5 | True | 6 | +| 128 | 16 | 4 | 8 | 4 | 0 | True | 6 | +| 256 | 32 | 4 | 8 | 4 | 0 | True | 6 | + +### Benchmark +Here are some benchmark results for DeepSeek / Qwen3-235B. All configurations match the recommended settings based on the number of GPUs. + +| model | num gpus | mean response length | rollout time(s) | GPU memory(GB) | CPU memory(GB) | MFU | step time(s) | +| -- | -- | -- | -- | -- | -- | -- | -- | +| DeepSeek 671b | 96 | 1960 | 1050 | 66 | 1500 | 0.19 | 1700 | + +### Qwen3-30B-A3B MOE + +For Qwen3-30b, please refer to [examples/grpo_trainer/run_qwen3_30b_a3b_megatron.sh](https://github.com/verl-project/verl/blob/main/examples/grpo_trainer/run_qwen3_30b_a3b_megatron.sh). + +To train your project, configure the following environment variables based on the number of available GPUs. These are recommended settings and can be adjusted based on your specific hardware. +| num gpus | NNODES | TP | PP | EP | OFFLOAD_FRACTION | OFFLOAD_OPTIM | MFU | +| -- | -- | -- | -- | -- | -- | -- | -- | +| 8 | 1 | 1 | 1 | 8 | 1. | True | 0.4 | +| 16 | 2 | 1 | 1 | 8 | 1. | True | 0.37 | +| 32 | 4 | 1 | 1 | 8 | 1. | True | 0.31 | + + +## Upcoming Optimizations + +The community continue to optimize large MoE models further, ongoing efforts include: +- further optimizing memory consumption, and provide recommended/tuned configurations with various machine types +- optimizing long context RL training performance +- performance improvement with SGLang x Megatron + +We invite the community to try and improve verl together. Get connected with us on [slack](https://join.slack.com/t/verlgroup/shared_invite/zt-2w5p9o4c3-yy0x2Q56s_VlGLsJ93A6vA)/[wechat](https://raw.githubusercontent.com/eric-haibin-lin/verl-community/refs/heads/main/WeChat.JPG)/[Github issues](https://github.com/verl-project/verl/issues/708)! + +## Acknowledgement +@vermouth1992 @ISEEKYAN @ETOgaosion @yzlnew @ShareLer @BearBiscuit05 @ccclyu @ann-qin-lu @SwordFaith @zzong2006 @zhaochenyang20 @ocss884 @eric-haibin-lin @chenhaiq @techkang diff --git a/verl/docs/perf/nsight_profiling.md b/verl/docs/perf/nsight_profiling.md new file mode 100644 index 0000000000000000000000000000000000000000..490de5e7e4f7b6ba6c0e372eb7c0c3bfce2a77b9 --- /dev/null +++ b/verl/docs/perf/nsight_profiling.md @@ -0,0 +1,94 @@ +# NVIDIA Nsight Systems profiling in verl + +Last updated: 06/20/2025. + +This guide explains how to use NVIDIA Nsight Systems for profiling verl training runs. + +## Configuration + +Profiling in verl can be configured through several parameters in the trainer configuration file (ppo_trainer.yaml or other files like dapo_trainer.yaml): + +### Prerequisites + +Nsight Systems version is important, please reference `docker/Dockerfile.vllm.sglang.megatron` for the version we used. + +### Global profiling control + +verl has one single controller process and multiple worker processes. Both controller and worker processes can be profiled. Since the controller process can be executed in any nodes in the cluster, there is a message printed in the logging to indicate the controller process node hostname and process id. + +In `global_profiler`, three new config entries control the profiler behaviors: + +* **`global_profiler.steps`**. List of step numbers at which profiling should be performed. For example: [1, 2, 5] will profile steps 1, 2, and 5. And ``null`` means no profiling. + +* **`global_profiler.profile_continuous_steps`**. If true, and the following `global_profiler.discrete==False`, then the continuous steps in `global_profiler.steps` will be combined into one database. For example the above step 1 and 2 are in one database, and 5 in another. If false, every step occupies at least one database. The reason for this config is to observe the program behaviors between steps. + +Nsys options in controller nodes and worker nodes are configured in `global_profiler.global_tool_config.nsys`: + +* **`global_profiler.global_tool_config.nsys.controller_nsight_options`**. This config group is for the single controller. All fields in this config group will be just sent to Nsight Systems when Ray starts the controller process. `ppo_trainer.yaml` provides a workable example. Users can reference [Nsight Systems manual](https://docs.nvidia.com/nsight-systems/UserGuide/index.html) and [Ray user guide](https://docs.ray.io/en/latest/ray-observability/user-guides/profiling.html) for more details. +* **`global_profiler.global_tool_config.nsys.worker_nsight_options`**. This config group is for the worker processes. Similarly all fields in this config group will be just sent to Nsight Systems when Ray starts the controller process. Capture range is used to control the profiler when to start and stop. So `capture-range: "cudaProfilerApi"` is fixed and does not change it. Users can change `capture-range-end` with some accurate calculation or just leave it `null`. + +### Worker process profiling + +Verl manages mulitiple RL roles, _Actor_, _Ref_, _Rollout_, _Critic_, _Reward_, which are implemented in different Worker classes. And these workers can be combined into one Ray Actor, running in a process group. Each RL role has its own profiling config group, `profiler`, which consists of three fields: + +* **`all_ranks` and `ranks`**. When `all_ranks` is set `True` then all ranks will be profiled; when set `False`, `ranks` will be profiled. By default, verl profiles the whole training process in a series ` worker_process_..nsys-rep` files for each process rank. PID is the process ID; RID is the capture range ID. +* **`discrete`**. When set `False`, all the roles actions in one training step will be dumped in one database. When set `True`, the actions annotated by `DistProfiler.annotate` will be dumped into a discrete database. In this case, each role's action occupies one ``. +* **Verl collocate mode**. Verl can combine two Worker sub classes to one Worker Actor. In this case, the user should take care that the combined Workers have consistent `discrete`. The Nsight Systems profiler uses a `torch.cuda.profiler.start()` and `stop()` pair to dump a `` database anyway. + +### where to find the profiling data + +By default the `*.nsys-rep` files are saved in the directory `/tmp/ray/session_latest/logs/nsight/` at each node. According to the Ray manual, this default directory is not changeable. ["however, Ray preserves the `--output` option of the default config"](https://docs.ray.io/en/latest/ray-observability/user-guides/profiling.html). + +Some users may think it is not convenient, but it is understandable that Ray may start hundreds of processes and it would be a big network file system pressure if we save the files in one central place. + +## Usage Example + +To enable profiling for specific components and steps, modify your ppo_trainer.yaml like this: + +### Disable profiler + +```yaml + profiler: + steps: null # disable profile +``` + +### Enable profiler and one database for one training step + +```yaml + global_profiler: + steps: [1, 2, 5] + discrete: False + actor_rollout_ref: + actor: + profiler: + enable: True + all_ranks: True + # rollout & ref follow actor settings + critic: + profiler: + enable: True + all_ranks: True + reward_model: + profiler: + enable: True + all_ranks: True +``` + +### Enable profiler and multiple databases for one training step + +```yaml + profiler: + steps: [1, 2, 5] + discrete: True +``` + +## Profiling Output + +When profiling is enabled, verl will generate Nsight Systems profiles for the specified components and steps. The profiles will include: + +- CUDA kernel execution +- Memory operations +- CPU-GPU synchronization +- NVTX markers for key operations + +Nsight Systems supports multi-report view, to open multiple databases together. In this mode, different processes and steps can be aligned in one time line for better analysis. diff --git a/verl/docs/perf/perf_tuning.rst b/verl/docs/perf/perf_tuning.rst new file mode 100644 index 0000000000000000000000000000000000000000..cf25e9e9c627ada3c2caf63e479955e590684bfd --- /dev/null +++ b/verl/docs/perf/perf_tuning.rst @@ -0,0 +1,224 @@ +Performance Tuning Guide +============================== + +Last updated: 07/17/2025. + +Author: `Guangming Sheng `_, `Jiali Zheng `_ + +In this section, we will discuss how to tune the performance of all the stages in verl, including: + +1. Rollout generation throughput. + +2. Enable ``use_remove_padding=True`` for sequence packing (i.e., data packing and remove padding). + +3. Batch size tuning for forward and backward computation + +4. Enable ``use_dynamic_bsz=True`` for higher throughput. + +5. Utilize Ulysses Sequence Parallel for Long Context Training + +6. LigerKernel for SFT performance optimization + +7. Forward prefetch in FSDP training backend + +8. Memory optimization for entropy calculation from logits + +Rollout Generation Tuning +-------------------------- + +verl currently supports two rollout backends: vLLM and TGI (with SGLang support coming soon). + +Below are key factors for tuning vLLM-based rollout. Before tuning, we recommend setting ``actor_rollout_ref.rollout.disable_log_stats=False`` so that rollout statistics are logged. + +- Increase ``gpu_memory_utilization``. + + - For vLLM v0.7.0 and later, the vLLM instance will only use gpu_memory_utilization of the **total** memory. + - For SGLang, it's the fraction of the free GPU memory used for **static** memory like model weights and KV cache. However, the remaining (1-gpu_memory_utilization) will also be used during inference. + + However, if model parameters and optimizer states are not offloaded, using too high a fraction can lead to OOM. + A value between 0.5 and 0.7 often strikes a good balance between high throughput and avoiding OOM. + + Note: since the definition of ``gpu_memory_utilization`` varies across inference engines, a value that works well for one engine may cause OOM for another. + +- Adjust ``max_num_seqs`` or ``max_num_batched_tokens``. + If the GPU cache utilization is relatively low in the log, increase ``max_num_seqs`` or ``max_num_batched_tokens`` + can enlarge the effective batch size in the decoding stage, allowing more concurrent requests per batch. + We recommend setting ``max_num_batched_tokens > 2048`` for higher throughput. + +- Use a smaller ``tensor_parallel_size``. + When GPU resources allow, a smaller tensor parallel size spawns more vLLM replicas. + Data parallelism (DP) can yield higher throughput than tensor parallelism (TP), but also increases KVCache consumption. + Carefully balance the trade-off between more replicas and higher memory usage. + Our experiment in Sec. 8.4 of `HybridFlow paper `_ evaluate this trade-off. + +- Balance performance and memory using ``cudagraph_capture_sizes``. + If ``cudagraph_capture_sizes`` is set, vLLM will try to capture the model execution graph for different batch sizes. + Since cudagraph memory can not be offloaded to cpu, The memory stay in gpu when update actor is running. + Using smaller batch sizes can avoid OOM but slightly reduce throughput. + Must to set ``enforce_eager=False`` to use ``cudagraph_capture_sizes``. + +More tuning details such as dealing with Preemption and Chunked-prefill +can be found in `vLLM official tuning guide `_ + +For optimal performance, we recommend using vLLM v0.8.3 or later. See https://github.com/verl-project/verl/blob/main/docs/README_vllm0.8.md for details. + +Enable remove padding (sequence packing) +----------------------------------------- + +Currently, for llama, mistral, gemma1 and qwen based models, users can enable `use_remove_padding=True` to utilize the +sequence packing implementation provided by transformers library. + +For other models, transformers library may also support it but we haven't tested it yet. +Users can add the desired model config to the `test_transformer.py `_ file. +And test its functionality by running the following command: + +.. code-block:: bash + + pytest -s tests/models/test_transformer.py + +If the test passes, you can add your desired model into the model `registry.py `_ file. +Then, you can enjoy the performance boost of sequence packing +and welcome to PR your tested model to verl! + + +Batch Size Tuning +----------------- + +To achieve higher throughput in experience preparation (i.e., model fwd) and model update (i.e., actor/critic fwd/bwd), +users may need to tune the ``*micro_batch_size_per_gpu`` for different computation. + +In verl, the core principle for setting batch sizes is: + +- **Algorithmic metrics** (train batch size, PPO mini-batch size) are *global* (from a single-controller perspective), + normalized in each worker. See the `normalization code `_. + +- **Performance-related parameters** (micro batch size, max token length for dynamic batch size) are *local* parameters that define the per-GPU data allocations. + See the `normalization code `_. + +.. note:: In your training script, please use ``*micro_batch_size_per_gpu`` instead of ``*micro_batch_size``. + So that you don't need to consider the normalization of the ``micro_batch_size`` and ``micro_batch_size`` will be deprecated. + +Batch Size Tuning tips +"""""""""""""""""""""" + +Therefore, users may need to tune the ``*micro_batch_size_per_gpu`` to accelerate training. Here're some tips: + +1. **Enable gradient checkpointing**: + Set ``actor_rollout_ref.model.enable_gradient_checkpointing=True`` and ``critic.model.enable_gradient_checkpointing=True``. + This often allows for larger micro-batch sizes and will be beneficial for large mini-batch training. + +2. Increase the ``*micro_batch_size_per_gpu`` as much as possible till equals to normalized ``mini_batch_size``. + +3. **Use larger forward-only parameters**: + Forward only parameter, such as ``actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu``, + ``actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu``, ``critic.forward_micro_batch_size_per_gpu`` could be larger (e.g., 2x) than training related micro batch sizes, + such as ``actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu``, ``critic.ppo_micro_batch_size_per_gpu``. + +4. **Allow larger micro-batch sizes for Critic and Reward models**: + micro batch size of Critic and Reward model could be larger than Actor model. This is because the actor model has much larger vocab size in the final layer. + +5. **Enable activation offloading**: + Set ``actor_rollout_ref.model.enable_activation_offload=True`` and ``critic.model.enable_activation_offload=True``. + This often works together with gradient checkpointing to get larger micro-batch sizes and it's only available in FSDP backend now. + +Tuning for Dynamic Batch Size +----------------------------- + +Dynamic batch size is a technique that allows the model to process similar number of tokens in a single forward pass (with different actual batch sizes). +This can significantly improve the training efficiency and reduce the memory usage. + +To utilize this technique, users can set ``use_dynamic_bsz=True`` in actor, ref, critic and reward models. +With ``use_dynamic_bsz=True``, users don't need to tune ``*micro_batch_size_per_gpu``. +Instead, users should tune the following parameters: + +- ``actor_rollout_ref.actor.ppo_max_token_len_per_gpu``, ``critic.ppo_max_token_len_per_gpu``: + The maximum number of tokens to be processed in fwd and bwd of ``update_policy`` and ``update_critic``. + +- ``actor_rollout_ref.ref.log_prob_max_token_len_per_gpu`` and ``actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu``: + The maximum number of tokens to be processed in a the fwd computation of ``compute_log_prob`` and ``compute_ref_log_prob``. + +- ``critic.forward_micro_batch_size_per_gpu``, ``reward_model.forward_micro_batch_size_per_gpu``: + The maximum number of tokens to be processed in a the fwd computation of ``compute_values``, ``compute_rm_score``. + +Dynamic Batch Size Tuning tips +"""""""""""""""""""""""""""""" + +Here're some tips to tune the above parameters: + +1. **Increase** ``actor_rollout_ref.actor.ppo_max_token_len_per_gpu`` + Make it at least 2 x (max_prompt_length + max_response_length). See `run_qwen3_8b_fsdp.sh `_ for an example. + Try to increase it to get higher throughput. + +2. **Forward-only parameters can be larger**: + Similar to the non-dynamic-batch scenario, forward-only token limits can exceed those used in forward/backward operations. + +3. **Use larger limits for Critic and Reward models**: + Critic and Reward parameters can be set at least 2× the Actor’s limits. See + `run_qwen3_8b_fsdp.sh `_ for an example. + +.. :math:`\text{critic.ppo_max_token_len_per_gpu} = 2 \times \text{actor.ppo_max_token_len_per_gpu})`. + +Ulysses Sequence Parallel for Long Context Training +---------------------------------------------------- + +To utilize this technique, users can set ``ulysses_sequence_parallel_size>1`` in actor, ref, critic and reward models. + +We support different model utilize different ulysses_sequence_parallel_size sizes. + +To train long sequence (>32k), users may need to decrease the ``*micro_batch_size_per_gpu`` and ``*max_token_len_per_gpu`` to avoid OOM. + +LigerKernel for training performance +-------------------------------------- + +LigerKernel provides fused Triton kernels (RMSNorm, SwiGLU, RoPE) that can improve training throughput. It works with both SFT and RL (PPO/GRPO) training, including vision-language models. + +1. Install liger-kernel via ``pip3 install liger-kernel``. Set ``use_liger`` in your configuration: + + .. code-block:: yaml + + model: + use_liger: True # Enable LigerKernel + +2. The default value is ``False``. When enabled, verl applies Liger's fused RMSNorm, SwiGLU, and RoPE kernels to the model. The ``fused_linear_cross_entropy`` optimization is disabled because verl computes log-probabilities via its own path. + +3. ``use_liger`` is compatible with ``use_fused_kernels`` — they operate at different levels (Liger optimizes model internals, fused kernels optimize the output head). Using both together gives the best speed-memory tradeoff. + +Forward prefetch in FSDP training backend +---------------------- + +During the training phase, users can enable forward prefetching in FSDP by setting ``fsdp_config.forward_prefetch=True``. For example, ``actor_rollout_ref.actor.fsdp_config.forward_prefetch=True``. This configuration prefetches the next forward-pass all-gather operation before completing the current forward computation, overlapping communication with computation and improving efficiency. For further details, refer to the `FSDP forward_prefetch `_ documentation. + +.. note:: + Backward prefetch is unsupported because the ``BACKWARD_POST`` policy may prefetch incorrectly in nested-module cases. For details, see the `FSDP documentation `_ + +Migrating to FSDP2 +---------------------- + +FSDP2 offers notable improvements over FSDP1. According to `PyTorch TorchTitan benchmarks `_: + +- 7% lower GPU memory usage on average +- 1.5% throughput improvement with BF16 training +- Better composability with DTensor and per-parameter sharding + +**Enabling FSDP2 in VERL:** + + .. code-block:: python + + # Enable FSDP2 in actor configuration + actor_rollout_ref.actor.strategy="fsdp2" + +.. note:: + FSDP2 requires PyTorch 2.1+ and is recommended for models with transformer architecture. + +Memory optimization for entropy calculation from logits +---------------------- + +The ``logits`` tensor (typically of shape ``[bsz*seq_len, voc]``) can consume significant memory. When using ``compute_entropy_from_logits``, memory usage reaches approximately ``[bsz*seq_len, voc] × (4 bytes (float32) + 2 bytes (autocast for softmax+logsumexp) + 1 byte (softmax output))``. + +To reduce this memory peak, enable chunked computation by setting: +``actor_rollout_ref.ref.entropy_from_logits_with_chunking = True`` +This processes the tensor in chunks of shape ``[chunk_size, voc]`` (e.g., 2048) rather than the full sequence length, exclusively during the model's forward pass. + +Additionally, during training, standard gradient checkpointing (``enable_gradient_checkpointing=True``) does not apply to entropy calculations. To reduce memory peaks in this context, set: +``actor_rollout_ref.actor.entropy_checkpointing = True`` +This enables entropy recomputation specifically for the entropy calculation, lowering memory usage during training. diff --git a/verl/docs/perf/torch_profiling.md b/verl/docs/perf/torch_profiling.md new file mode 100644 index 0000000000000000000000000000000000000000..cef455de3d94dffb4080d4cdb60c09206386bd15 --- /dev/null +++ b/verl/docs/perf/torch_profiling.md @@ -0,0 +1,111 @@ +# PyTorch Profiling in verl + +Last updated: 01/13/2026. + +This guide explains how to use the native [PyTorch Profiler](https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html) for profiling verl training runs. + +## Configuration + +Profiling in verl can be configured through parameters in the trainer configuration file (e.g., `ppo_trainer.yaml`). + +### Global Profiling Control + +In `global_profiler`, you can control when and how profiling occurs globally: + +* **`global_profiler.steps`**: List of step numbers to profile. E.g., `[1, 2, 5]` profiles steps 1, 2, and 5. Set to `null` to disable. +* **`global_profiler.save_path`**: Directory to save the profiling results. Default is `outputs/profile`. + +### Role Profiling Control + +Each RL role (Actor, Critic, etc.) has its own `profiler` configuration: + +* **`enable`**: Whether to enable profiling for this role. +* **`all_ranks`**: If `True`, profiles all ranks. +* **`ranks`**: List of specific ranks to profile if `all_ranks` is `False`. +* **`tool_config.torch`**: Configuration specific to the PyTorch Profiler. + +#### PyTorch Profiler Options (`tool_config.torch`) + +You can customize the PyTorch Profiler behavior using the following fields under `tool_config.torch`: + +* **`contents`**: List of contents to profile. + * **`cpu`**: Profile CPU activities. + * **`cuda`**: Profile CUDA activities. + * **`memory`**: Track tensor memory allocation/free. + * **`shapes`**: Record shapes of operator inputs. + * **`stack`**: Record source code file and line number. +* **`schedule`**: (Advanced) configuration for `wait`, `warmup`, `active`, `repeat` cycles. + +## Examples + +### 1. End-to-End Collection + +Collects performance data for all steps in a single trace file. + +```yaml +global_profiler: + steps: [1, 2, 5] + save_path: ./outputs/profile + +actor_rollout_ref: + actor: + profiler: + enable: True + all_ranks: True + tool_config: + torch: + discrete: False + contents: [cpu, cuda] + # rollout & ref follow actor settings +``` + +### 2. Discrete Mode Collection + +Discrete mode saves separate trace files for each step. This is useful for detailed analysis and is **mandatory** when using Agent Loop. + +**Configuration Example** + +This configuration supports profiling both Training (Actor) and Inference (Rollout). You can enable/disable them independently. + +```yaml +actor_rollout_ref: + actor: + profiler: + enable: True # Set to True to profile training + all_ranks: False + ranks: [0] # Global Rank 0 + tool_config: + torch: + discrete: True + contents: [cpu, cuda] + rollout: + profiler: + enable: True # Set to True to profile inference + all_ranks: False + ranks: [0] # In Agent Loop, this is the Replica Rank (e.g. 0-th instance) + tool_config: + torch: + discrete: True # REQUIRED + # ref follow actor settings +``` + +**Agent Loop Mode Description** + +When Rollout runs in [Agent Loop](../advance/agent_loop.rst) mode, performance data for the Rollout phase **must be collected using discrete mode**. In this case, the Profiler is triggered by the inference engine backend. + +1. Rank Definition: ranks in the Rollout configuration refers to Replica Rank (inference instance index), not Global Rank. + +2. Inference Engine Support: Currently, vLLM and SGLang engines are supported without additional settings. Specific details are as follows: + + * **vLLM Engine**: Automatically collects AsyncLLM scheduling stacks and inference process performance data. + * **SGLang Engine**: Automatically collects inference process performance data. Does not support the memory option in contents. + +## Visualization + +Collected trace files (usually `.json` or `.json.gz`) are stored in the configured `save_path`. + +You can visualize them using: + +1. **Chrome Tracing**: Open `chrome://tracing` in a Chrome browser and load the JSON file. +2. **Perfetto**: Open [ui.perfetto.dev](https://ui.perfetto.dev/) and load the file (recommended for large traces). +3. **TensorBoard**: If using the TensorBoard plugin for PyTorch Profiler. diff --git a/verl/docs/perf/verl_profiler_system.md b/verl/docs/perf/verl_profiler_system.md new file mode 100644 index 0000000000000000000000000000000000000000..fc7ecc38eed92ca5e05274e23f40b6f1ce7033b0 --- /dev/null +++ b/verl/docs/perf/verl_profiler_system.md @@ -0,0 +1,36 @@ +# verl Profiler System + +Last updated: 08/18/2025. + +## Architecture + +The architecture of verl profiler system is like below: + +![verl-profiler-arch](https://raw.githubusercontent.com/eric-haibin-lin/verl-community/2bc7ed0ba2f37f21707bfac3b241eca4b86d1bc6/docs/verl_profiler_arch.png) + +There is a global profiler and tool configuration to set some common config in single controller level, deciding + +- `tool`: which tool to use +- `steps`: which steps to profile +- `save_path`: results saving path + +When some tool need to profile behavior of each role, configurations in role-level is needed: + +- `tool`: which tool to use +- `enable`: whether enable profiling on this role +- rank info: `all_ranks` and `rank` to decide which rank to profile or log output + +For tool config in role-level, there are some detailed behavior needed to control, like the `discrete` mode in nsys profiler. + +Every role has a profiler config, and by default, rollout/ref/reward models follow the Actor's behavior. + +## To Add a new profiling tool + +New added profiling tool shall reuse the current APIs as much as possible. + +1. The logic of **whether to use the tool**: `tool == [new tool]`. +2. Add the global and local tool config to `ppo_trainer.yaml`/`ppo_megatron_trainer.yaml` and each `[role].yaml`, under `global_tool_config.[new tool]` and `tool_config.[new tool]` +3. The tool config should be implemented in `verl/utils/profiler/config.py`, inherit the `BaseConfig` class. +4. Implement profiling tool initialization logic using configurations in `global_profiler.global_tool_config.[new tool]` and the results saving logics (can also save in role-level profile) +5. For role function-level profiling, please follow the nsys profiler way in `nvtx_profiler.py`, implement a profiler class inherit `DistProfiler` and import new profiler in `verl/utils/profiler/__init__.py` +6. Add unit test and examples for others to use in convinience. \ No newline at end of file diff --git a/verl/docs/preparation/prepare_data.rst b/verl/docs/preparation/prepare_data.rst new file mode 100644 index 0000000000000000000000000000000000000000..c429e4b167967652a0c3fb52d9e0029f1b9899d4 --- /dev/null +++ b/verl/docs/preparation/prepare_data.rst @@ -0,0 +1,128 @@ +Prepare Data for Post-Training +======================================== + +Last updated: 02/09/2025. + +Before starting the post-training job, we need to prepare the data for +the policy training. The data should be stored in the parquet format. + +We provide several data preprocess scripts for different datasets, +including GSM8K, MATH, HelloSwag, Full_hh_rlhf. To prepare other datasets, we need +to follow the following steps: The data preprocess script can be divided +into two parts: + +1. The first part is the common part, which loads the dataset from + huggingface's ``datasets`` package. Then preprocess the datasets with + the ``make_map_fn`` and then store in the parquet format. + +.. code:: python + + import re + import os + import datasets + + from verl.utils.hdfs_io import copy, makedirs + import argparse + + # To extract the solution for each prompts in the dataset + # def extract_solution(solution_str): + # ... + + + if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--local_dir', default='/opt/tiger/gsm8k') + parser.add_argument('--hdfs_dir', default=None) + + args = parser.parse_args() + + num_few_shot = 5 + data_source = 'openai/gsm8k' + + dataset = datasets.load_dataset(data_source, 'main') + + train_dataset = dataset['train'] + test_dataset = dataset['test'] + + # Construct a `def make_map_fn(split)` for the corresponding datasets. + # ... + + train_dataset = train_dataset.map(function=make_map_fn('train'), with_indices=True) + test_dataset = test_dataset.map(function=make_map_fn('test'), with_indices=True) + + local_dir = args.local_dir + hdfs_dir = args.hdfs_dir + + train_dataset.to_parquet(os.path.join(local_dir, 'train.parquet')) + test_dataset.to_parquet(os.path.join(local_dir, 'test.parquet')) + + makedirs(hdfs_dir) + + copy(src=local_dir, dst=hdfs_dir) + +2. The users are required to implement the ``make_map_fn()`` function + (as well as the ``extract_solution``) on their own to support + different datasets or tasks. + +We already implemented the data preprocess of GSM8k, MATH, Hellaswag and Full_hh_rlhf +datasets. And we take the GSM8k dataset as an example: + +**GSM8K** + +In the ``make_map_fn``, each data field should consist of the following +5 fields: + +1. ``data_source``: The name of the dataset. To index the corresponding + reward function in the ``RewardModel`` +2. ``prompt``: This field should be constructed in the format of + huggingface chat_template. The tokenizer in ``RLHFDataset`` will + apply chat template and tokenize the prompt. +3. ``ability``: Define the task category. +4. ``reward_model``: Currently, we only utilize the ``ground_truth`` + field during evaluation. The ``ground_truth`` is computed by the + ``extract_solution`` function. **NOTED** that the implementation of + the corresponding reward function should align with this extracted + ``ground_truth``. +5. ``extra_info``: Record some information of the current prompt. Not + use for now. + +.. code:: python + + def extract_solution(solution_str): + solution = re.search("#### (\\-?[0-9\\.\\,]+)", solution_str) # extract the solution after #### + assert solution is not None + final_solution = solution.group(0) + final_solution = final_solution.split('#### ')[1].replace(',', '') + return final_solution + + instruction_following = "Let's think step by step and output the final answer after \"####\"." + + # add a row to each data item that represents a unique id + def make_map_fn(split): + + def process_fn(example, idx): + question = example.pop('question') + + question = question + ' ' + instruction_following + + answer = example.pop('answer') + solution = extract_solution(answer) + data = { + "data_source": data_source, + "prompt": [{ + "role": "user", + "content": question + }], + "ability": "math", + "reward_model": { + "style": "rule", + "ground_truth": solution + }, + "extra_info": { + 'split': split, + 'index': idx + } + } + return data + + return process_fn diff --git a/verl/docs/preparation/reward_function.rst b/verl/docs/preparation/reward_function.rst new file mode 100644 index 0000000000000000000000000000000000000000..760898cd98c3956b9624c0ad99c29a9dfc838dcf --- /dev/null +++ b/verl/docs/preparation/reward_function.rst @@ -0,0 +1,71 @@ +Implement Reward Function for Dataset +====================================== + +Last updated: 06/02/2025. + +For each dataset, we need to implement a reward function or utilize a reward model to compute the rewards for the generated responses. +We already pre-implemented some reward functions in `reward_score directory `_. +You can also use customized reward functions. + +Currently, we support reward functions for GSM8k and MATH datasets. For RLHF datasets (e.g., +full_hh_rlhf) and Code Generation (e.g., APPS), we utilize reward model +and SandBox (will opensource soon) for evaluation respectively. + +RewardManager +------------- + +In the entrypoint of the PPO Post-Training script `main_ppo.py `_, +we implement a ``RewardManager`` that utilize pre-implemented reward functions to compute the scores for each response. + +In the ``RewardManager``, we implemented a ``__call__`` function to +compute the score for each response. +All the reward functions are executed by ``compute_score_fn``. +The input is a ``DataProto``, which includes: + +- ``input_ids``, ``attention_mask``: ``input_ids`` and ``attention_mask`` after applying + chat_template, including prompt and response +- ``responses``: response tokens +- ``ground_truth``: The ground truth string of the current prompt. + Stored in ``non_tensor_batch`` in the ``DataProto``, which should be + preprocessed in the parquet files. +- ``data_source``: The dataset name of the current prompt. Stored in + ``non_tensor_batch`` in the ``DataProto``, which should be + preprocessed in the parquet files. + +After detokenize the responses, the responses string and the ground +truth string will be input to the ``compute_score_fn`` to compute the +score for each response. + +Reward Functions +---------------- + +Pre-implemented +~~~~~~~~~~~~~~~ + +We already pre-implemented some reward functions in `reward_score directory `_. + +- In the `GSM8k example `_, we + force the response to output the final answer after four ####, then + use string matching to compare with the ground truth. If completely + correct, score 1 point; if the format is correct, score 0.1 points; if + the format is incorrect, score 0 points. +- In the `MATH example `_, we follow + the implementation in `lm-evaluation-harness repository `_. + +Customized +~~~~~~~~~~ + +You can implement customized reward functions in a separate file and specify them using ``custom_reward_function.path`` and ``custom_reward_function.name``. For the set of them, please refer to :ref:`config-explain-page`. + +The parameters of your reward function should be ``data_source``, ``solution_str``, ``ground_truth``, and ``extra_info``. +For example: + +.. code:: python + + def my_reward_fn(data_source, solution_str, ground_truth, extra_info=None): + return len(solution_str)/100 + +If you are testing only a single customized reward function, you can simply name it 'compute_score' and leave ``custom_reward_function.name`` unset. + +To run multiple tests with different customized reward functions, you can modify both ``custom_reward_function.path`` and ``custom_reward_function.name`` for each trial. +For instance, you might create a single `my_reward.py` file and implement multiple reward functions within it. This way, for different trials, you only need to adjust ``custom_reward_function.name``, making it more convenient to conduct multiple tests within scripts. diff --git a/verl/docs/requirements-docs.txt b/verl/docs/requirements-docs.txt new file mode 100644 index 0000000000000000000000000000000000000000..55ccdb8f7149bd6b774b592dca068e63e87256db --- /dev/null +++ b/verl/docs/requirements-docs.txt @@ -0,0 +1,13 @@ +# markdown support +recommonmark +myst_parser +# markdown table support +sphinx-markdown-tables + +# theme default rtd + +# crate-docs-theme +sphinx-rtd-theme + +# pin tokenizers version to avoid env_logger version req +tokenizers==0.21 diff --git a/verl/docs/sglang_multiturn/multiturn.rst b/verl/docs/sglang_multiturn/multiturn.rst new file mode 100644 index 0000000000000000000000000000000000000000..b40418b1ad559ffd2013f68754c84464337c23a8 --- /dev/null +++ b/verl/docs/sglang_multiturn/multiturn.rst @@ -0,0 +1,341 @@ +Multi-turn Rollout Support +========================== + +Last updated: 05/09/2026. + +Basic Configuration +~~~~~~~~~~~~~~~~~~~ + +To enable multi-turn rollout, make sure to configure the following fields in your rollout configuration: + +.. code-block:: yaml + + actor_rollout_ref: + rollout: + multi_turn: True + name: "sglang" + +These configuration activates the sglang engine for multi-turn interaction during rollout. + +Custom Tool Configuration +~~~~~~~~~~~~~~~~~~~~~~~~~ + +For custom environment interaction tools, you can implement your own tools based on ``verl.tools.base_tool.BaseTool``. Then, specify your tool configurations in a YAML file: + +.. code-block:: yaml + + tools: + - class_name: "" + config: + type: native + tool_schema: + +For stateless tools that don't need ``BaseTool``'s ``create``/``release`` lifecycle, see the `Function Tool Configuration`_ section below for the simpler ``@function_tool`` API. + +Finally, set the ``tools_config_file`` in your rollout config: + +.. code-block:: yaml + + actor_rollout_ref: + rollout: + tool_kwargs: + tools_config_file: + +This allows integration of customized tool behaviors during actor rollout steps. + +If your tool creates multi-modal inputs, you should return a list of multi-modal inputs in your tool.execute() implementation. + +Image and video should be processed before returning. For example, if you are using Qwen2.5-VL, you can use the following code to get the representations: + +.. code-block:: python + + async def create(self, ...) -> tuple[str, ToolResponse]: + ... + from verl.utils.dataset.vision_utils import process_image, process_video + + img1 = process_image(img1) + video1 = process_video(video1) + + # due to the (image | video) key is ("image" | "video") instead of ("images" | "videos") in vllm, we need to use ("image" | "video") to specify list of images/videos + # link: https://github.com/vllm-project/vllm/blob/3c545c0c3b98ee642373a308197d750d0e449403/vllm/multimodal/parse.py#L205 + return instance_id, ToolResponse(image=[img1, ...], video=[video1, ...], text="...") + + async def execute(self, ...) -> Tuple[str | Dict[str, Any], float, dict]: + ... + from verl.utils.dataset.vision_utils import process_image, process_video + + img1 = process_image(img1) + video1 = process_video(video1) + + # due to the (image | video) key is ("image" | "video") instead of ("images" | "videos") in vllm, we need to use ("image" | "video") to specify list of images/videos + # link: https://github.com/vllm-project/vllm/blob/3c545c0c3b98ee642373a308197d750d0e449403/vllm/multimodal/parse.py#L205 + return ToolResponse(image=[img1, ...], video=[video1, ...], text="..."), 0, {} + +remeber to set ``return_multi_modal_inputs: False`` in your dataset config in order to process the multi-modal inputs in the rollout correctly. +Refer to the `Handling Multi-Modal Inputs in Datasets`_ section for more details. + +Function Tool Configuration +~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +For stateless tools, defining a full ``BaseTool`` subclass plus a yaml schema is overkill. The ``@function_tool`` decorator lets you register a plain Python function as a tool; verl delegates schema inference to :func:`transformers.utils.get_json_schema`, which reads the function signature and a Google-style docstring. + +A typical configuration: + +.. code-block:: yaml + + actor_rollout_ref: + rollout: + mode: async + multi_turn: + enable: True + format: hermes # or any other format your model's chat template supports + function_tool_path: path/to/your_tools.py + agent: + default_agent_loop: tool_agent + +Define your tools in ``path/to/your_tools.py``. The decorator works either bare or with an explicit name; the function name is used as the tool name when bare: + +.. code-block:: python + + from verl.tools.function_tool import function_tool + + @function_tool + def get_weather(city: str) -> dict: + """Get the current weather for a city. + + Args: + city: The city to look up, e.g. "Tokyo" or "San Francisco". + """ + return {"temperature_c": 17.3, "condition": "drizzle"} + + @function_tool("calculator") # explicit name overrides the function name + def calculator(expression: str) -> str: + """Evaluate a Python-style arithmetic expression. + + Args: + expression: A Python-style arithmetic expression, e.g. "(3+4)*5". + """ + return str(eval(expression, {"__builtins__": {}}, {})) + +A few notes on the inferred schema: + +- Parameter types come from the function's type annotations. Beyond the primitives (``str``, ``int``, ``float``, ``bool``), generic and union forms work too: ``list[T]`` / ``dict[K, V]``, ``Optional[X]`` / ``X | None`` (yields ``nullable``), ``int | float`` (yields the JSON ``["integer", "number"]`` type), ``Literal["a", "b"]`` (yields ``enum``). +- Per-argument descriptions come from the ``Args:`` section of the docstring. +- Parameters without a default value are marked ``required``. +- The function may be sync or async; sync functions are dispatched through ``asyncio.to_thread`` so they don't block the event loop. +- ``*args`` / ``**kwargs`` are not representable in JSON Schema and will be rejected at registration time. Use a ``param: list[T]`` argument instead for variable-length inputs. +- Pass ``schema=`` to the decorator if you want to bypass inference entirely and supply your own ``OpenAIFunctionToolSchema`` (or a dict with the same shape). + +If the function violates ``transformers.get_json_schema``'s contract -- no docstring, missing type hint on a parameter, or a parameter that isn't documented in ``Args:`` -- registration raises a ``DocstringParsingException`` or ``TypeHintParsingException`` that points at the offending function. Fix the function rather than catching the exception. + +Return values are normalised the same way for every function tool: + +- ``str`` → wrapped as ``ToolResponse(text=...)`` +- ``dict`` → JSON-serialised into ``ToolResponse(text=...)`` +- ``ToolResponse`` → passed through unchanged +- ``(response, reward)`` or ``(response, reward, metrics)`` tuples are accepted, with ``None`` reward / metrics treated as ``0.0`` / ``{}`` + +``function_tool_path`` and ``tool_config_path`` can be set together. ``AgentLoopWorker`` merges both into one registry once on startup; name collisions across the two paths raise an error. ``RLHFDataset`` shares the same loader so prompt-length filtering sees the exact tool schemas the rollout will. + +The function-tool path intentionally exposes no framework-level lifecycle hooks: each call is just ``fn(**parameters)``, with no ``create``/``release`` or per-trajectory ``instance_id``. For *per-trajectory* state that must be torn down between rollouts (sandbox VMs, scratch directories, DB rows) or that needs ``tools_kwargs`` injected from the dataset, keep using ``BaseTool`` via ``tool_config_path``. + +Multi-turn Tokenization +~~~~~~~~~~~~~~~~~~~~~~~ + +Tokenizing multi-turn rollouts poses a challenge: after applying the chat template and tokenizing the full message list, it's hard to identify which tokens belong to assistant messages. Since the token list is flat, it lacks direct alignment with the message roles. + +To address this, we adopt a **delta-based tokenization** strategy. Each time the LLM generates a new message, we: + +1. Apply the chat template to all prior messages (`messages[:i]`). +2. Apply the chat template again including the latest message (`messages[:i+1]`). +3. Tokenize only the *delta* between these two serialized message strings. + +This ensures that only tokens generated by the assistant are included in the loss mask. + +.. code-block:: python + + # When using tokenizer + # Exclude the assistant prompt (e.g., "<|im_start|>assistant") from the loss by setting add_generation_prompt=True + prev = tokenizer.apply_chat_template(messages[:i], add_generation_prompt=True, tokenize=False) + curr = tokenizer.apply_chat_template(messages[:i+1], add_generation_prompt=False, tokenize=False) + token_ids += tokenizer.encode(curr[len(prev):], add_special_tokens=False) + loss_mask += [1] * len(token_ids) # Mask only the new assistant tokens + +.. code-block:: python + + # When using processor + # Exclude the assistant prompt (e.g., "<|im_start|>assistant") from the loss by setting add_generation_prompt=True + prev = processor.apply_chat_template(messages[:i], add_generation_prompt=True, tokenize=False) + prev_model_inputs = processor(text=prev, images=images, videos=videos, return_tensors="pt")[0].tolist() + curr = processor.apply_chat_template(messages[:i+1], add_generation_prompt=False, tokenize=False) + curr_model_inputs = processor(text=curr, images=images, videos=videos, return_tensors="pt")[0].tolist() + token_ids += curr_model_inputs["input_ids"][len(prev_model_inputs["input_ids"]):] + loss_mask += [1] * len(token_ids) # Mask only the new assistant tokens + +While we've validated this produces consistent results with full message tokenization, future models' chat template could break compatibility. To guard against silent inconsistencies, we compare the delta-based tokenization with full-tokenization results by default at the end of each rollout. + +If you see the following warning, you can check the mismatched substring in the log: + +.. code-block:: + + Inconsistent training and inference tokenization detected. This may lead to unexpected behavior during training. Please review your chat template to determine if this is intentional. For more information, refer to the multiturn README.md. + +The tokenization sanity check mode can be configured using the ``actor_rollout_ref.rollout.multi_turn.tokenization_sanity_check_mode`` parameter, which accepts the following values: + +- ``strict`` (default): Performs strict comparison between delta-based and full tokenization results, raising warnings for any differences. + +- ``ignore_strippable``: Ignores differences in whitespace characters (``\n``, ``\t``, ``\r``, spaces) while still checking for meaningful text mismatches. This is useful when debugging chat template issues where whitespace variations are expected and acceptable. + +- ``disable``: Completely disables the tokenization sanity check. Only use this if you have thoroughly validated that tokenization discrepancies are expected and won't impact training. + +Example configuration: + +.. code-block:: yaml + + actor_rollout_ref: + rollout: + multi_turn: + tokenization_sanity_check_mode: "ignore_strippable" # Choose from: "disable", "ignore_strippable", "strict" + +Handling Multi-Modal Inputs in Datasets +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +If your dataset includes multi-modal inputs (such as images or videos), you can control whether these are pre-processed and included in each sample by setting the return_multi_modal_inputs flag in your dataset config (used by RLHFDataset). + +- ``return_multi_modal_inputs: True`` (default): The dataset will pre-process and include a multi_modal_inputs dictionary for each sample. This dict contains the model-ready representations (e.g., image tensors, video tensors, etc.) as produced by your processor. This is useful for single-turn or SFT-style training, where the model expects all modalities to be present in the batch. + +- ``return_multi_modal_inputs: False``: The dataset will not include the multi_modal_inputs field. This is recommended for multi-turn RL or tool-augmented rollouts, where the model may generate new multi-modal inputs dynamically during rollout, and you want to avoid conflicts or redundant data in the batch. + + +Special Cases +^^^^^^^^^^^^^ + +Some models (e.g., Qwen/QwQ-32B and Qwen3 series) remove internal reasoning content during chat template rendering. As a result, the message content can vary across turns, making the delta-based tokenization inaccurate. + +For example, for the following conversation: + +.. code-block:: python + + messages = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "What is 2 + 2?"}, + {"role": "assistant", "content": "user asked about a simple math question. 2 + 2 = 4."}, + {"role": "user", "content": "Explain why."}, + {"role": "assistant", "content": "user wants to know the reasoning behind the answer. Search for a good explanation", + "tool_calls": [{"id": "tool1", "type": "search", "arguments": {"query": "Why is 2 + 2 = 4?"}}]}, + {"role": "tool", "content": "The sum of two and two is four because it is a basic arithmetic operation."}, + {"role": "assistant", "content": "The tool provided a good explanation.The sum of two and two is four because it is a basic arithmetic operation."} + ] + +1. Qwen/QwQ-32B will remove all reasoning content except the last assistant message after applying the chat template. + +.. code-block:: text + + <|im_start|>system + You are a helpful assistant.<|im_end|> + <|im_start|>user + What is 2 + 2?<|im_end|> + <|im_start|>assistant + 2 + 2 = 4.<|im_end|> + <|im_start|>user + Explain why.<|im_end|> + <|im_start|>assistant + + {"name": "", "arguments": {"query": "Why is 2 + 2 = 4?"}} + <|im_end|> + <|im_start|>user + + The sum of two and two is four because it is a basic arithmetic operation. + <|im_end|> + <|im_start|>assistant + The tool provided a good explanation. The sum of two and two is four because it is a basic arithmetic operation.<|im_end|> + +2. Qwen3 series will remove all reasoning content before the last user message. + +.. code-block:: text + + <|im_start|>system + You are a helpful assistant.<|im_end|> + <|im_start|>user + What is 2 + 2?<|im_end|> + <|im_start|>assistant + 2 + 2 = 4.<|im_end|> + <|im_start|>user + Explain why.<|im_end|> + <|im_start|>assistant + + user wants to know the reasoning behind the answer. Search for a good explanation + + + + {"name": "", "arguments": {"query": "Why is 2 + 2 = 4?"}} + <|im_end|> + <|im_start|>user + + The sum of two and two is four because it is a basic arithmetic operation. + <|im_end|> + <|im_start|>assistant + + The tool provided a good explanation. + + + The sum of two and two is four because it is a basic arithmetic operation.<|im_end|> + +To handle this, we fall back to a **fixed base conversation** containing only a single system and user message. Since this base doesn't include assistant messages or reasoning content, it remains consistent across turns. + +.. code-block:: python + + BASE_CHAT_HISTORY = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "I am a user."} + ] + prev = tokenizer.apply_chat_template(BASE_CHAT_HISTORY, add_generation_prompt=True, tokenize=False) + curr = tokenizer.apply_chat_template([*BASE_CHAT_HISTORY, messages[i]], add_generation_prompt=False, tokenize=False) + token_ids += tokenizer.encode(curr[len(prev):], add_special_tokens=False) + loss_mask += [1] * len(token_ids) + +This method works well for Qwen3 series. However, Qwen/QwQ-32B currently has a bug in its chat template. A fix_ has been proposed but not yet adopted. Until then, use the following command to download the fixed model revision: + +.. code-block:: bash + + pip install huggingface_hub + hf download Qwen/QwQ-32B --revision refs/pr/81 + +.. _fix: https://huggingface.co/Qwen/QwQ-32B/discussions/81 + +Discrepancy Between Training and Inference Templates +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Although the above approach fixes the delta mismatch issue, the removal of reasoning content in the inference-time chat template introduces a new discrepancy: training uses the full reasoning content, while inference does not. + +This mismatch can affect model performance in unpredictable ways. To avoid it, we default to using the full response (including reasoning) for both training and rollout. + +However, this approach comes with trade-offs: + +1. Long reasoning contents can easily exceed the model's context window, especially in multi-turn rollout. +2. There's a mismatch between rollout and production environment now—models will not have reasoning content from past turns if you use the default chat template in production. + +We are still evaluating the impact of these issues. If you experience context length problems or prefer rollouts that match production (i.e., exclude reasoning), you can enable: + +``actor_rollout_ref.rollout.multi_turn.use_inference_chat_template = True`` + +GSM8K Multi-turn Training Performance +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +See the training performance of multi-turn rollout on the GSM8K task HERE_. + +.. _HERE: https://wandb.ai/zhaochenyang20/gsm8k_async_rl/runs/1ro1r7om?nw=nwuserzhaochenyang20 + +.. _function_tool_examples: https://github.com/verl-project/verl/blob/main/tests/experimental/agent_loop/function_tool_examples.py + +Search Tool Integration +~~~~~~~~~~~~~~~~~~~~~~~ + +.. toctree:: + :maxdepth: 1 + + search_tool_example + +Code Walkthrough +~~~~~~~~~~~~~~~~~~~~~~~ +If you want to learn more in depth about the code execution flow, please read https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/tree/main/rlhf/verl/multi-turn/code-walk-through diff --git a/verl/docs/sglang_multiturn/sandbox_fusion.rst b/verl/docs/sglang_multiturn/sandbox_fusion.rst new file mode 100644 index 0000000000000000000000000000000000000000..995635627ed24d2e172f9913189a9eefcd869d79 --- /dev/null +++ b/verl/docs/sglang_multiturn/sandbox_fusion.rst @@ -0,0 +1,312 @@ +=============================== +Sandbox Fusion Tool Integration +=============================== + +Last updated: 05/09/2026. + +.. note:: + + The in-tree ``verl.tools.sandbox_fusion_tools.SandboxFusionTool`` reference + implementation discussed in this guide has been removed. The integration + pattern below remains accurate for users implementing their own + :class:`verl.tools.base_tool.BaseTool` subclass; for the previous + implementation see git history at commit ``f5e21df6`` and earlier. + +Motivations +=========== + +- As users of verl, we want to allow the model to call certain tools during Actor rollout, incorporating the results into the training process. +- A colleague from ByteDance proposed a paper aimed at enhancing model capability through code execution tools. +- We aim to support tool-calling capabilities of inference engines using `sandbox-fusion` as the code execution system, providing the community with a reimplementation of `retools`. + +Reward Compute with Sandbox Fusion + FaaS Integration +===================================================== + +- In current datasets and tasks, similar work already exists (e.g., Prime), which uses local processes as runners to execute model-generated code for reward computation. +- On this basis, #1429 has advanced the design by integrating FaaS as the runner for reward computation. + +Goals +===== + +- Adapt to the `sglang` tool-calling protocol and define tools for sandbox fusion. +- Integrate with the `async-rollout` process, ensuring sandbox fusion tools follow asyncIO conventions. +- Design and implement a basic rate limiter to prevent issues such as 429 errors. + +Non-Goals +========= + +- Training effectiveness is out of scope. +- Observability metrics are not considered. +- Distributed failover and component fault tolerance are not addressed. + +Design Details +============== + +Tool Schema Definition +---------------------- + +- Currently, only code execution is considered, requiring a `code` field in the JSON from the model. +- Only Python code is supported for now, so no `language` parameter is defined. + +.. code-block:: python + + OpenAIFunctionToolSchema( + type="function", + function=OpenAIFunctionSchema( + name="code_interpreter", + description="A tool for executing code.", + parameters=OpenAIFunctionParametersSchema( + type="object", + properties={ + "code": OpenAIFunctionPropertySchema( + type="string", + description="The code to execute.", + enum=None, + ) + }, + required=["code"], + ), + strict=False, + ) + ) + +Configuration Parameters +-------------------------- + ++----------------------------+--------------------------------------------------------------+ +| Parameter Name | Description | ++============================+==============================================================+ +| `num_workers` | Number of worker threads/processes per DP to request runner. | ++----------------------------+--------------------------------------------------------------+ +| `rate_limit` | Global limit of concurrent code executions. Default: 10 | ++----------------------------+--------------------------------------------------------------+ +| `default_timeout` | Timeout (in seconds) for each code execution. Default: 30 | ++----------------------------+--------------------------------------------------------------+ +| `default_language` | Default programming language. Default: "python" | ++----------------------------+--------------------------------------------------------------+ +| `enable_global_rate_limit` | Whether to enable global rate limiting. Default: True | ++----------------------------+--------------------------------------------------------------+ +| `sandbox_fusion_url` | URL for the veFaas sandbox execution service | ++----------------------------+--------------------------------------------------------------+ + +Rate Limiting Design +----------------------- + +Objective: + +- Limit the number of inflight requests using a token bucket model. + +- Ensure ordered submission to code runners to avoid starvation due to backoff. + +Design Highlights: + +- Use Ray Global Actor as a singleton distributed counter at cluster level. + +- Semaphore used for counting, with `acquire` and `release` in separate thread pools to preserve order. + +- Use Ray’s cloud-pickle to serialize functions for decoupled `ExecutionWorker`. + +.. code-block:: python + + @ray.remote(concurrency_groups={"acquire": 1,"release": 10}) + class TokenBucketWorker: + def __init__(self, rate_limit: int): + self.rate_limit = rate_limit + self.current_count = 0 + self._semaphore = threading.Semaphore(rate_limit) + + @ray.method(concurrency_group="acquire") + def acquire(self): + self._semaphore.acquire() + self.current_count += 1 + + @ray.method(concurrency_group="release") + def release(self): + self._semaphore.release() + self.current_count -= 1 + + def get_current_count(self): + return self.current_count + + class ExecutionWorker: + def __init__(self, enable_global_rate_limit=True, rate_limit=10): + self.rate_limit_worker = self._init_rate_limit(rate_limit) if enable_global_rate_limit else None + + def _init_rate_limit(self, rate_limit): + return TokenBucketWorker.options(name="rate-limiter", get_if_exists=True).remote(rate_limit) + + def execute(self, fn: Callable[..., T], *fn_args, **fn_kwargs) -> T: + with ExitStack() as stack: + stack.callback(self.rate_limit_worker.release.remote) + ray.get(self.rate_limit_worker.acquire.remote()) + try: + return fn(*fn_args, **fn_kwargs) + except Exception as e: + logger.warning(f"Error when executing code: {e}") + + def init_execution_pool(num_workers: int, enable_global_rate_limit=True, rate_limit=10, mode: PoolMode=PoolMode.ThreadMode): + if mode == PoolMode.ThreadMode: + return ray.remote(ExecutionWorker).options(max_concurrency=num_workers).remote( + enable_global_rate_limit=enable_global_rate_limit, + rate_limit=rate_limit + ) + else: + raise NotImplementedError("Process mode is not implemented yet") + +Tool Implementation +------------------- + +- Use `instance_id` to identify requests across multiple dialogue rounds. + +- Use `execution_pool` to implement async invocation. + +- Cleanup state after rollout completion. + +.. code-block:: python + + class SandboxFusionTool(BaseTool): + def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema): + ... + self.execution_pool = init_execution_pool(...) + ... + + async def create(self, instance_id: Optional[str] = None, ...): + ... + + async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> Tuple[str, float, dict]: + code = parameters.get("code", "") + timeout = parameters.get("timeout", self.default_timeout) + language = parameters.get("language", self.default_language) + if not isinstance(code, str): + code = str(code) + + result = await self.execution_pool.execute.remote(self.execute_code,instance_id,code,timeout,language) + self._instance_dict[instance_id]["reward"].append(result.strip()) + + return result, result, {} + + def execute_code(self,instance_id,code,timeout=30,language="python"): + result_status, metadata = _process_single_case(0, None, None,self.sandbox_fusion_url, code, timeout, language) + # we should always expect this since we don't have correct answer + if metadata["run_status"] == "Finished": + actual_output = metadata["stdout"] if metadata["stdout"] is not None else "" + return actual_output + else: + return "no stdout here" + + async def calc_reward(self, instance_id: str, ...): + ... + + async def release(self, instance_id: str, ...): + ... + +Test Plan +========= + +Unit Tests +---------- + +- **test_tools_registration**: Test tool registration and initialization. +- **test_rollout_req_creation**: Validate that `AsyncRolloutReq` is built correctly. +- **test_over_size_case**: Ensure rollout terminates early when exceeding `max_seq_len`. +- **test_tool_call_basic_case**: Mock `sglang` output, validate tool call and result. +- **test_tool_call_batch_case**: Test batch processing of tool calls. +- **test_basic_multi_process_init**: Validate Ray global actor behaves as singleton. +- **TestSingleNodeRateLimiterCase**: Verify rate limiter works in single-node mode. +- **test_rotten_execution**: Ensure rate limiter recovers from function errors. +- **TestMultiNodeRateLimiterCase**: Verify behavior in multi-node environments. + +e2e Tests +---------- +we provide e2e test scripts in `tests/special_e2e` folder, named `tests/special_e2e/run_gsm8k_fsdp_sgl_multiturn_sf_tool.sh` + +by setting 'trainer.rollout_data_dir' you can dump the rollout data to local disk. here is an sample taken from the rollout data: + +.. code-block:: python + + { + "input": " + + system\nYou are a math expert. You are given a question and you need to solve it step by step. Reasoning step by step before any tool call. You should use the `calc_gsm8k_reward` tool after step by step solving the question, before generate final answer at least once and refine your answer if necessary. Put your final answer in the format of `#### `.\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within XML tags:\n\n{\"type\": \"function\", \"function\": {\"name\": \"code_interpreter\", \"description\": \"A tool for executing code.\", \"parameters\": {\"type\": \"object\", \"properties\": {\"code\": {\"type\": \"string\", \"description\": \"The code to execute.\", \"enum\": null}}, \"required\": [\"code\"]}, \"strict\": false}}\n\n\nFor each function call, return a json object with function name and arguments within XML tags:\n\n{\"name\": , \"arguments\": }\n\n + + user\nJohn gets a bonus that's the same percentage every year. Last year he made $100,000 and got a $10,000 bonus. This year he makes $200,000. How much will John make this year when adding both his total pay and bonus together? Let's think step by step and output the final answer after `####`.\n + + assistant\n", + "output": "assistant\n\nOkay, let's see. John gets a bonus that's the same percentage every year. Last year, he made $100,000 and got a $10,000 bonus. This year, he makes $200,000. The question is asking how much he'll make this year when adding both his total pay and bonus together.\n\nFirst, I need to figure out the percentage bonus he gets. Since the bonus is the same percentage each year, if I can find that percentage based on last year's details, I can apply it to this year's earnings to find this year's bonus. Then, I'll add that bonus to his total pay this year.\n\nLast year, he earned $100,000 and got a $10,000 bonus. To find the percentage, I can divide the bonus by his total earnings. So, 10,000 divided by 100,000. Let me calculate that.\n\n10,000 / 100,000 equals 0.1, which is 10% as a percentage. So, his bonus is 10% of his pay each year.\n\nThis year, he makes $200,000. So, the bonus this year would be 10% of 200,000. Let me compute that using the code interpreter to ensure accuracy.\n + + \n{\"name\": \"code_interpreter\", \"arguments\": {\"code\": \"total_pay_this_year = 200000\\nbonus_percentage = 10 / 100\\nbonus_this_year = total_pay_this_year * bonus_percentage\\ntotal_income_this_year = total_pay_this_year + bonus_this_year\\nprint(total_income_this_year)\", \"executes\": \"True\"}}\n\n + + tool\n220000.0\n\n + + assistant\n\nSo the code shows that John's total income this year would be $220,000. This is consistent with the 10% bonus on last year's pay of $100,000, which was $10,000. Therefore, this year's bonus of $20,000 added to his total pay of $200,000 results in a total income of $220,000.\n\nI need to verify that there are no other factors at play. The problem states it's a constant percentage bonus each year, so there's no mention of changing percentages or other variables. Therefore, applying a straightforward percentage of 10% on this year's partial salary amount would result in $20,000. Adding that to his total pay gives $220,000. The calculation is correct and aligns with the question's parameters.\n\n\n#### 220000.0", + "score": 0, + "step": 1 + } + +here is the readable format version: + +.. code-block:: python + + [system] + + You are a math expert. You are given a question and you need to solve it step by step. Reasoning step by step before any tool call. You should use the `calc_gsm8k_reward` tool after step by step solving the question, before generate final answer at least once and refine your answer if necessary. Put your final answer in the format of `#### `. + + # Tools + + You may call one or more functions to assist with the user query. + + You are provided with function signatures within XML tags: + + {"type": "function", "function": {"name": "code_interpreter", "description": "A tool for executing code.", "parameters": {"type": "object", "properties": {"code": {"type": "string", "description": "The code to execute.", "enum": null}}, "required": ["code"]}, "strict": false}} + + + For each function call, return a json object with function name and arguments within XML tags: + + {"name": , "arguments": } + + + [user] + + John gets a bonus that's the same percentage every year. Last year he made $100,000 and got a $10,000 bonus. This year he makes $200,000. How much will John make this year when adding both his total pay and bonus together? Let's think step by step and output the final answer after `####`. + + [assistant] + + + Okay, let's see. John gets a bonus that's the same percentage every year. Last year, he made $100,000 and got a $10,000 bonus. This year, he makes $200,000. The question is asking how much he'll make this year when adding both his total pay and bonus together. + + First, I need to figure out the percentage bonus he gets. Since the bonus is the same percentage each year, if I can find that percentage based on last year's details, I can apply it to this year's earnings to find this year's bonus. Then, I'll add that bonus to his total pay this year. + + Last year, he earned $100,000 and got a $10,000 bonus. To find the percentage, I can divide the bonus by his total earnings. So, 10,000 divided by 100,000. Let me calculate that. + + 10,000 / 100,000 equals 0.1, which is 10% as a percentage. So, his bonus is 10% of his pay each year. + + This year, he makes $200,000. So, the bonus this year would be 10% of 200,000. Let me compute that using the code interpreter to ensure accuracy. + + {"name": "code_interpreter", "arguments": {"code": "total_pay_this_year = 200000\nbonus_percentage = 10 / 100\nbonus_this_year = total_pay_this_year * bonus_percentage\ntotal_income_this_year = total_pay_this_year + bonus_this_year\nprint(total_income_this_year)", "executes": "True"}} + + + [tool] + + 220000.0 + + [assistant] + + + So the code shows that John's total income this year would be $220,000. This is consistent with the 10% bonus on last year's pay of $100,000, which was $10,000. Therefore, this year's bonus of $20,000 added to his total pay of $200,000 results in a total income of $220,000. + + I need to verify that there are no other factors at play. The problem states it's a constant percentage bonus each year, so there's no mention of changing percentages or other variables. Therefore, applying a straightforward percentage of 10% on this year's partial salary amount would result in $20,000. Adding that to his total pay gives $220,000. The calculation is correct and aligns with the question's parameters. + + + #### 220000.0 + + +You can also use the `RolloutViewer` TUI tool to view the dumped rollout data: + + +.. code-block:: bash + + python scripts/rollout_viewer.py ${trainer.rollout_data_dir} + + +.. image:: https://github.com/user-attachments/assets/e34e5157-2880-4a21-afb2-73885d0dfb11 + :alt: RolloutViewer screenshot \ No newline at end of file diff --git a/verl/docs/sglang_multiturn/search_tool_example.rst b/verl/docs/sglang_multiturn/search_tool_example.rst new file mode 100644 index 0000000000000000000000000000000000000000..13dfeebdfe3489c66ea45150bccce5edfb9d7ab3 --- /dev/null +++ b/verl/docs/sglang_multiturn/search_tool_example.rst @@ -0,0 +1,272 @@ +======================= +Search Tool Integration +======================= + +Last updated: 05/09/2026. + +.. note:: + + The in-tree ``verl.tools.search_tool.SearchTool`` reference implementation + used throughout this guide has been removed. The end-to-end recipe + (config, retrieval server, training script) still works as documentation + of the integration pattern, but the ``SearchTool`` class must now be + provided by users (see :class:`verl.tools.base_tool.BaseTool`). + +Introduction +------------ +- We have added a search tool calling function to Multi-Turn RL, enabling the model to initiate retrieval requests during Actor rollout and directly use retrieval results for training. **We support using a local dense retriever as the retrieval tool, as well as integrating with your own local retrieval engine.** + + + +Quick Reproduction +------------------ + +Create a New Docker Container +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. code:: bash + + docker run \ + -it \ + --shm-size 32g \ + --gpus all \ + -v {Huggingface-Cache-Path}:/root/.cache \ + --ipc=host \ + --network=host \ + --privileged \ + --name sglang_{your-name} \ + lmsysorg/sglang:dev \ + /bin/zsh + +If you need to restart after exiting the container: + +.. code:: bash + + docker start -i sglang_{your-name} + +Update Python and Configure the Virtual Environment using uv +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. code:: bash + + apt update + apt install -y python3.10 python3.10-venv + + # Create a virtual environment + python3 -m venv ~/.python/verl-multiturn-rollout + + # Activate the virtual environment + source ~/.python/verl-multiturn-rollout/bin/activate + + # Install uv + python3 -m pip install uv + +Install verl Upstream +~~~~~~~~~~~~~~~~~~~~~ + +.. code:: bash + + cd ~ + git clone https://github.com/verl-project/verl.git + cd verl + + # Install verl + python3 -m uv pip install . + python3 -m uv pip install -r ./requirements_sglang.txt + + # Manually install flash-attn + python3 -m uv pip install wheel + python3 -m uv pip install packaging + python3 -m uv pip install flash-attn --no-build-isolation --no-deps + +Set Up a Local Retrieval Engine +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +If you are using your own local retrieval service, you can skip this +step. We chose the local dense retriever provided in the search-R1 +example; detailed instructions are in the `searchR1 +docs `__. +In brief: + +- The GPU version offers higher accuracy and speed; each GPU uses about + 5–7 GB of memory. +- The CPU version can be used for simple testing but has lower + retrieval precision, which will degrade training performance. See the + `retriever + documentation `__ + in search-R1 for details. +- Recommend using Conda to install faiss-gpu=1.8.0; venv may cause errors. + +**Note**: To start both the training process and the local retrieval +service, we launch two separate Python environments. The training uses +uv in the verl-multiturn-rollout environment, while the retriever uses +conda to install ``faiss-gpu``. + +.. code:: bash + + # Download the Miniconda installer script + wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda.sh + + # Install to $HOME/miniconda3 in batch mode + bash ~/miniconda.sh -b -p $HOME/miniconda3 + + # Activate conda (only in the current shell) + eval "$($HOME/miniconda3/bin/conda shell.bash hook)" + + # (Optional) Add conda to your default shell startup + conda init + + # Reload shell config + source ~/.bashrc + + # Create and activate the retriever environment with Python 3.10 + conda create -n retriever python=3.10 -y + conda activate retriever + + # Install PyTorch (with GPU support) and related libraries + conda install pytorch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 pytorch-cuda=12.1 -c pytorch -c nvidia -y + + # Install other Python packages + pip install transformers datasets pyserini huggingface_hub + + # Install the GPU version of faiss + conda install faiss-gpu=1.8.0 -c pytorch -c nvidia -y + + # Install the API service framework + pip install uvicorn fastapi + +Download the Indexing and Corpus +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The local retrieval files are large—prepare sufficient disk space. +Downloading is about 60–70 GB, and uncompressed takes about 132 GB: + +.. code:: bash + + conda activate retriever + + save_path=/the/path/to/save + python examples/sglang_multiturn/search_r1_like/local_dense_retriever/download.py --save_path $save_path + cat $save_path/part_* > $save_path/e5_Flat.index + gzip -d $save_path/wiki-18.jsonl.gz + +Start the Local flat e5 Retrieval Server +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +1. The first startup will download models and load the index. +2. Apart from the download, startup takes about 1–2 minutes. +3. After startup, each GPU uses about 5–7 GB of memory, leaving the rest + for multi-turn RL training. + +.. code:: bash + + conda activate retriever + + index_file=$save_path/e5_Flat.index + corpus_file=$save_path/wiki-18.jsonl + retriever_name=e5 + retriever_path=intfloat/e5-base-v2 + + python examples/sglang_multiturn/search_r1_like/local_dense_retriever/retrieval_server.py \ + --index_path $index_file \ + --corpus_path $corpus_file \ + --topk 3 \ + --retriever_name $retriever_name \ + --retriever_model $retriever_path \ + --faiss_gpu + +Set Up WANDB_API_KEY +~~~~~~~~~~~~~~~~~~~~ + +.. code:: bash + + export WANDB_API_KEY={YOUR_WANDB_API_KEY} + + # Define a timestamp function + function now() { + date '+%Y-%m-%d-%H-%M' + } + +**Preprocess the Dataset** +~~~~~~~~~~~~~~~~~~~~~~~~~~ + + **Note:** The following data processing and training commands must be + run in the verl-multiturn-rollout environment. + +.. code:: bash + + python3 examples/data_preprocess/preprocess_search_r1_dataset.py + +Testing on 8 x H20 +~~~~~~~~~~~~~~~~~~ + +.. code:: bash + + # Ensure the now() function is defined + # Create a logs directory + mkdir -p logs + + # Set GPUs and run with a suitable log path + export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 + + nohup bash examples/sglang_multiturn/search_r1_like/run_qwen2_5_3b_search_multiturn_fsdp.sh \ + trainer.experiment_name=qwen2.5-3b-it_rm-searchR1-like-sgl-multiturn-$(now) \ + > logs/searchR1-like$(now).log 2>&1 & + +Custom Search Configuration +--------------------------- + +To enable multi-turn reasoning, set the following fields in your config: + +.. code:: yaml + + actor_rollout_ref: + rollout: + name: "sglang" + multi_turn: + enable: True + +You must specify ``retrieval_service_url`` in ``examples/sglang_multiturn/config/tool_config/search_tool_config.yaml``, and properly configure concurrency. For more details on concurrency, refer to the Sandbox Fusion example: + +.. code:: yaml + + tools: + - class_name: verl.tools.search_tool.SearchTool + config: + retrieval_service_url: http://127.0.0.1:8000/retrieve + num_workers: 120 + rate_limit: 120 + timeout: 30 + +The retriever input/output formats are as follows. If your service +parameters match, only modify ``retrieval_service_url``. You can also +customize in ``search_r1_like_utils.py``. + +.. code:: python + + Input format: + { + "queries": ["What is Python?", "Tell me about neural networks."], + "topk": 3, + "return_scores": true + } + + Output format (when return_scores=True, similarity scores are returned): + { + "result": [ + [ # Results for each query + { + "document": doc, "score": score + }, + # ... more documents + ], + # ... results for other queries + ] + } + +Notes +----- + +1. The total training time is about 27 hours; meanwhile, the validation + dataset is very large (51 k), and each validation takes about 6000 s. + (Therefore, ``val_before_train=False`` by default) diff --git a/verl/docs/single_controller.rst b/verl/docs/single_controller.rst new file mode 100644 index 0000000000000000000000000000000000000000..03cde0f4b1a4cb6ed748290ca1de0f709b41eb33 --- /dev/null +++ b/verl/docs/single_controller.rst @@ -0,0 +1,336 @@ +The Design of ``verl.single_controller`` +============================================== + +Last updated: 05/21/2025. + +**Author:**\ `Wang Zhang `__ + +Preface +------- + +We prepared this document for developers of ``verl``, particularly those +interested in understanding or contributing to the +``verl.single_controller`` module. It is not intended for end users, but +for contributors seeking to understand the architectural rationale and +internal mechanics. + +-------------- + +Origin +------ + +The ``single_controller`` module originated from a request I received — +to adapt a toy single-process RLHF script into a distributed system with +minimal changes, while maintaining ease of debugging. + +Common practice — such as using PyTorch’s Distributed Data Parallel +(DDP) — typically involves wrapping ``nn.Module`` and launching multiple +processes that execute the same function under different ranks. However, +this approach presents two main limitations in the context of +distributed RLHF: - Difficulty representing multiple DAGs as required by +PPO; - Difficulty inspecting intermediate tensors during training. + +To maintain debuggability, we opted for a different approach — breaking +the training loop into well-defined stages like ``generate_sequences``, +``compute_advantages``, and so on. + +We selected `Ray `__ as the initial backend for +``verl`` due to its ability to expose Python class methods as RPC +endpoints. However, Ray’s default model only supports **one method call, +one RPC**, while training LLMs typically requires coordination across +multiple processes. + +To hide this multi-Ray actors invocation for a single method from users, +we introduced the following components: + +- ``WorkerGroup`` – manages a group of remote workers and provides + a unified interface for multi-process distributed computation; +- ``ResourcePool`` – binds computational resources to worker + processes; +- ``ClassWithArgs`` – enables delayed remote instantiation with + specified initialization arguments. + +-------------- + +A Running Example: ``generate_sequences`` +----------------------------------------- + +To illustrate the design, we walk through how the ``generate_sequences`` +method in the ``ActorRolloutRefWorker`` class is registered and invoked +across distributed workers. + +-------------- + +Step 1: Register with a Decorator +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The first step is to define the ``generate_sequences`` and decorate it +with ``@register`` as it will be called in driver script. + +**Source:** +`engine_workers.py `__ + +.. code:: python + + class ActorRolloutRefWorker(Worker): + ... + @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO) + def generate_sequences(self, prompts: DataProto): + prompts = prompts.to(torch.cuda.current_device()) + ... + +The ``@register`` decorator adds metadata to the ``generate_sequences`` +method. Currently, it doesn’t alter functionality, but attaches +attributes via a magic key (``MAGIC_ATTR``): + +**Source:** +`decorator.py `__ + +.. code:: python + + def register(dispatch_mode=Dispatch.ALL_TO_ALL, execute_mode=Execute.ALL, blocking=True, materialize_futures=True): + ... + def decorator(func): + @wraps(func) + def inner(*args, **kwargs): + if materialize_futures: + args, kwargs = _materialize_futures(*args, **kwargs) + return func(*args, **kwargs) + + attrs = {"dispatch_mode": dispatch_mode, "execute_mode": execute_mode, "blocking": blocking} + setattr(inner, MAGIC_ATTR, attrs) + return inner + + return decorator + +As the code shows, values of ``dispatch_mode``, ``execute_mode`` and +``blocking`` is attached the ``generate_sequences`` method. + +-------------- + +Step 2: Binding During Initialization +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +These attached attributes are extracted and utilized when +``ActorRolloutRefWorker``, wrapped in a ``RayClassWithArgs``, is passed +into a ``RayWorkerGroup``. + +**Source:** +`main_generation.py `__ + +.. code:: python + + ray_cls_with_init = RayClassWithInitArgs(cls=ray.remote(ActorRolloutRefWorker), config=config, role="rollout") + resource_pool = RayResourcePool(process_on_nodes=[config.trainer.n_gpus_per_node] * config.trainer.nnodes) + wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init) + +During the +`initialization `__ +of ``RayWorkerGroup``, two key steps occur: + +1. Worker instances (Ray actors) are created: + `RayWorkerGroup._init_with_resource_pool `__ +2. Methods decorated with ``@register`` are bound to ``RayWorkerGroup``: + `RayWorkerGroup._bind_worker_method `__ + +.. figure:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/worker_group_init.png?raw=true + :alt: initialization_and_binding_of_worker_group + + initialization_and_binding_of_worker_group + +The binding procedure is the heart of ``verl.single_controller``. + +**Key function:** +`WorkerGroup._bind_worker_method `__ + +.. code:: python + + def _bind_worker_method(self, user_defined_cls, func_generator): + ... + for method_name in dir(user_defined_cls): + try: + method = getattr(user_defined_cls, method_name) + assert callable(method) + except Exception: + continue # Skip properties + <<>> + +When a method has the ``MAGIC_ATTR``, the attributes set by +``@register`` are extracted: + +.. code:: python + + <<>> + if hasattr(method, MAGIC_ATTR): + attribute = getattr(method, MAGIC_ATTR) + dispatch_mode = attribute["dispatch_mode"] + execute_mode = attribute["execute_mode"] + blocking = attribute["blocking"] + + <<>> + +As show in the flow chart above, these attributes are fed into +``func_generator``. However, ``func_generator`` takes ``method_name``, +``dispatch_fn``, ``collect_fn``, ``execute_fn``, ``blocking``. We need +to find the corresponding ``dispatch_fn`` and ``collect_fn`` associated +with the ``dispatch_mode`` (``DP_COMPUTE_PROTO``) from +`DISPATCH_MODE_FN_REGISTRY `__: + +.. code:: python3 + + DISPATCH_MODE_FN_REGISTRY = { + Dispatch.ONE_TO_ALL: { + "dispatch_fn": dispatch_one_to_all, + "collect_fn": collect_all_to_all, + }, + ... + Dispatch.DP_COMPUTE_PROTO: { + "dispatch_fn": dispatch_dp_compute_data_proto, + "collect_fn": collect_dp_compute_data_proto, + }, + ... + } + +Similarly, the ``execute_fn`` is selected by ``execute_mode`` and +extracted by: + +.. code:: python + + <<>> + # get execute_fn_name + execute_mode = get_predefined_execute_fn(execute_mode=execute_mode) + wg_execute_fn_name = execute_mode["execute_fn_name"] + + # get execute_fn from string + try: + execute_fn = getattr(self, wg_execute_fn_name) + assert callable(execute_fn), "execute_fn must be callable" + except Exception: + print(f"execute_fn {wg_execute_fn_name} is invalid") + raise + <<>> + +In this ``generate_sequences`` cases: - +``dispatch_mode = Dispatch.DP_COMPUTE_PROTO`` - +``dispatch_fn = dispatch_dp_compute_data_proto`` - +``collect_fn = collect_dp_compute_data_proto`` - +``execute_fn = RayWorkerGroup.execute_all`` + +ONE_TO_ALL v.s. DP_COMPUTE_PROTO +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +``dispatch_mode`` is associated with a ``dispatch_fn`` and a +``collect_fn``. As the name implies, ``dispatch_fn`` processes the input +arguments in ``WorkerGroup`` and generate a batch (list) of input +arguments, each of which will be fed into a worker attached to the +``WorkerGroup``. + +``dispatch_fn`` of ``ONE_TO_ALL`` is +`dispatch_one_to_all `__, +which just duplicates all the input arguments into N replicas, where N +equals the number of Workers attached to the ``worker_group``: + +.. code:: python + + def dispatch_one_to_all(worker_group, *args, **kwargs): + args = tuple([arg] * worker_group.world_size for arg in args) + kwargs = {k: [v] * worker_group.world_size for k, v in kwargs.items()} + return args, kwargs + +``dispatch_fn`` of ``DP_COMPUTE_PROTO`` is +`dispatch_dp_compute_data_proto `__, +which uses ``DataProto.chunk`` to split a large ``DataProto`` into N +smaller ``DataProto``, where N equals the world_size (number of the +workers) of the ``worker_group``: + +.. code:: python + + def dispatch_dp_compute_data_proto(worker_group, *args, **kwargs): + from verl.single_controller.base.worker_group import WorkerGroup + + assert isinstance(worker_group, WorkerGroup) + # Note: enable auto padding for dp compute DatapProto + splitted_args, splitted_kwargs = _split_args_kwargs_data_proto_with_auto_padding( + worker_group.world_size, + *args, + **kwargs, + ) + return splitted_args, splitted_kwargs + +The ``collect_fn`` follows the same pattern and process a batch (list) +of returned value from all workers of a ``WorkerGroup`` and merge it +into a list as ``collect_all_to_all`` does or a large ``DataProto`` as +``collect_dp_compute_data_proto`` does. + +Finally, a new method is dynamically generated using ``func_generator`` +and added to the ``WorkerGroup`` instance: + +.. code:: python + + <<>> + # bind a new method to the RayWorkerGroup + func = func_generator( + self, + method_name, + dispatch_fn=dispatch_fn, + collect_fn=collect_fn, + execute_fn=execute_fn, + blocking=blocking, + ) + + try: + setattr(self, method_name, func) + method_names.append(method_name) + except Exception as e: + raise ValueError(f"Fail to set method_name {method_name}") from e + +This makes the method invocable via the ``WorkerGroup`` interface. + +-------------- + +Step 3: Call Chain +~~~~~~~~~~~~~~~~~~ + +All the machinery above ensures that distributed calls feel identical to +single-process ones. In the original single-process script, the code +looks like: + +.. code:: python + + rollout = Rollout() + rollout.generate_sequences(batch) + +With ``verl``, the multiprocess program becomes: + +.. code:: python + + rollout = RayWorkerGroup(resource_pool=[4], RayClassWithArgs(Rollout)) + rollout.generate_sequences(batch) + +.. figure:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/call_generate_sequences.png?raw=true + :alt: call_chain_of_generate_sequences + + call_chain_of_generate_sequences + +Behind this simple call: - ``dispatch_fn`` splits input across workers - +``execute_fn`` performs the actual remote invocation - ``collect_fn`` +gathers the results + +All of this is abstracted away, enabling developers to write distributed +code with minimal changes to their existing logic. + +-------------- + +Beyond RL Post-Training: Generalizing ``verl.single_controller`` +---------------------------------------------------------------- + +The ``verl.single_controller`` module generalizes well beyond +reinforcement learning. It provides a clean abstraction to batch-process +remote method calls, with automatic input/output handling. + +By minimizing the gap between single-process and multi-process scripts, +``verl.single_controller`` opens the door to distributed computing in +broader domains — not limited to RL post-training. + +We hope this design inspires more examples and extensions from the +community. diff --git a/verl/docs/start/agentic_rl.rst b/verl/docs/start/agentic_rl.rst new file mode 100644 index 0000000000000000000000000000000000000000..46ca53d447fabcec7702138c878ecc4dc5084504 --- /dev/null +++ b/verl/docs/start/agentic_rl.rst @@ -0,0 +1,133 @@ +Agentic RL Training +=================== + +Last updated: 07/15/2025. + +Overview +---------- +The goal of Agentic RL is to improve the performance of backend models from reinforcement learning to the Agent. During the training process, a series of features are developed: + +1. Server-based asynchronous rollout +2. Multi-turn conversations and tool calls +3. LangGraph-based Agent + + +This document explains the system principles and usage involved to help users implement Agentic RL. + + +Server-based Asynchronous Rollout +--------------------------------- + +Since Agents need to interact with the environment through various tool calls, in order to avoid GPU idling while waiting for tool call return results, an asyncio based co-routing mechanism is utilized to execute each rollout requests asynchronously, thereby improving training performance. To support asynchronous rollout, the inference engine (server) and the agent (client) are architecturally separated, implementing a server-based system with the following objectives: + +1. Enabling load balancing mechanisms to balance loads across multiple GPUs and reduce the impact of long-tail requests on performance. For this purpose, scheduling capabilities in stream mode (recipe\stream_mode) are implemented as a recipe. +2. Preventing agent specific features such as tracing from affecting the inference engine. + +System Architecture +~~~~~~~~~~~~~~~~~~~ + +.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/agent_loop.png?raw=true + +For more detail on internal design, please refer to :doc:`Agent Loop<../advance/agent_loop>`. + +System Components +~~~~~~~~~~~~~~~~~ + ++--------------------------+----------------------------------------------------------------------------+ +| Component | Role | ++==========================+============================================================================+ +| AgentLoop | Client, implements Agent functions | ++--------------------------+----------------------------------------------------------------------------+ +| LLMServerClient | Inference gateway, provides generate interface for AgentLoop | ++--------------------------+----------------------------------------------------------------------------+ +| AsyncServer | Server, each instance is connected to one DP group of the inference engine | ++--------------------------+----------------------------------------------------------------------------+ + +**"generate" Interface** + +The "generate" function based on ray actor is used between the Client and Server instead of the standard chat completion API. This is because the conversion between tokens and text can be irreversible. For example, the token converted from "" will be different from that generated by the LLM. During the training phase, it is necessary to strictly use the tokens generated by LLM inference to avoid inaccurate in computing advantage, which may affect model performance. Having the Server provide a token-based API helps the Client maintain the relationship between the text generated by tool calls and the tokens returned by the LLM, so as to output correct tokens for training. + + +**Inference Engine Adaptation** +AsyncServer uniformly provides a generate function to the upper layer, with separate implementations for SGLang and vLLM to hide underlying differences: + +1. The SGLang AsyncServer uses the async_generate interface of the SGLang engine, which is located on the first GPU of each TP group. Therefore, AsyncServer needs to remotely call async_generate through ray actor. +2. The vLLM AsyncServer uses the generate interface of the vLLM engine, which can communicate with the GPUs in the TP group through ZMQ and can be directly called in AsyncServer. + + +Usage Example +~~~~~~~~~~~~~ + +Follow :doc:`GSM8K example<../examples/gsm8k_example>` to prepare the dataset and model checkpoints. + +There are two options required to use agent loop: + +- `data.return_raw_chat=True` +- `actor_rollout_ref.rollout.mode=async` + +This example uses the sglang inference engine by default, and you can also modify rollout_name to use vllm. + +.. code-block:: bash + + bash examples/grpo_trainer/run_qwen3_8b_fsdp.sh + + +Multi-turn Conversations and Tool Calls +--------------------------------------- + +Follow :doc:`Multi-turn Rollout Support<../sglang_multiturn/multiturn>` to prepare tool and configuration files. + +The Tool Agent Loop has an additional requirement: adding an "agent_name" field to the dataset. During rollout, it will choose to use tool_agent_loop or single_turn_agent (default) based on this field. + +Usage Example +~~~~~~~~~~~~~ + +.. code-block:: bash + + # install mlflow to view toolcall and llm trace + pip install mlflow + + # This will download and preprocess the GSM8K dataset into ~/data/gsm8k/ and add the "agent_name" field. + python examples/data_preprocess/gsm8k_tool_agent_loop.py + + # Start training with tool calls and enabled mlflow based trace helping to debug the rollout details + bash examples/sglang_multiturn/run_qwen2_5_3b_gsm8k_tool_agent_mlflow_fsdp.sh + + # When training is done, start a mlflow server to view trace + mlflow ui -h 0.0.0.0 -p 5000 --backend-store-uri sqlite:////tmp/mlruns.db + + # then you can open http://:5000 from browser to view trace + + +Note: During training, because the model may sometimes fail to generate correct toolcall tags, an error message "Failed to decode tool call" will be output to the console, which does not indicate an abnormality in training. + + +Follow :doc:`Rollout trace<../advance/rollout_trace>` to known more about trace feature. + + + +Agent Framework +--------------- + +System Architecture +~~~~~~~~~~~~~~~~~~~ + +.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/langgraph_agent.png?raw=true + +System Components +~~~~~~~~~~~~~~~~~ + ++--------------------------+-----------------------------------------------------------------------------------------------+ +| Component | Role | ++==========================+===============================================================================================+ +| ChatModel | LLM object of LangChain, used to adapt to the “generate” api provided by LLMServerClient | ++--------------------------+-----------------------------------------------------------------------------------------------+ +| ReactAgentLoop | Agent adaptation layer, which by default supports a naive LangGraph Agentic. | +| | New classes can be derived to support user-defined Agents, and the run function needs to be | +| | implemented to complete Agent calls. | ++--------------------------+-----------------------------------------------------------------------------------------------+ +| AsyncServer | Server, each instance is connected to one DP group of the inference engine. | ++--------------------------+-----------------------------------------------------------------------------------------------+ + + +Follow doc "recipe/langgraph_agent/example/README.md" for more details. \ No newline at end of file diff --git a/verl/docs/start/install.rst b/verl/docs/start/install.rst new file mode 100644 index 0000000000000000000000000000000000000000..4b0f3a166615ab6be9acb4cb9a66c3b5e0fb921b --- /dev/null +++ b/verl/docs/start/install.rst @@ -0,0 +1,317 @@ +Installation +============ + +Requirements +------------ + +- **Python**: Version >= 3.10 +- **CUDA**: Version >= 12.8 + +verl supports various backends. Currently, the following configurations are available: + +- **FSDP** and **Megatron-LM** (optional) for training. +- **SGLang**, **vLLM** and **TGI** for rollout generation. + +Choices of Backend Engines +---------------------------- + +1. Training: + +We recommend using the **FSDP / FSDP2** backend to investigate, research and prototype different models, datasets and RL algorithms. For users who pursue better scalability, we recommend the **Megatron-LM** backend. Currently, we support `Megatron-LM v0.13.1 `_. Both backends are served through the same unified worker layer – see :doc:`Engine Workers<../workers/engine_workers>` for the worker-level API and :doc:`Model Engine<../workers/model_engine>` for the engine-level design. + + +2. Inference: + +For inference, vllm 0.8.3 and later versions have been tested for stability. We recommend turning on env var `VLLM_USE_V1=1` for optimal performance. + +For SGLang, refer to the :doc:`SGLang Backend<../workers/sglang_worker>` for detailed installation and usage instructions. SGLang rollout is under extensive development and offers many advanced features and optimizations. We encourage users to report any issues or provide feedback via the `SGLang Issue Tracker `_. + +For huggingface TGI integration, it is usually used for debugging and single GPU exploration. + +Install from docker image +------------------------- + +Start from v0.6.0, we use vllm and sglang release image as our base image. + +Base Image +:::::::::: + +- vLLM: https://hub.docker.com/r/vllm/vllm-openai +- SGLang: https://hub.docker.com/r/lmsysorg/sglang + +Application Image +::::::::::::::::: + +Upon base image, the following packages are added: + +- flash_attn +- Megatron-LM +- Apex +- TransformerEngine +- DeepEP + +Latest docker file: + +- `Dockerfile.stable.vllm `_ +- `Dockerfile.stable.sglang `_ + +All pre-built images are available in dockerhub: `verlai/verl `_. For example, ``verlai/verl:sgl055.latest``, ``verlai/verl:vllm011.latest``. + +You can find the latest images used for development and ci in our github workflows: + +- `.github/workflows/vllm.yml `_ +- `.github/workflows/sgl.yml `_ + + +Installation from Docker +:::::::::::::::::::::::: + +After pulling the desired Docker image and installing desired inference and training frameworks, you can run it with the following steps: + +1. Launch the desired Docker image and attach into it: + +.. code:: bash + + docker create --runtime=nvidia --gpus all --net=host --shm-size="10g" --cap-add=SYS_ADMIN -v .:/workspace/verl --name verl sleep infinity + docker start verl + docker exec -it verl bash + + +2. If you use the images provided, you only need to install verl itself without dependencies: + +.. code:: bash + + # install the nightly version (recommended) + git clone https://github.com/verl-project/verl && cd verl + pip3 install --no-deps -e . + +[Optional] If you hope to switch between different frameworks, you can install verl with the following command: + +.. code:: bash + + # install the nightly version (recommended) + git clone https://github.com/verl-project/verl && cd verl + pip3 install -e ".[vllm]" + pip3 install -e ".[sglang]" + + +Install from custom environment +--------------------------------------------- + +We recommend to use docker images for convenience. However, if your environment is not compatible with the docker image, you can also install verl in a python environment. + +.. note:: + + - Dockerfile provides more details than this installation instructions. You can find examples in each Dockerfile, for example `verl0.6-cu128-torch2.8.0-fa2.7.4 Dockerfile.base `_ . + + +Pre-requisites +:::::::::::::: + +For training and inference engines to utilize better and faster hardware support, CUDA/cuDNN and other dependencies are required, +and some of the dependencies are easy to be overridden when installing other packages, +so we put them in the :ref:`Post-installation` step. + +.. note:: + + - The installation steps below are recommended configurations for the latest version of verl. + + If you are trying to customize your own environment, please ignore the strict constraints. + +We need to install the following pre-requisites: + +- **CUDA**: Version >= 12.8 +- **cuDNN**: Version >= 9.10.0 +- **Apex** + +CUDA above 12.8 is recommended to use as the docker image, +please refer to `NVIDIA's official website `_ for other version of CUDA. + +.. code:: bash + + # change directory to anywher you like, in verl source code directory is not recommended + wget https://developer.download.nvidia.com/compute/cuda/12.8.1/local_installers/cuda-repo-ubuntu2204-12-8-local_12.8.1-570.124.06-1_amd64.deb + dpkg -i cuda-repo-ubuntu2204-12-8-local_12.8.1-570.124.06-1_amd64.deb + cp /var/cuda-repo-ubuntu2204-12-8-local/cuda-*-keyring.gpg /usr/share/keyrings/ + apt-get update + apt-get -y install cuda-toolkit-12-8 + update-alternatives --set cuda /usr/local/cuda-12-8 + + +cuDNN can be installed via the following command, +please refer to `NVIDIA's official website `_ for other version of cuDNN. + +.. code:: bash + + # change directory to anywher you like, in verl source code directory is not recommended + wget https://developer.download.nvidia.com/compute/cudnn/9.10.2/local_installers/cudnn-local-repo-ubuntu2204-9.10.2_1.0-1_amd64.deb + dpkg -i cudnn-local-repo-ubuntu2204-9.10.2_1.0-1_amd64.deb + cp /var/cudnn-local-repo-ubuntu2204-9.10.2/cudnn-*-keyring.gpg /usr/share/keyrings/ + apt-get update + apt-get -y install cudnn-cuda-12 + +Install dependencies +:::::::::::::::::::: + +.. note:: + + We recommend to use a fresh new conda environment to install verl and its dependencies. + + **Notice that the inference frameworks often strictly limit your pytorch version and will directly override your installed pytorch if not paying enough attention.** + + As a countermeasure, it is recommended to install inference frameworks first with the pytorch they needed. For vLLM, if you hope to use your existing pytorch, + please follow their official instructions + `Use an existing PyTorch installation `_ . + + +1. First of all, to manage environment, we recommend using conda: + +.. code:: bash + + conda create -n verl python==3.12 + conda activate verl + + +2. Then, execute the ``install.sh`` script that we provided in verl: + +.. code:: bash + + # Make sure you have activated verl conda env + # If you need to run with megatron + bash scripts/install_vllm_sglang_mcore.sh + # Or if you simply need to run with FSDP + USE_MEGATRON=0 bash scripts/install_vllm_sglang_mcore.sh + + +If you encounter errors in this step, please check the script and manually follow the steps in the script. + +[Optional] NVIDIA Apex is recommended for Megatron-LM training, but it's not needed if you only use FSDP backend. +You can install it via the following command, but notice that this steps can take a very long time. +It is recommended to set the ``MAX_JOBS`` environment variable to accelerate the installation process, +but do not set it too large, otherwise the memory will be overloaded and your machines may hang. + +.. code:: bash + + # change directory to anywher you like, in verl source code directory is not recommended + git clone https://github.com/NVIDIA/apex.git && \ + cd apex && \ + MAX_JOB=32 pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./ + +Install verl +:::::::::::: + +For installing the latest version of verl, the best way is to clone and +install it from source. Then you can modify our code to customize your +own post-training jobs. + +.. code:: bash + + git clone https://github.com/verl-project/verl.git + cd verl + pip install --no-deps -e . + + +Post-installation +::::::::::::::::: + +Please make sure that the installed packages are not overridden during the installation of other packages. + +The packages worth checking are: + +- **torch** and torch series +- **vLLM** +- **SGLang** +- **pyarrow** +- **tensordict** +- **nvidia-cudnn-cu12**: For Magetron backend + +If you encounter issues about package versions during running verl, please update the outdated ones. + + +Install with AMD GPUs - ROCM kernel support +------------------------------------------------------------------ + +When you run on AMD GPUs (MI300) with ROCM platform, you cannot use the previous quickstart to run verl. You should follow the following steps to build a docker and run it. +If you encounter any issues in using AMD GPUs running verl, feel free to contact me - `Yusheng Su `_. + +Find the docker for AMD ROCm: `docker/Dockerfile.rocm `_ +:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::: + +.. code-block:: bash + + # Build the docker in the repo dir: + # docker build -f docker/Dockerfile.rocm -t verl-rocm:03.04.2015 . + # docker images # you can find your built docker + FROM rocm/vllm:rocm6.2_mi300_ubuntu20.04_py3.9_vllm_0.6.4 + + # Set working directory + # WORKDIR $PWD/app + + # Set environment variables + ENV PYTORCH_ROCM_ARCH="gfx90a;gfx942" + + # Install vllm + RUN pip uninstall -y vllm && \ + rm -rf vllm && \ + git clone -b v0.6.3 https://github.com/vllm-project/vllm.git && \ + cd vllm && \ + MAX_JOBS=$(nproc) python3 setup.py install && \ + cd .. && \ + rm -rf vllm + + # Copy the entire project directory + COPY . . + + # Install dependencies + RUN pip install "tensordict<0.6" --no-deps && \ + pip install accelerate \ + codetiming \ + datasets \ + dill \ + hydra-core \ + liger-kernel \ + numpy \ + pandas \ + datasets \ + peft \ + "pyarrow>=15.0.0" \ + pylatexenc \ + "ray[data,train,tune,serve]" \ + torchdata \ + transformers \ + wandb \ + orjson \ + pybind11 && \ + pip install -e . --no-deps + +Build the image +:::::::::::::::::::::::: + +.. code-block:: bash + + docker build -t verl-rocm . + +Launch the container +:::::::::::::::::::::::::::: + +.. code-block:: bash + + docker run --rm -it \ + --device /dev/dri \ + --device /dev/kfd \ + -p 8265:8265 \ + --group-add video \ + --cap-add SYS_PTRACE \ + --security-opt seccomp=unconfined \ + --privileged \ + -v $HOME/.ssh:/root/.ssh \ + -v $HOME:$HOME \ + --shm-size 128G \ + -w $PWD \ + verl-rocm \ + /bin/bash + +If you do not want to root mode and require assign yourself as the user, +Please add ``-e HOST_UID=$(id -u)`` and ``-e HOST_GID=$(id -g)`` into the above docker launch script. + +verl with AMD GPUs currently supports FSDP as the training engine, vLLM and SGLang as the inference engine. We will support Megatron in the future. diff --git a/verl/docs/start/more_resources.rst b/verl/docs/start/more_resources.rst new file mode 100644 index 0000000000000000000000000000000000000000..aa8cb2a62b46579ee4bef2880d7f62485175495e --- /dev/null +++ b/verl/docs/start/more_resources.rst @@ -0,0 +1,7 @@ +More Resources +============== + +Last updated: 06/30/2025. + +- Introduction to verl (`Slides `_) +- verl Code Walkthrough (`Slides `_, `Talk in Chinese `_) diff --git a/verl/docs/start/multinode.rst b/verl/docs/start/multinode.rst new file mode 100644 index 0000000000000000000000000000000000000000..8148da24bd7e9bb489a648941d241745529ff226 --- /dev/null +++ b/verl/docs/start/multinode.rst @@ -0,0 +1,821 @@ +Multinode Training +================== + +Last updated: 06/10/2025. + +.. _wuxibin89: https://github.com/wuxibin89 + +Author: `Xibin Wu `_, `Yusheng Su `_. + +Option 1: Launch Manually +------------------------------ + +Set up multinode ray cluster +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +1. Start head node with ``ray start --head --dashboard-host=0.0.0.0``, there're 2 address you should care about: + +- GCS address: ``ray start --address=
``, where worker node should connect to. +- Dashboard address: ``
:8265``, where you should submit job to the cluster. + +.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/head.png?raw=true + +2. Start worker node with ``ray start --address=
`` you get above. + +.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/worker.png?raw=true + +3. Now you should see the cluster have 2 nodes with ``ray status``. + +.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/status.png?raw=true + +4. Additionally, you can access dashboard in the browser with the address you get above. + +*Firewall rules maybe need configure to access the dashboard, if there's any trouble, please contact your network administrator.* + +.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/overview.png?raw=true + +Submit job to ray cluster +~~~~~~~~~~~~~~~~~~~~~~~~~ +1. Submit ray job to cluster with the dashboard address you get above. + +.. code-block:: bash + + ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env=verl/trainer/runtime_env.yaml \ + --no-wait \ + -- \ + python3 -m verl.trainer.main_ppo \ + trainer.n_gpus_per_node=8 \ + trainer.nnodes=2 \ + ... + +.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/submit.png?raw=true + +2. Then you can check the job status with the following commands: + +- ray job list: list all jobs submitted to the cluster. +- ray job logs : query the logs of the job. +- ray job status : query the status of the job. +- ray job stop : request the job to be stopped. +- ray job list | grep submission_id | grep JobStatus | grep RUNNING | grep -oP 'raysubmit_[^'\''"]+' | head -n 1: get the latest job submission ID of the running job. +- ray job logs --follow: added ``--follow`` parameter to ray job logs command to enable continuous log streaming. + +3. You can also access driver/task/actor logs in ``/tmp/ray/session_latest/logs/``, driver log is ``job-driver-raysubmit_.log``. + +4. We strongly recommend you to view job detail from dashboard in multinode training, because it provide more structure way to view the job information. + +.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/job.png?raw=true +.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/job_detail.png?raw=true + +Option 2: Launch via SkyPilot on Kubernetes or clouds +------------------------------------------------------ + +.. note:: + Ready-to-use SkyPilot example configurations are available in the `examples/tutorial/skypilot/ `_ directory: + + - ``verl-ppo.yaml`` - PPO training with GSM8K dataset + - ``verl-grpo.yaml`` - GRPO training with MATH dataset + - ``verl-multiturn-tools.yaml`` - Multi-turn tool usage training + + See the `SkyPilot examples README `_ for detailed usage instructions. + +Step 1: Setup SkyPilot +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +SkyPilot can support different clouds, here we use GCP as example. `install skypilot `_ + +.. code-block:: bash + + conda create -y -n sky python=3.10 + conda activate sky + pip install "skypilot[gcp]" + + conda install -c conda-forge google-cloud-sdk + gcloud init + + # Run this if you don't have a credential file. + # This will generate ~/.config/gcloud/application_default_credentials.json. + gcloud auth application-default login + + # Check if the GCP credential is correctly setup. + sky check gcp + +.. image:: https://github.com/yottalabsai/open-source/blob/main/static/verl/setup_skypilot.png?raw=true + +Step 2: Prepare dataset +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. code-block:: bash + + git clone https://github.com/verl-project/verl.git + cd examples/data_preprocess + python3 gsm8k.py --local_save_dir ~/data/gsm8k + + +Step 3: Submit a job with SkyPilot +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +1. Create a SkyPilot YAML ``verl-cluster.yml`` with the following content: + +.. parsed-literal:: workdir: . will sync all the data in the current dir to the remote cluster. + +.. code-block:: yaml + + resources: + accelerators: L4:1 # every node has 1 L4 GPU + image_id: docker:verlai/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.0-fa2.7.4 + memory: 64+ # every node has 64 GB memory + ports: 8265 # expose port for ray dashboard + + num_nodes: 2 # cluster size + + # --------------- Work Directory Synchronization (workdir) --------------- + # Defines the local working directory to be synchronized to the remote cluster. + # Here, '.' means synchronizing the directory where the sky submit command is currently run. + workdir: . + + # --------------- (secrets) --------------- + secrets: + ## your wandb api key ## + WANDB_API_KEY: null + + # --------------- File Mounts/Data Upload (file_mounts) --------------- + # If your dataset (gsm8k folder) is local, it needs to be uploaded to the remote cluster. + file_mounts: + # Remote path (relative to remote user's home directory): Local path + # /remote/dir1/file: /local/dir1/file + data/gsm8k: ~/data/gsm8k + + # --------------- Environment Setup (setup) --------------- + # Commands run on each node of the remote cluster to set up the environment (e.g., install dependencies). These are run directly inside Docker. + setup: | + rm -rf verl + git clone https://github.com/verl-project/verl.git + cd verl + pip3 install -v -e .[vllm] + + # --------------- Run Command (run) --------------- + # The actual task commands to be executed on the remote cluster. + # This script will first start the Ray cluster (different ray start commands are executed on Head and Worker nodes). + # Then, your training script will only be run on the Head node (SKYPILOT_NODE_RANK == 0). + run: | + # Get the Head node's IP and total number of nodes (environment variables injected by SkyPilot). + head_ip=`echo "$SKYPILOT_NODE_IPS" | head -n1` + num_nodes=`echo "$SKYPILOT_NODE_IPS" | wc -l` # Here num_nodes should be equal to 2. + + # login wandb + python3 -c "import wandb; wandb.login(relogin=True, key='$WANDB_API_KEY')" + + # Start Ray based on node role (Head=0, Worker>0). + # This logic is a standard Ray cluster startup script. + if [ "$SKYPILOT_NODE_RANK" == "0" ]; then + # Head node starts Ray Head. + echo "Starting Ray head node..." + # Check if a Ray Head is already running to avoid duplicate starts. + ps aux | grep ray | grep 6379 &> /dev/null || ray start --head --disable-usage-stats \ + --port=6379 \ + --dashboard-host=0.0.0.0 \ + --dashboard-port=8265 + + # Wait for all worker nodes to join the cluster. + while [ $(ray nodes | grep NODE_ID | wc -l) -lt $num_nodes ]; do + echo "Waiting for all nodes to join... ($(ray nodes | grep NODE_ID | wc -l)/$num_nodes)" + sleep 5 + done + + # Head node executes the training script. + echo "Executing training script on head node..." + + python3 -m verl.trainer.main_ppo \ + data.train_files=data/gsm8k/train.parquet \ + data.val_files=data/gsm8k/test.parquet \ + data.train_batch_size=256 \ + data.max_prompt_length=512 \ + data.max_response_length=256 \ + actor_rollout_ref.model.path=Qwen/Qwen2.5-0.5B-Instruct \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.actor.ppo_mini_batch_size=64 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=1 \ + actor_rollout_ref.rollout.name=vllm \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \ + critic.optim.lr=1e-5 \ + critic.model.path=Qwen/Qwen2.5-0.5B-Instruct \ + critic.ppo_micro_batch_size_per_gpu=4 \ + algorithm.kl_ctrl.kl_coef=0.001 \ + trainer.logger=['console','wandb'] \ + trainer.val_before_train=False \ + trainer.default_hdfs_dir=null \ + trainer.n_gpus_per_node=1 \ + trainer.nnodes=2 \ + trainer.save_freq=20 \ + trainer.test_freq=20 \ + trainer.total_epochs=2 \ + trainer.project_name=verl_examples \ + trainer.experiment_name=experiment_name_gsm8k + + else + # Wait for Ray Head to start. + sleep 10 # Increase waiting time to ensure Head finishes starting. + # Worker node starts Ray Worker. + echo "Starting Ray worker node..." + + # Check if a Ray Worker is already running to avoid duplicate starts. + ps aux | grep ray | grep $head_ip:6379 &> /dev/null || ray start --address $head_ip:6379 --disable-usage-stats + + # Add sleep to after `ray start` to give ray enough time to daemonize + sleep 5 # Ensure Worker successfully connects to Head. + fi + + # No commands are added to the Worker node here; the Worker's main task is to start Ray and wait for the Head node to assign tasks. + echo "Node setup and Ray start script finished for rank $SKYPILOT_NODE_RANK." + + +.. code-block:: bash + + export WANDB_API_KEY= + sky launch -c verl --secret WANDB_API_KEY verl-cluster.yml + +.. image:: https://github.com/yottalabsai/open-source/blob/main/static/verl/running_job.png?raw=true +.. image:: https://github.com/yottalabsai/open-source/blob/main/static/verl/running_job_1.png?raw=true +.. image:: https://github.com/yottalabsai/open-source/blob/main/static/verl/finished.png?raw=true + +**Check the cluster on GCP** + +.. image:: https://github.com/yottalabsai/open-source/blob/main/static/verl/gcp_instances.png?raw=true + +**Check Ray Dashboard** + +We can see the cluster on the RAY Dashboard with the GCP head node: + +```console +$ sky status --endpoint 8265 verl +1.2.3.4:8265 +``` + +.. image:: https://github.com/yottalabsai/open-source/blob/main/static/verl/ray_dashboard_overview.png?raw=true +.. image:: https://github.com/yottalabsai/open-source/blob/main/static/verl/ray_dashboard_jobs.png?raw=true +.. image:: https://github.com/yottalabsai/open-source/blob/main/static/verl/ray_dashboard_cluster.png?raw=true + + +**Check the checkpoint of model** + +.. code-block:: bash + + # login the head node + ssh verl + # The global step will vary. Find the correct path from the training logs. + cd ~/sky_workdir/checkpoints/verl_examples/gsm8k/ + # Then list contents to find the checkpoint, e.g.: + ls -R . + +.. image:: https://github.com/yottalabsai/open-source/blob/main/static/verl/saved_model.png?raw=true + + +Option 3: Launch via Slurm +------------------------------ + +Ray provides users with `this `_ official +tutorial to start a Ray cluster on top of Slurm. We have verified the :doc:`GSM8K example<../examples/gsm8k_example>` +on a Slurm cluster under a multi-node setting with the following steps. + +1. [Optional] If your cluster support `Apptainer or Singularity `_ and you wish +to use it, convert verl's Docker image to an Apptainer image. Alternatively, set up the environment with the package +manager available on your cluster or use other container runtimes (e.g. through `Slurm's OCI support `_) available to you. + +.. code:: bash + + apptainer pull /your/dest/dir/vemlp-th2.4.0-cu124-vllm0.6.3-ray2.10-te1.7-v0.0.3.sif docker://verlai/verl:vemlp-th2.4.0-cu124-vllm0.6.3-ray2.10-te1.7-v0.0.3 + +2. Follow :doc:`GSM8K example<../examples/gsm8k_example>` to prepare the dataset and model checkpoints. + +3. Modify `examples/tutorial/slurm/ray_on_slurm.slurm `_ with your cluster's own information. + +4. Submit the job script to the Slurm cluster with `sbatch`. + +Please note that Slurm cluster setup may vary. If you encounter any issues, please refer to Ray's +`Slurm user guide `_ for common caveats. + +If you changed Slurm resource specifications, please make sure to update the environment variables in the job script if necessary. + + +Option 4: Launch via dstack +------------------------------ + +`dstackai/dstack `_ is an open-source container orchestrator that simplifies distributed training across cloud providers and on-premises environments +without the need to use K8S or Slurm. + +Prerequisite +~~~~~~~~~~~~ +Once dstack is `installed `_, initialize the directory as a repo with ``dstack init``. + +.. code-block:: bash + + mkdir myproject && cd myproject + dstack init + +**Create a fleet** + +Before submitting distributed training jobs, create a `dstack` `fleet `_. + +Run a Ray cluster task +~~~~~~~~~~~~~~~~~~~~~~ + +Once the fleet is created, define a Ray cluster task, e.g. in ``ray-cluster.dstack.yml``: + +.. code-block:: yaml + + type: task + name: ray-verl-cluster + + nodes: 2 + + env: + - WANDB_API_KEY + - PYTHONUNBUFFERED=1 + - CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 + + image: verlai/verl:app-verl0.6-transformers4.56.1-sglang0.5.2-mcore0.13.0-te2.2 + commands: + - git clone https://github.com/verl-project/verl + - cd verl + - pip install --no-deps -e . + - pip install hf_transfer hf_xet + - | + if [ $DSTACK_NODE_RANK = 0 ]; then + python3 examples/data_preprocess/gsm8k.py --local_save_dir ~/data/gsm8k + python3 -c "import transformers; transformers.pipeline('text-generation', model='Qwen/Qwen2.5-7B-Instruct')" + ray start --head --port=6379; + else + ray start --address=$DSTACK_MASTER_NODE_IP:6379 + fi + + # Expose Ray dashboard port + ports: + - 8265 + + resources: + gpu: 80GB:8 + shm_size: 128GB + + # Save checkpoints on the instance + volumes: + - /checkpoints:/checkpoints + +Now, if you run this task via `dstack apply`, it will automatically forward the Ray's dashboard port to `localhost:8265`. + +.. code-block:: bash + + dstack apply -f ray-cluster.dstack.yml + +As long as the `dstack apply` is attached, you can use `localhost:8265` to submit Ray jobs for execution + +Submit Ray jobs +~~~~~~~~~~~~~~~ + +Before you can submit Ray jobs, ensure to install `ray` locally: + +.. code-block:: shell + + pip install ray + +Now you can submit the training job to the Ray cluster which is available at ``localhost:8265``: + +.. code-block:: shell + + $ RAY_ADDRESS=http://localhost:8265 + $ ray job submit \ + -- python3 -m verl.trainer.main_ppo \ + data.train_files=/root/data/gsm8k/train.parquet \ + data.val_files=/root/data/gsm8k/test.parquet \ + data.train_batch_size=256 \ + data.max_prompt_length=512 \ + data.max_response_length=256 \ + actor_rollout_ref.model.path=Qwen/Qwen2.5-7B-Instruct \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.actor.ppo_mini_batch_size=64 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=1 \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \ + critic.optim.lr=1e-5 \ + critic.model.path=Qwen/Qwen2.5-7B-Instruct \ + critic.ppo_micro_batch_size_per_gpu=4 \ + algorithm.kl_ctrl.kl_coef=0.001 \ + trainer.project_name=ppo_training \ + trainer.experiment_name=qwen-2.5-7B \ + trainer.val_before_train=False \ + trainer.n_gpus_per_node=8 \ + trainer.nnodes=2 \ + trainer.default_local_dir=/checkpoints \ + trainer.save_freq=10 \ + trainer.test_freq=10 \ + trainer.total_epochs=15 2>&1 | tee verl_demo.log \ + trainer.resume_mode=disable + + +For more details on how `dstack` works, check out its `documentation `_. + +How to debug? +--------------------- + + +Ray Distributed Debugger VSCode Extension (Recommended) +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +1. Starting with Ray 2.39, Anyscale has introduced the `Ray Distributed Debugger `_ VSCode extension. Follow the extension’s installation instructions, then add your cluster using the dashboard URL you obtained earlier. + + .. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/debugger.png?raw=true + :alt: Ray Distributed Debugger VSCode extension screenshot + +2. Prerequisites. + + Ensure the following are installed (see the extension README for more detail): + + - Visual Studio Code + - `ray[default]` >= 2.9.1 + - `debugpy` >= 1.8.0 + + .. image:: https://github.com/aoshen524/verl/blob/main/docs/start/c7098b755ff689859837773a916c857.png?raw=true + :alt: VSCode with Ray prerequisites + +3. Environment Variables. + + To enable post‑mortem debugging, set: + + .. code-block:: bash + + export RAY_DEBUG_POST_MORTEM=1 + + .. admonition:: Note + :class: important + + Be sure to remove any legacy flags before starting Ray: + + - `RAY_DEBUG=legacy` + - `--ray-debugger-external` + +4. Configuring BreakpointsSet up breakpoint() in your code, and submit job to cluster. Then the extension will show the breakpoint information. + + + 1. Insert `breakpoint()` calls into your remote functions. + 2. Submit your job to the cluster. + + The extension will detect active breakpoints and display them in VSCode. + + .. image:: https://github.com/aoshen524/verl/blob/main/docs/start/4ddad74395c79a1402331c0ce73316f.png?raw=true + :alt: Detected breakpoint in VSCode + + **Note:** Breakpoints are only supported inside functions decorated with `@ray.remote`. + +5. Launching the Debugger. + + Run your job directly from the command line (do not use a `launch.json`): + + .. code-block:: bash + + python job.py + +6. Attaching to a Breakpoint. + + Once the process hits the first `breakpoint()`, click the Ray Distributed Debugger icon in the VSCode sidebar to attach the debugger. + + .. image:: https://github.com/aoshen524/verl/blob/main/docs/start/4ddad74395c79a1402331c0ce73316f.png?raw=true + :alt: Attaching VSCode debugger to Ray process + +7. Debugging With Multiple breakpoint(). + + For each subsequent task, first disconnect the current debugger session, then click the extension icon again to attach to the next breakpoint. + + .. image:: https://github.com/aoshen524/verl/blob/main/docs/start/6e83c910a62c82fecb89c6619e001cd.png?raw=true + :alt: Disconnecting and reconnecting the debugger + +Legacy Ray Debugger +~~~~~~~~~~~~~~~~~~~ +1. Ray has a builtin legacy `debugger `_ that allows you to debug your distributed applications. To enable debugger, start ray cluster with ``RAY_DEBUG=legacy`` and ``--ray-debugger-external``. + +.. code-block:: bash + + # start head node + RAY_DEBUG=legacy ray start --head --dashboard-host=0.0.0.0 --ray-debugger-external + # start worker node + RAY_DEBUG=legacy ray start --address='10.124.46.192:6379' --ray-debugger-external + +2. Set up breakpoint in your code, and submit job to cluster. Then run ``ray debug`` to wait breakpoint: + +.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/legacy.png?raw=true + + +Multi-node training on AMD clusters +--------------------------------------------------------------------------------------- + +If you want to run multi-node training with slurm with Docker/Podman container on AMD Cluster, you can use the following script. + +If you encounter any issues in using AMD GPUs running verl, please contact `Yusheng Su `_. + +.. note:: + 1. You need to use ``podman`` or ``docker`` in the following script. We will release the apptainer script later. + 2. If you want to use ``podman``, you just replace ``docker`` with ``podman`` in the following script. + +The script includes the following steps: + +1. SLURM Configuration +2. Environment Setup +3. Docker/Podman Container Setup +4. Ray Cluster Initialization +5. Data Preprocessing +6. Model Setup +7. Training Launch + + +slurm_script.sh +~~~~~~~~~~~~~~~~~~~~ + +.. code-block:: bash + + #!/bin/bash + + #SBATCH --job-name=verl-ray-on-slurm + #SBATCH --nodes=2 + #SBATCH --ntasks-per-node=2 + #SBATCH --mem=200G + #SBATCH --time=30-00:00:00 + #SBATCH --gpus-per-node=8 + #SBATCH --cpus-per-task=28 + #SBATCH --output=../verl_log/slurm-%j.out + #SBATCH --error=../verl_log/slurm-%j.err + #SBATCH --nodelist=gpu-[0,1] + + + # load necessary modules + ### Run this setup + # [Cluster]: Use docker + # docker pull docker.io/rocm/vllm:rocm6.2_mi300_ubuntu20.04_py3.9_vllm_0.6.4 + + + ########################################################################## + ###The following setting should be set in different project and cluster### + ########################################################################## + + ### Project + CONTAINER_NAME="multinode_verl_training" + IMG="verl.rocm" + DOCKERFILE="docker/Dockerfile.rocm" + # echo $PWD + verl_workdir="${HOME}/projects/verl_upstream" + export TRANSFORMERS_CACHE="${HOME}/.cache/huggingface" + export HF_HOME=$TRANSFORMERS_CACHE + + ### Cluster Network Setting + export NCCL_DEBUG=TRACE + export GPU_MAX_HW_QUEUES=2 + export TORCH_NCCL_HIGH_PRIORITY=1 + export NCCL_CHECKS_DISABLE=1 + # export NCCL_IB_HCA=rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7 + export NCCL_IB_HCA=mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_8,mlx5_9 + export NCCL_IB_GID_INDEX=3 + export NCCL_CROSS_NIC=0 + export CUDA_DEVICE_MAX_CONNECTIONS=1 + export NCCL_PROTO=Simple + export RCCL_MSCCL_ENABLE=0 + export TOKENIZERS_PARALLELISM=false + export HSA_NO_SCRATCH_RECLAIM=1 + ########################################################################## + + ### For rocm and training script + export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 + export ROCR_VISIBLE_DEVICES=$HIP_VISIBLE_DEVICES + export CUDA_VISIBLE_DEVICES=$HIP_VISIBLE_DEVICES + + + # Build and launch the Docker container + srun bash -c " + # Exit on any error + set -e + + # Clean up dangling images (images with tag) + docker image prune -f + + # Need to pull the docker first + docker pull docker.io/rocm/vllm:rocm6.2_mi300_ubuntu20.04_py3.9_vllm_0.6.4 + + if ! docker images --format "{{.Repository}}:{{.Tag}}" | grep -q "${IMG}"; then + echo \"Building ${IMG} image...\" + docker build -f \"${DOCKERFILE}\" -t \"${IMG}\" . + else + echo \"${IMG} image already exists, skipping build\" + fi + + # Removing old container if exists + docker rm \"${CONTAINER_NAME}\" 2>/dev/null || true + + # Checking network devices + ibdev2netdev + + # Launch the docker + docker run --rm -d \ + -e HYDRA_FULL_ERROR=1 \ + -e HIP_VISIBLE_DEVICES=${HIP_VISIBLE_DEVICES} \ + -e ROCR_VISIBLE_DEVICES=${ROCR_VISIBLE_DEVICES} \ + -e CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} \ + -e NCCL_DEBUG=${NCCL_DEBUG} \ + -e GPU_MAX_HW_QUEUES=${GPU_MAX_HW_QUEUES} \ + -e TORCH_NCCL_HIGH_PRIORITY=${TORCH_NCCL_HIGH_PRIORITY} \ + -e NCCL_CHECKS_DISABLE=${NCCL_CHECKS_DISABLE} \ + -e NCCL_IB_HCA=${NCCL_IB_HCA} \ + -e NCCL_IB_GID_INDEX=${NCCL_IB_GID_INDEX} \ + -e NCCL_CROSS_NIC=${NCCL_CROSS_NIC} \ + -e CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS} \ + -e NCCL_PROTO=${NCCL_PROTO} \ + -e RCCL_MSCCL_ENABLE=${RCCL_MSCCL_ENABLE} \ + -e TOKENIZERS_PARALLELISM=${TOKENIZERS_PARALLELISM} \ + -e HSA_NO_SCRATCH_RECLAIM=${HSA_NO_SCRATCH_RECLAIM} \ + -e TRANSFORMERS_CACHE=${TRANSFORMERS_CACHE} \ + -e HF_HOME=${HF_HOME} \ + --network host \ + --device /dev/dri \ + --device /dev/kfd \ + --device /dev/infiniband \ + --group-add video \ + --cap-add SYS_PTRACE \ + --security-opt seccomp=unconfined \ + --privileged \ + -v \${HOME}:\${HOME} \ + -v \${HOME}/.ssh:/root/.ssh \ + -w "${verl_workdir}" \ + --shm-size 128G \ + --name \"${CONTAINER_NAME}\" \ + \"${IMG}\" \ + tail -f /dev/null + + echo \"Container setup completed\" + " + # (Optional): If you do not want to root mode and require assign yuorself as the user + # Please add `-e HOST_UID=$(id -u)` and `-e HOST_GID=$(id -g)` into the above docker launch script. + + + + + + ### Ray launch the nodes before training + + # Getting the node names + nodes_array=($(scontrol show hostnames "$SLURM_JOB_NODELIST" | tr '\n' ' ')) + + head_node=${nodes_array[0]} + head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address) + + # if we detect a space character in the head node IP, we'll + # convert it to an ipv4 address. This step is optional. + if [[ "$head_node_ip" == *" "* ]]; then + IFS=' ' read -ra ADDR <<<"$head_node_ip" + if [[ ${#ADDR[0]} -gt 16 ]]; then + head_node_ip=${ADDR[1]} + else + head_node_ip=${ADDR[0]} + fi + echo "IPV6 address detected. We split the IPV4 address as $head_node_ip" + fi + + port=6379 + ip_head=$head_node_ip:$port + export ip_head + echo "IP Head: $ip_head" + + # make sure we set environment variables before Ray initialization + + # Print out all env variables + printenv + + echo "Starting HEAD at $head_node" + srun --nodes=1 --ntasks=1 -w "$head_node" \ + docker exec "${CONTAINER_NAME}" \ + ray start --head --node-ip-address="$head_node_ip" --port=$port \ + --dashboard-port=8266 \ + --num-cpus "${SLURM_CPUS_PER_TASK}" --num-gpus "${SLURM_GPUS_PER_NODE}" --block & + # optional, though may be useful in certain versions of Ray < 1.0. + sleep 10 + + # number of nodes other than the head node + worker_num=$((SLURM_JOB_NUM_NODES - 1)) + + for ((i = 1; i <= worker_num; i++)); do + node_i=${nodes_array[$i]} + echo "Debug: Starting worker on node_i = ${node_i}" + if [ -z "$node_i" ]; then + echo "Error: Empty node name for worker $i" + continue + fi + echo "Starting WORKER $i at $node_i" + srun --nodes=1 --ntasks=1 -w "$node_i" \ + docker exec "${CONTAINER_NAME}" \ + ray start --address "$ip_head" --num-cpus "${SLURM_CPUS_PER_TASK}" --num-gpus "${SLURM_GPUS_PER_NODE}" --block & + sleep 5 + done + + + + + # Ray initlization test (See whether any error in the above execution) + echo "Testing Ray initialization in the slurm nodes..." + docker exec "${CONTAINER_NAME}" python3 -c ' + import ray + try: + ray.init(address="auto") + print("\n=== Ray Cluster Status ===") + print(f"Number of nodes: {len(ray.nodes())}") + for node in ray.nodes(): + print("Node: {}, Status: {}".format(node["NodeManagerHostname"], node["Alive"])) + # print(f"Node: {node}") + ray.shutdown() + print("Ray initialization successful!") + except Exception as e: + print(f"Ray initialization failed: {str(e)}") + ' + echo "=== Ray test completed ===" + ###### + + + + # Run data preprocessing + + echo "Starting data preprocessing..." + docker exec "${CONTAINER_NAME}" \ + python3 "examples/data_preprocess/gsm8k.py" "--local_save_dir" "../data/gsm8k" + + echo "Starting data preprocessing..." + docker exec "${CONTAINER_NAME}" \ + python3 "examples/data_preprocess/math_dataset.py" "--local_dir" "../data/math" + + train_files="../data/gsm8k/train.parquet" + val_files="../data/gsm8k/test.parquet" + + # Download and test model + echo "Loading model..." + docker exec "${CONTAINER_NAME}" \ + python3 -c "import transformers; transformers.pipeline('text-generation', model='Qwen/Qwen2-7B-Instruct')" + MODEL_PATH="Qwen/Qwen2-7B-Instruct" + + # Set model path after pipeline test + MODEL_PATH="Qwen/Qwen2.5-0.5B-Instruct" + + echo "== Data and model loading Done ==" + + echo "Start to train..." + + docker exec "${CONTAINER_NAME}" \ + python3 -c "import transformers; transformers.pipeline('text-generation', model='Qwen/Qwen2-7B-Instruct')" + MODEL_PATH="Qwen/Qwen2-7B-Instruct" + + + PYTHONUNBUFFERED=1 srun --overlap --nodes=${SLURM_NNODES} --ntasks=1 -w "$head_node" \ + docker exec "${CONTAINER_NAME}" \ + python3 -m verl.trainer.main_ppo \ + data.train_files=$train_files \ + data.val_files=$val_files \ + data.train_batch_size=1024 \ + data.max_prompt_length=1024 \ + data.max_response_length=1024 \ + actor_rollout_ref.model.path=$MODEL_PATH \ + actor_rollout_ref.model.enable_gradient_checkpointing=False \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.model.use_remove_padding=True \ + actor_rollout_ref.actor.ppo_mini_batch_size=256 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=8 \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.fsdp_config.param_offload=False \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=2 \ + actor_rollout_ref.rollout.name=vllm \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.9 \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \ + actor_rollout_ref.ref.fsdp_config.param_offload=True \ + critic.optim.lr=1e-5 \ + critic.model.use_remove_padding=True \ + critic.model.path=$MODEL_PATH \ + critic.model.enable_gradient_checkpointing=False \ + critic.ppo_micro_batch_size_per_gpu=8 \ + critic.fsdp.param_offload=False \ + critic.fsdp.optimizer_offload=False \ + algorithm.kl_ctrl.kl_coef=0.0001 \ + trainer.critic_warmup=0 \ + trainer.logger='["console","wandb"]' \ + trainer.project_name='verl_example' \ + trainer.experiment_name='Qwen2.5-32B-Instruct_function_rm' \ + trainer.n_gpus_per_node=${SLURM_GPUS_PER_NODE} \ + trainer.val_before_train=False \ + trainer.nnodes=${SLURM_NNODES} \ + trainer.save_freq=-1 \ + trainer.test_freq=10 \ + trainer.total_epochs=15 + + +Run multi-node training with above slurm_script.sh +~~~~~~~~~~~~~~~~~~~~ +Just sbatch your slurm_script.sh + +.. code-block:: bash + + sbatch slurm_script.sh + diff --git a/verl/docs/start/quickstart.rst b/verl/docs/start/quickstart.rst new file mode 100644 index 0000000000000000000000000000000000000000..492fded9cafcd4452ca8fad336676f58d274cebd --- /dev/null +++ b/verl/docs/start/quickstart.rst @@ -0,0 +1,151 @@ +.. _quickstart: + +========================================================= +Quickstart: PPO training on GSM8K dataset +========================================================= + +Post-train a LLM using GSM8K dataset. + +Introduction +------------ + +.. _hf_dataset_gsm8k: https://huggingface.co/datasets/openai/gsm8k + +In this example, we train an LLM to tackle the `GSM8k `_ task with function-based rewards. [1]_ + +Prerequisite: + +- the latest version of ``verl`` and its dependencies installed following the installation guide. Using the docker image is recommended. + +- a GPU with at least 24 GB HBM + + +Dataset Introduction +-------------------- + +GSM8k is a math problem dataset. The prompt is an elementary school +problem. The LLM model is asked to solve the math problem. Below is an example: + +Prompt + + Katy makes coffee using teaspoons of sugar and cups of water in the + ratio of 7:13. If she used a total of 120 teaspoons of sugar and cups + of water, calculate the number of teaspoonfuls of sugar she used. + +Solution + + The total ratio representing the ingredients she used to make the + coffee is 7+13 = <<7+13=20>>20 Since the fraction representing the + number of teaspoons she used is 7/20, she used 7/20\ *120 = + <<7/20*\ 120=42>>42 #### 42 + +Step 1: Prepare the dataset +---------------------------- + +We preprocess the dataset in parquet format so that (1) it contains necessary fields for computing RL rewards and (2) is faster to read. + +.. code-block:: bash + + python3 examples/data_preprocess/gsm8k.py --local_save_dir ~/data/gsm8k + +Step 2: Download a model for post-training +------------------------------------------- + +In this example, we start with the ``Qwen2.5-0.5B-Instruct`` model. + +If you want to perform SFT before RL, refer to the :doc:`Complete GSM8K Example<../examples/gsm8k_example>`, the `sft directory `_ and `SFT Trainer `_ for further details. + +.. code-block:: bash + + python3 -c "import transformers; transformers.pipeline('text-generation', model='Qwen/Qwen2.5-0.5B-Instruct')" + +Step 3: Perform PPO training with the instruct model +---------------------------------------------------------------------- + +**Reward Model/Function** + +We use a pre-defined rule-based reward model. We force the model to produce a final +answer following 4 “#” as shown in the solution. We extract the final +answer from both the solution and model's output using regular +expression matching. We assign a reward of 1 to correct +answer, 0.0 to incorrect answer and 0 to no answer. + +For more details, please refer to `verl/utils/reward_score/gsm8k.py `_. + +**Training Script** + +Now let's run PPO training with the dataset and model above. [2]_ + + +Set the ``data.train_files`` ,\ ``data.val_files``, ``actor_rollout_ref.model.path`` and ``critic.model.path`` based on your dataset and model names or paths. +You may set ``VERL_USE_MODELSCOPE=True`` to download models from `modelscope `_ instead of `huggingface `_. + +.. code-block:: bash + + PYTHONUNBUFFERED=1 python3 -m verl.trainer.main_ppo \ + data.train_files=$HOME/data/gsm8k/train.parquet \ + data.val_files=$HOME/data/gsm8k/test.parquet \ + data.train_batch_size=256 \ + data.max_prompt_length=512 \ + data.max_response_length=512 \ + actor_rollout_ref.model.path=Qwen/Qwen2.5-0.5B-Instruct \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.actor.ppo_mini_batch_size=64 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \ + actor_rollout_ref.rollout.name=vllm \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=1 \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \ + critic.optim.lr=1e-5 \ + critic.model.path=Qwen/Qwen2.5-0.5B-Instruct \ + critic.ppo_micro_batch_size_per_gpu=4 \ + algorithm.kl_ctrl.kl_coef=0.001 \ + trainer.logger=console \ + trainer.val_before_train=False \ + trainer.n_gpus_per_node=1 \ + trainer.nnodes=1 \ + trainer.save_freq=10 \ + trainer.test_freq=10 \ + trainer.total_epochs=15 2>&1 | tee verl_demo.log + +You are expected to see the following logs, indicating training in progress. The key metric ``val/test_score/openai/gsm8k`` is computed every ``trainer.test_freq`` steps: + +.. code-block:: bash + + step:0 - timing/gen:21.470 - timing/ref:4.360 - timing/values:5.800 - actor/reward_kl_penalty:0.000 - actor/reward_kl_penalty_coeff:0.001 - timing/adv:0.109 - timing/update_critic:15.664 - critic/vf_loss:14.947 - critic/vf_clipfrac:0.000 - critic/vpred_mean:-2.056 - critic/grad_norm:1023.278 - critic/lr(1e-4):0.100 - timing/update_actor:20.314 - actor/entropy_loss:0.433 - actor/pg_loss:-0.005 - actor/pg_clipfrac:0.000 - actor/ppo_kl:0.000 - actor/grad_norm:1.992 - actor/lr(1e-4):0.010 - critic/score/mean:0.004 - critic/score/max:1.000 - critic/score/min:0.000 - critic/rewards/mean:0.004 - critic/rewards/max:1.000 - critic/rewards/min:0.000 - critic/advantages/mean:-0.000 - critic/advantages/max:2.360 - critic/advantages/min:-2.280 - critic/returns/mean:0.003 - critic/returns/max:0.000 - critic/returns/min:0.000 - critic/values/mean:-2.045 - critic/values/max:9.500 - critic/values/min:-14.000 - response_length/mean:239.133 - response_length/max:256.000 - response_length/min:77.000 - prompt_length/mean:104.883 - prompt_length/max:175.000 - prompt_length/min:68.000 + step:1 - timing/gen:23.020 - timing/ref:4.322 - timing/values:5.953 - actor/reward_kl_penalty:0.000 - actor/reward_kl_penalty:0.001 - timing/adv:0.118 - timing/update_critic:15.646 - critic/vf_loss:18.472 - critic/vf_clipfrac:0.384 - critic/vpred_mean:1.038 - critic/grad_norm:942.924 - critic/lr(1e-4):0.100 - timing/update_actor:20.526 - actor/entropy_loss:0.440 - actor/pg_loss:0.000 - actor/pg_clipfrac:0.002 - actor/ppo_kl:0.000 - actor/grad_norm:2.060 - actor/lr(1e-4):0.010 - critic/score/mean:0.000 - critic/score/max:0.000 - critic/score/min:0.000 - critic/rewards/mean:0.000 - critic/rewards/max:0.000 - critic/rewards/min:0.000 - critic/advantages/mean:0.000 - critic/advantages/max:2.702 - critic/advantages/min:-2.616 - critic/returns/mean:0.000 - critic/returns/max:0.000 - critic/returns/min:0.000 - critic/values/mean:-2.280 - critic/values/max:11.000 - critic/values/min:-16.000 - response_length/mean:232.242 - response_length/max:256.000 - response_length/min:91.000 - prompt_length/mean:102.398 - prompt_length/max:185.000 - prompt_length/min:70.000 + +Checkout ``Algorithm Baselines`` page for full training and validation logs for reference. + +The checkpoint is saved at the following dir by default: ``checkpoints/${trainer.project_name}/${trainer.experiment_name}``. You can merge the saved checkpoints to huggingface model using ``verl.model_merger`` module, for example: + +.. code-block:: bash + + python3 -m verl.model_merger merge \ + --backend fsdp \ + --local_dir checkpoints/${trainer.project_name}/${trainer.experiment_name}/global_step_1/actor \ + --target_dir checkpoints/${trainer.project_name}/${trainer.experiment_name}/global_step_1/actor/huggingface + +For more details about checkpoint and model merging, please refer to :ref:`checkpoint-page`. + +To enable ``wandb`` for experiment tracking, set the following configs: + +.. code-block:: bash + + trainer.logger='["console","wandb"]' \ + trainer.project_name=$YOUR_PROJECT_NAME \ + trainer.experiment_name=$YOUR_RUN_NAME \ + +If you encounter out of memory issues with HBM less than 32GB, enable the following configs would help: + +.. code-block:: bash + + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \ + critic.ppo_micro_batch_size_per_gpu=1 \ + +For the full set of configs, please refer to :ref:`config-explain-page` for detailed explanation and performance tuning. + + +.. [1] The original paper (https://arxiv.org/pdf/2110.14168) mainly focuses on training a verifier (a reward model) to solve math problems via Best-of-N sampling. In this example, we train an RL agent using a rule-based reward model. +.. [2] More training script examples for FSDP and Megatron-LM backend are stored in `examples/ppo_trainer `_ directory. diff --git a/verl/docs/start/ray_debug_tutorial.rst b/verl/docs/start/ray_debug_tutorial.rst new file mode 100644 index 0000000000000000000000000000000000000000..9e7c87dfaee0c04f24bdb6921717b8068d1ee6a2 --- /dev/null +++ b/verl/docs/start/ray_debug_tutorial.rst @@ -0,0 +1,96 @@ +Ray Debug Tutorial +================== + +Last updated: 04/23/2025 + + +.. _wuxibin89: https://github.com/wuxibin89 + +Author: `Ao Shen `_. + +How to debug? +--------------------- + + +Ray Distributed Debugger VSCode Extension (Recommended) +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +1. Starting with Ray 2.39, Anyscale has introduced the `Ray Distributed Debugger `_ VSCode extension. Follow the extension’s installation instructions, then add your cluster using the dashboard URL you obtained earlier. + + .. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/debugger.png?raw=true + :alt: Ray Distributed Debugger VSCode extension screenshot + +2. Prerequisites. + + Ensure the following are installed (see the extension README for more detail): + + - Visual Studio Code + - `ray[default]` >= 2.9.1 + - `debugpy` >= 1.8.0 + + .. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/readme.png?raw=true + :alt: VSCode with Ray prerequisites + +3. Environment Variables. + + To enable post‑mortem debugging, set: + + .. code-block:: bash + + export RAY_DEBUG_POST_MORTEM=1 + + .. admonition:: Note + :class: important + + Be sure to remove any legacy flags before starting Ray: + + - `RAY_DEBUG=legacy` + - `--ray-debugger-external` + +4. Configuring BreakpointsSet up breakpoint() in your code, and submit job to cluster. Then the extension will show the breakpoint information. + + + 1. Insert `breakpoint()` calls into your remote functions. + 2. Submit your job to the cluster. + + The extension will detect active breakpoints and display them in VSCode. + + **Note:** Breakpoints are only supported inside functions decorated with `@ray.remote`. + +5. Launching the Debugger. + + Run your job directly from the command line (do not use a `launch.json`): + + .. code-block:: bash + + python job.py + +6. Attaching to a Breakpoint. + + Once the process hits the first `breakpoint()`, click the Ray Distributed Debugger icon in the VSCode sidebar to attach the debugger. + + .. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/launch.png?raw=true + :alt: Attaching VSCode debugger to Ray process + +7. Debugging With Multiple breakpoint(). + + For each subsequent task, first disconnect the current debugger session, then click the extension icon again to attach to the next breakpoint. + + .. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/disconnect.png?raw=true + :alt: Disconnecting and reconnecting the debugger + +Legacy Ray Debugger +~~~~~~~~~~~~~~~~~~~ +1. Ray has a builtin legacy `debugger `_ that allows you to debug your distributed applications. To enable debugger, start ray cluster with ``RAY_DEBUG=legacy`` and ``--ray-debugger-external``. + +.. code-block:: bash + + # start head node + RAY_DEBUG=legacy ray start --head --dashboard-host=0.0.0.0 --ray-debugger-external + # start worker node + RAY_DEBUG=legacy ray start --address='10.124.46.192:6379' --ray-debugger-external + +2. Set up breakpoint in your code, and submit job to cluster. Then run ``ray debug`` to wait breakpoint: + +.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/legacy.png?raw=true + diff --git a/verl/docs/workers/automodel_workers.rst b/verl/docs/workers/automodel_workers.rst new file mode 100644 index 0000000000000000000000000000000000000000..e12bee5270469ee4f41d8e4174a7739b17a4e582 --- /dev/null +++ b/verl/docs/workers/automodel_workers.rst @@ -0,0 +1,65 @@ +Automodel Backend +================= + +Last updated: 03/07/2026. + +We support the Automodel (nemo_automodel) backend by implementing the +``AutomodelEngine`` and ``AutomodelEngineWithLMHead`` engine classes. +The Automodel backend delegates model building, parallelization, optimizer +sharding, LR scheduling, gradient clipping, and checkpointing to +nemo_automodel's infrastructure while using verl's training loop, +data pipeline, and loss function. + +**Requirements** + +- Automodel r0.3.0 +- transformers v5.0.0 + +**Pros** + +- Supports FSDP2 and TP distributed strategies out of + the box. + +- Native support for Mixture-of-Experts (MoE) models with Expert + Parallelism (EP) via DeepEP. + +- TransformerEngine (TE) integration for optimized attention, linear + layers, and RMSNorm. + +- Readily supports any HuggingFace model without checkpoint conversion. + +**Cons** + +- Pipeline parallelism is not yet supported. + + +SFT Examples +------------ + +We provide example SFT training scripts using the Automodel backend in +`examples/sft/gsm8k/ `_. + +Basic: Qwen2.5-0.5B with FSDP2 +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +A minimal example using ``Qwen/Qwen2.5-0.5B-Instruct`` with FSDP2 and +no parallelism: + +.. code:: shell + + bash examples/sft/gsm8k/run_qwen2_5_0_5b_automodel.sh 4 /tmp/automodel_sft_test + +See `run_qwen2_5_0_5b_automodel.sh `_. + +Advanced: Qwen3-30B MoE with Expert Parallelism +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +A larger-scale example using ``Qwen/Qwen3-30B-A3B-Base`` (MoE model) +with Expert Parallelism (EP=8), DeepEP, TransformerEngine backend, and +torch_mm experts backend: + +.. code:: shell + + bash examples/sft/gsm8k/run_qwen3_30b_automodel.sh 8 /tmp/automodel_sft_30b + +See `run_qwen3_30b_automodel.sh `_. diff --git a/verl/docs/workers/engine_workers.rst b/verl/docs/workers/engine_workers.rst new file mode 100644 index 0000000000000000000000000000000000000000..31e5a6b5e01932a86e55925e818b0d30b54f7344 --- /dev/null +++ b/verl/docs/workers/engine_workers.rst @@ -0,0 +1,246 @@ +Engine Workers +============== + +Last updated: 04/20/2026. + +:mod:`verl.workers.engine_workers` provides the worker-layer classes that +``RayWorkerGroup`` instantiates for PPO / GRPO / SFT style RL training. +They are **engine agnostic** – FSDP, FSDP2, Megatron-LM, Automodel, +TorchTitan and VeOmni are all wired in through the same entry points. +The specific backend is selected at runtime from ``actor.strategy`` / +``critic.strategy`` and resolved by +:class:`verl.workers.engine.EngineRegistry`. + +For the engine-layer design (how ``BaseEngine`` subclasses implement +``forward_step``, parallelism, checkpointing, weight export, etc.) see +:doc:`model_engine`. + +Class Hierarchy +--------------- + +:: + + ActorRolloutRefWorker # hybrid worker, co-locates actor + rollout + optional ref + ├── self.actor : TrainingWorker (built if role contains "actor") + ├── self.ref : TrainingWorker (built if role contains "ref") + ├── self.rollout: BaseRollout (vLLM / SGLang, built if role contains "rollout") + └── self.checkpoint_engine (built if role contains "actor") + + TrainingWorker # generic "one engine + optimizer + profiler" worker + └── self.engine : BaseEngine (fsdp / fsdp2 / megatron / automodel / veomni / torchtitan) + +``TrainingWorker`` is also used standalone for the critic, reference +model, reward model and SFT / DPO training – it's essentially a +Ray-wrapped ``BaseEngine`` with a Tinker-like API +(https://thinkingmachines.ai/tinker/) exposed as RPCs. + +ActorRolloutRefWorker +--------------------- + +:class:`verl.workers.engine_workers.ActorRolloutRefWorker` is the +hybrid worker used for actor, rollout and (optional) reference policy. +The ``role`` argument selects which sub-workers are constructed: + +========================= =========================================================================== +role What is built inside ``init_model`` +========================= =========================================================================== +``actor`` ``self.actor`` (``TrainingWorker``) + checkpoint engine +``rollout`` ``self.rollout`` (``BaseRollout``) +``ref`` ``self.ref`` (``TrainingWorker`` with ``forward_only`` engine config) +``actor_rollout`` actor + rollout + checkpoint engine (most common for colocated PPO) +``actor_rollout_ref`` all three +========================= =========================================================================== + +Key RPCs +^^^^^^^^ + +1. ``init_model`` + + .. code:: python + + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def init_model(self): + + ``ONE_TO_ALL``: the driver calls ``init_model`` and the same routine + runs on every worker. It builds the ``TrainingWorker`` (which in turn + builds the ``BaseEngine`` via ``EngineRegistry.new``), the rollout + engine, and the checkpoint engine used for trainer→rollout weight + sync. + +2. ``compute_log_prob`` / ``compute_ref_log_prob`` + + .. code:: python + + @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="actor")) + def compute_log_prob(self, data: TensorDict) -> TensorDict: + return self.actor.infer_batch(data) + + @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="ref")) + def compute_ref_log_prob(self, data: TensorDict) -> TensorDict: + return self.ref.infer_batch(data) + + ``TrainingWorker.infer_batch`` drives ``BaseEngine.infer_batch`` (eval + mode + ``no_grad``). The n-d dispatch function is built from the + engine's actual parallel topology, so Megatron's PP dimension is + surfaced as an extra DP dimension to the single controller without + needing a backend-specific dispatch mode. + +3. ``update_actor`` + + .. code:: python + + @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="actor")) + def update_actor(self, data: TensorDict) -> TensorDict: + return self.actor.train_mini_batch(data=data) + + ``train_mini_batch`` splits the batch into mini-batches, iterates + over PPO epochs, and calls ``TrainingWorker.train_batch`` for each + mini-batch (one optimizer step per mini-batch). The PPO loss + or distillation loss is wired by ``init_model`` via + ``TrainingWorker.set_loss_fn``. + +4. ``update_weights`` + + .. code:: python + + @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False) + async def update_weights(self, global_steps: int = None): + + Push the freshest trainer weights to the rollout engine. + + - For **colocated sync training** (``checkpoint_engine.backend == + "naive"``): export per-tensor parameters via + ``engine.get_per_tensor_param`` and call ``rollout.update_weights`` + directly. LoRA adapters are merged into base weights up-front when + ``model.lora.merge=True``. + - For **disaggregated async training**: send the weights through + ``self.checkpoint_engine.send_weights`` instead. + +5. ``save_checkpoint`` / ``load_checkpoint`` + + Both delegate to the actor ``TrainingWorker``, which in turn calls + ``BaseEngine.save_checkpoint`` / ``load_checkpoint``. The backend + engine is responsible for sharded model + optimizer + scheduler state + (and HuggingFace export when applicable). + +TrainingWorker +-------------- + +:class:`verl.workers.engine_workers.TrainingWorker` is the generic +worker for a single engine + optimizer + profiler. It is used: + +- As ``self.actor`` / ``self.ref`` inside ``ActorRolloutRefWorker``. +- As the critic / reward worker (via ``add_critic_worker`` / + ``add_reward_model_worker`` in ``verl/trainer/main_ppo.py``). +- Standalone for SFT / DPO training. + +Construction takes a single +:class:`verl.workers.config.TrainingWorkerConfig` which bundles the +``model_config``, ``engine_config``, ``optimizer_config``, +``checkpoint_config`` and ``profiler_config``. The backend is chosen +from ``engine_config.strategy`` (``fsdp``, ``fsdp2``, ``megatron``, +``automodel``, ``veomni``, ``torchtitan``). + +Key RPCs +^^^^^^^^ + +- ``reset()`` – first call initializes the engine; subsequent calls + reload weights and reset optimizer / scheduler state. +- ``to(device, model=True, optimizer=True, grad=True)`` – manual + load/offload control. ``device`` must be either ``"cpu"`` or + ``"device"`` (which is mapped to the actual accelerator name). +- ``set_loss_fn(loss_fn)`` – install the loss closure (PPO loss, + distillation loss, or any custom callable that + accepts ``(model_output, batch)``). +- ``train_mini_batch(data)`` – mini-batch + PPO-epoch loop; one + optimizer step per mini-batch; allgather metrics across DP. +- ``train_batch(data)`` – single mini-batch train step. Usually invoked + indirectly via ``train_mini_batch``. +- ``infer_batch(data)`` – forward-only step used for log-prob / value / + reward / distillation-teacher computation. Supports + ``no_lora_adapter=True`` to temporarily disable the adapter at + inference. +- ``save_checkpoint`` / ``load_checkpoint`` – delegate to + ``BaseEngine``. + +Backend Selection +----------------- + +Set the ``strategy`` field on ``actor.engine`` / ``critic.engine`` / +``ref.engine`` in your Hydra config: + +.. code-block:: yaml + + actor_rollout_ref: + actor: + strategy: fsdp2 # or: fsdp, megatron, automodel, veomni, torchtitan + engine: + strategy: fsdp2 + param_offload: False + # ... + +The ``EngineRegistry`` dispatches on ``(model_type, backend, device)`` – +for example ``(language_model, fsdp2, cuda)`` or +``(language_model, megatron, npu)``: + +===================== ===================== ===================== ============================================================= +model_type backend device Engine class +===================== ===================== ===================== ============================================================= +``language_model`` ``fsdp`` / ``fsdp2`` ``cuda`` / ``npu`` ``verl.workers.engine.fsdp.FSDPEngineWithLMHead`` +``language_model`` ``megatron`` ``cuda`` ``verl.workers.engine.megatron.MegatronEngineWithLMHead`` +``language_model`` ``megatron`` ``npu`` ``verl.workers.engine.mindspeed.MindspeedEngineWithLMHead`` +``language_model`` ``mindspeed_llm`` ``npu`` ``verl.workers.engine.mindspeed.MindSpeedLLMEngineWithLMHead`` +``language_model`` ``automodel`` ``cuda`` ``verl.workers.engine.automodel.AutomodelEngineWithLMHead`` +``language_model`` ``veomni`` ``cuda`` / ``npu`` ``verl.workers.engine.veomni.VeOmniEngineWithLMHead`` +``language_model`` ``torchtitan`` ``cuda`` / ``npu`` ``verl.workers.engine.torchtitan.TorchTitanEngineWithLMHead`` +``value_model`` ``fsdp`` / ``fsdp2`` ``cuda`` / ``npu`` ``verl.workers.engine.fsdp.FSDPEngineWithValueHead`` +``value_model`` ``megatron`` ``cuda`` ``verl.workers.engine.megatron.MegatronEngineWithValueHead`` +===================== ===================== ===================== ============================================================= + +Migrating from Legacy Workers +----------------------------- + +The legacy ``verl.workers.fsdp_workers`` / ``verl.workers.megatron_workers`` +modules (together with ``verl.workers.actor`` / ``verl.workers.critic`` +/ ``verl.workers.sharding_manager`` / ``verl.workers.legacy``) have been +removed. The table below summarises the equivalent entry points: + +============================================================== ========================================================================= +Legacy (removed) Current (``verl.workers.engine_workers``) +============================================================== ========================================================================= +``verl.workers.fsdp_workers.ActorRolloutRefWorker`` ``ActorRolloutRefWorker`` (``strategy=fsdp``/``fsdp2``) +``verl.workers.megatron_workers.ActorRolloutRefWorker`` ``ActorRolloutRefWorker`` (``strategy=megatron``) +``verl.workers.fsdp_workers.CriticWorker`` ``TrainingWorker`` (with critic config + value-model engine) +``verl.workers.megatron_workers.CriticWorker`` ``TrainingWorker`` (with critic config + value-model engine) +``verl.workers.actor.DataParallelPPOActor`` ``FSDPEngineWithLMHead`` + ``TrainingWorker`` +``verl.workers.actor.MegatronPPOActor`` ``MegatronEngineWithLMHead`` + ``TrainingWorker`` +``verl.workers.critic.DataParallelPPOCritic`` ``FSDPEngineWithValueHead`` + ``TrainingWorker`` +``verl.workers.critic.MegatronPPOCritic`` ``MegatronEngineWithValueHead`` + ``TrainingWorker`` +``verl.workers.sharding_manager.FSDPUlyssesShardingManager`` ``verl.utils.ulysses.FSDPUlyssesShardingManager`` +``Dispatch.MEGATRON_PP_AS_DP_PROTO`` ``make_nd_compute_dataproto_dispatch_fn(mesh_name=...)`` (derived from engine) +``use_legacy_worker_impl: True`` (removed; only the unified engine is available) +============================================================== ========================================================================= + +Extending +--------- + +To add a new backend, implement a ``BaseEngine`` subclass under +``verl/workers/engine//`` and register it with +``@EngineRegistry.register(model_type=..., backend=...)``. The worker +layer (``TrainingWorker`` / ``ActorRolloutRefWorker``) is already +engine-agnostic and will pick up the new backend as soon as +``engine_config.strategy`` is set accordingly. See :doc:`model_engine` +for the detailed extension guide and the test harness under +``tests/special_e2e/sft/``. + +Source +------ + +- :mod:`verl.workers.engine_workers` – + `engine_workers.py `__ +- :mod:`verl.workers.engine` – + `engine/ `__ +- :mod:`verl.workers.rollout` – + `rollout/ `__ +- Driver-side PPO glue – + `verl/trainer/main_ppo.py `__ diff --git a/verl/docs/workers/model_engine.rst b/verl/docs/workers/model_engine.rst new file mode 100644 index 0000000000000000000000000000000000000000..53653e581444cc61bc5f8dc9f9e28869d8859273 --- /dev/null +++ b/verl/docs/workers/model_engine.rst @@ -0,0 +1,125 @@ +Model Engine +============ + +.. _vermouth: https://github.com/vermouth1992 + +Author: `Chi Zhang `_ + +Last updated: 09/25/2025. + +Current Support Matrix +---------------------- + ++----------+-----------+--------------+-------------+--------------------------+ +| Backends | Model | Scalability | Model | Pain points | +| | Supported | | Definition | | +| | | | | | ++==========+===========+==============+=============+==========================+ +| FSDP | Day 1 | - Dense is OK| Huggingface | Monkey patch can be | +| + | support | | + monkey | easily impacted by | +| ulysses | HF model | - MoE is bad | patch | transformers version | ++----------+-----------+--------------+-------------+--------------------------+ +| MCore | Limited | Best | GPTModel | Supporting new models is | +| | | | (One model | difficult | +| | | | for all) | | ++----------+-----------+--------------+-------------+--------------------------+ + +- We monkey patch attention function to support ulysses +- We monkey patch VLM models to support FSDP with mixed data with and + without images + +Class Hierarchy +--------------- + +Note that all the workers and trainers run in **SPMD** mode. SFT/DPO/RM +trainer is directly invoked by ``torchrun``. The Actor/Critic worker can +also be invoked by a RayWorkerGroup and provides APIs to a single +controller. + +- Base Engine level: implement model init, optimizer init, lr scheduler + init, sharding, checkpoint manager. +- Full Engine level: subclass base engine and implement + ``forward_step``. +- Worker/SPMD trainer level: **engine agnostic**, implement training + logics using abstract engine APIs + +RL trainer utilizes workers to construct HybridFlow program. This is out +of the scope of model engine. + +Existing Model Types +-------------------- + +========== ====================== ====================== +Model type Language model Value model +========== ====================== ====================== +Input text/image/video/audio text/image/video/audio +Output logits for next token logits as value +========== ====================== ====================== + +Currently, we have two model types: language model and value model. We +expect to expand the category to include Qwen-Omni family (output both +text and audio) and VLA models. + +Data Format +----------- + +Currently, verl adopts left-right padding data format in RL trainer. +This creates massive padding when the discrepancy between response +length is large. We will start to implement no-padding format throughout +the whole system. + +.. image:: https://github.com/vermouth1992/verl-data/blob/master/images/data_format.png?raw=true + :alt: Data Format + +Here is the migration plan: +- Implement no-padding format in engine +- Add a transformation layer in Actor/Critic worker. +- Replace Actor/Critic Worker in RL trainer +- Implement no-padding throughput system + +Checkpoint System +----------------- + +.. image:: https://github.com/vermouth1992/verl-data/blob/master/images/verl-ckpt.png?raw=true + :alt: Model Engine Checkpoint System + +The engine constructs the model using huggingface config, then load +weights from huggingface checkpoint. If the engine directly uses +huggingface model definition, it can use function provided by +``transformers``. Otherwise, each engine has to write their own +checkpoint load logic (e.g., +`mbridge `__). During model +training, each engine has to implement save_checkpoint and +load_checkpoint that save/load intermediate sharded checkpoint including +model, optimizer and lr scheduler states. Each engine has to implement a +checkpoint merge script, that merges the intermediate sharded checkpoint +back to huggingface format. + +API +--- + +A tentative model engine API can be found: +https://github.com/verl-project/verl/blob/main/verl/workers/engine/base.py#L24 + +Extension +--------- + +Add a new backend +~~~~~~~~~~~~~~~~~ + +- Start a new folder under ``verl/workers/engine``. Then, implement + ``transformer_impl.py``. If you want to implement a non-transformer + model, please contact us in advance. +- Add the engine config to the GSM8k SFT trainer script: + https://github.com/verl-project/verl/blob/main/tests/special_e2e/sft/run_sft_engine_gsm8k.sh +- Invoke the tests with your backend: + https://github.com/verl-project/verl/blob/main/tests/special_e2e/sft/test_sft_engine_all.sh. + This test script will run various backends and various + configurations, and compare the loss and grad norm of the first step + to make sure they are close. + +Add a new model type +~~~~~~~~~~~~~~~~~~~~ + +- This is mainly reserved for models whose the output is not just text + (e.g., Qwen3-Omni). Please discuss with us before you proceed. diff --git a/verl/docs/workers/ray_trainer.rst b/verl/docs/workers/ray_trainer.rst new file mode 100644 index 0000000000000000000000000000000000000000..5b085ce59068fc2f45a055b1cce25145614c551b --- /dev/null +++ b/verl/docs/workers/ray_trainer.rst @@ -0,0 +1,241 @@ +PPO Ray Trainer +=============== + +Last updated: 02/12/2025. + +We implement the RayPPOTrainer, which is a trainer runs on the driver +process on a single CPU/GPU node (default is CPU). + +The PPORayTrainer include 3 core functions for data preparation, +WorkerGroup initialization and PPO training loop. + +Data Preparation +---------------- + +The ``PPORayTrainer``, as a single process, is responsible for loading a +complete batch of samples (prompts) from the dataset and then dispatch +to different worker_groups running on different GPUs. + +To generalize the data loading, we implement the ``RLHFDataset`` class +to load the preprocessed parquet files, apply chat templates to the +prompts, add padding, truncate prompts that exceed max prompt length and +then tokenize. + +.. code:: python + + self.train_dataset = RLHFDataset(data_files=self.config.data.train_files, + tokenizer=self.tokenizer, + config=self.config.data) + +Then, the dataloader will iterate the dataset under PPO mini batch size. + +WorkerGroup Initialization +-------------------------- + +We first introduce a basic implementation of initializing the +``WorkerGroup`` of the actor model on a given set of GPUs. + +.. code:: python + + # max_colocate_count means the number of WorkerGroups (i.e. processes) in each RayResourcePool + # For FSDP backend, we recommend using max_colocate_count=1 that merge all WorkerGroups into one. + # For Megatron backend, we recommend using max_colocate_count>1 that can utilize different WorkerGroup for differnt models + resource_pool = RayResourcePool(process_on_nodes=[config.trainer.n_gpus_per_node] * config.trainer.nnodes, + use_gpu=True, + max_colocate_count=1) + # define actor rollout cls to be init on remote + actor_rollout_cls = RayClassWithInitArgs(cls=ActorRolloutWorker) + # define actor_rollout worker group + actor_rollout_worker_group = MegatronRayWorkerGroup(resource_pool=resource_pool, + ray_cls_with_init=actor_rollout_cls, + default_megatron_kwargs=config.actor_rollout.megatron) + +Different WorkerGroups, like ``actor_rollout_worker_group`` , +``critic_worker_group`` and ``ref_worker_group`` lies on a separate +process in the above implementation. + +The driver process can then call the distributed compute function within +the ``actor_rollout_worker_group`` and other roles to construct the RL +training loop. + +For models colocated in the same set of GPUs, we further provide a +fine-grain optimization, which merge the ``worker_group`` of different roles +in the same process. This optimization can save the redundant +CUDA/distributed context in different processes. + +.. code:: python + + # initialize WorkerGroup + # NOTE: if you want to use a different resource pool for each role, which can support different parallel size, + # you should not use `create_colocated_worker_cls`. Instead, directly pass different resource pool to different worker groups. + # See TODO(url) for more information. + all_wg = {} + for resource_pool, class_dict in self.resource_pool_to_cls.items(): + worker_dict_cls = create_colocated_worker_cls(class_dict=class_dict) + wg_dict = self.ray_worker_group_cls(resource_pool=resource_pool, ray_cls_with_init=worker_dict_cls) + spawn_wg = wg_dict.spawn(prefix_set=class_dict.keys()) + all_wg.update(spawn_wg) + + if self.use_critic: + self.critic_wg = all_wg['critic'] + self.critic_wg.init_model() + + if self.use_reference_policy: + self.ref_policy_wg = all_wg['ref'] + self.ref_policy_wg.init_model() + + if self.use_rm: + self.rm_wg = all_wg['rm'] + self.rm_wg.init_model() + + # we should create rollout at the end so that vllm can have a better estimation of kv cache memory + self.actor_rollout_wg = all_wg['actor_rollout'] + self.actor_rollout_wg.init_model() + +.. note:: For megatron backend, if we merge the ``worker_groups`` into the same processes, all the roles will utilize the same 3D parallel size. To optimize this, we may need to maintain several 3D process groups for each role in the same distributed context. If you want to use different 3D parallel size for different roles, please follow the similar architecture of the first code block to initialize each role's ``worker_group`` + + +PPO Training Loop +----------------- + +We implement the PPO training loop by calling the functions in +worker_group of each role. The input and output data of each function is +a ``DataProto`` object implemented in `protocol.py `_. In the training +loop, trainer will dispatch/collect the data to/from different GPUs +following the transfer protocols wrapped in the workers' functions. The +computation of PPO micro batches is processed in ``update_actor`` and +``update_critic`` functions. + +To extend to other RLHF algorithms, such as DPO, GRPO, please refer to +:doc:`../advance/dpo_extension`. + +.. code:: python + + def fit(self): + """ + The training loop of PPO. + The driver process only need to call the compute functions of the worker group through RPC to construct the PPO dataflow. + The light-weight advantage computation is done on the driver process. + """ + from verl.utils.tracking import Tracking + from omegaconf import OmegaConf + + logger = Tracking(project_name=self.config.trainer.project_name, + experiment_name=self.config.trainer.experiment_name, + default_backend=self.config.trainer.logger, + config=OmegaConf.to_container(self.config, resolve=True)) + + global_steps = 0 + + # perform validation before training + # currently, we only support validation using the reward_function. + if self.val_reward_fn is not None: + val_metrics = self._validate() + pprint(f'Initial validation metrics: {val_metrics}') + + for epoch in range(self.config.trainer.total_epochs): + for batch_dict in self.train_dataloader: + metrics = {} + + batch: DataProto = DataProto.from_single_dict(batch_dict) + # batch = batch.to('cuda') + + # pop those keys for generation + gen_batch = batch.pop(batch_keys=['input_ids', 'attention_mask', 'position_ids']) + + # generate a batch + with Timer(name='gen', logger=None) as timer: + gen_batch_output = self.actor_rollout_wg.generate_sequences(gen_batch) + metrics['timing/gen'] = timer.last + + batch = batch.union(gen_batch_output) + + if self.use_reference_policy: + # compute reference log_prob + with Timer(name='ref', logger=None) as timer: + ref_log_prob = self.ref_policy_wg.compute_ref_log_prob(batch) + batch = batch.union(ref_log_prob) + metrics['timing/ref'] = timer.last + + # compute values + with Timer(name='values', logger=None) as timer: + values = self.critic_wg.compute_values(batch) + batch = batch.union(values) + metrics['timing/values'] = timer.last + + with Timer(name='adv', logger=None) as timer: + # compute scores. Support both model and function-based. + # We first compute the scores using reward model. Then, we call reward_fn to combine + # the results from reward model and rule-based results. + if self.use_rm: + # we first compute reward model score + reward_tensor = self.rm_wg.compute_rm_score(batch) + batch = batch.union(reward_tensor) + + # we combine with rule-based rm + reward_tensor = self.reward_fn(batch) + batch.batch['token_level_scores'] = reward_tensor + + # compute rewards. apply_kl_penalty if available + batch, kl_metrics = apply_kl_penalty(batch, + kl_ctrl=self.kl_ctrl_in_reward, + kl_penalty=self.config.algorithm.kl_penalty) + metrics.update(kl_metrics) + + # compute advantages, executed on the driver process + batch = compute_advantage(batch, + self.config.algorithm.gamma, + self.config.algorithm.lam, + adv_estimator=self.config.algorithm.adv_estimator) + metrics['timing/adv'] = timer.last + + # update critic + if self.use_critic: + with Timer(name='update_critic', logger=None) as timer: + critic_output = self.critic_wg.update_critic(batch) + metrics['timing/update_critic'] = timer.last + critic_output_metrics = reduce_metrics(critic_output.meta_info['metrics']) + metrics.update(critic_output_metrics) + + # implement critic warmup + if self.config.trainer.critic_warmup <= global_steps: + # update actor + with Timer(name='update_actor', logger=None) as timer: + actor_output = self.actor_rollout_wg.update_actor(batch) + metrics['timing/update_actor'] = timer.last + actor_output_metrics = reduce_metrics(actor_output.meta_info['metrics']) + metrics.update(actor_output_metrics) + + # validate + if self.val_reward_fn is not None and (global_steps + 1) % self.config.trainer.test_freq == 0: + with Timer(name='testing', logger=None) as timer: + val_metrics: dict = self._validate() + val_metrics = {f'val/{key}': val for key, val in val_metrics.items()} + metrics['timing/testing'] = timer.last + metrics.update(val_metrics) + + # collect metrics + data_metrics = compute_data_metrics(batch=batch) + metrics.update(data_metrics) + + # TODO: make a canonical logger that supports various backend + logger.log(data=metrics, step=global_steps) + + if self.config.trainer.save_freq > 0 and (global_steps + 1) % self.config.trainer.save_freq == 0: + actor_local_path = os.path.join(self.config.trainer.default_local_dir, 'actor', + f'global_step_{global_steps}') + actor_remote_path = os.path.join(self.config.trainer.default_hdfs_dir, 'actor') + self.actor_rollout_wg.save_checkpoint(actor_local_path, actor_remote_path) + + if self.use_critic: + critic_local_path = os.path.join(self.config.trainer.default_local_dir, 'critic', + f'global_step_{global_steps}') + critic_remote_path = os.path.join(self.config.trainer.default_hdfs_dir, 'critic') + self.critic_wg.save_checkpoint(critic_local_path, critic_remote_path) + + global_steps += 1 + + # perform validation after training + if self.val_reward_fn is not None: + val_metrics = self._validate() + pprint(f'Final validation metrics: {val_metrics}') diff --git a/verl/docs/workers/sglang_worker.rst b/verl/docs/workers/sglang_worker.rst new file mode 100644 index 0000000000000000000000000000000000000000..6edb8871730d3e9bd99d6a04e21f725baa2a3a74 --- /dev/null +++ b/verl/docs/workers/sglang_worker.rst @@ -0,0 +1,237 @@ +SGLang Backend +============== + +Last updated: 05/31/2025. + +**Authored By SGLang RL Team and listed alphabetically by last name** + +`Jingyi Chen `_, `Yitong Guan `_, `Zhuobin Huang `_, `Jiajun Li `_, `Ji Li `_, `Shenggui Li `_, `Junrong Lin `_, `Xiang Long `_, `Rui Lu `_, `Jin Pan `_, `Shuai Shi `_, `Yushen Su `_, `Xinyuan Tong `_, `Chendong Wang `_, `Hanchen Zhang `_, `Haoran Wang `_, `Yongan Xiang `_, `Chengxing Xie `_, `Yuhao Yang `_, `Jinwei Yao `_, `Qiaolin Yu `_, `Yuzhen Zhou `_, `Chenyang Zhao `_ + + + +Introduction +------------ +`SGLang `_ is an open-source state-of-the-art inference service engine, fully adopted by xAI to support all inference needs of Grok during research and serving processes. + +Currently, verl fully supports using SGLang as the inference engine during the rollout phase. As a rollout engine, SGLang provides the same feature coverage as vLLM., including memory saving and multi-node rollout features. After installing verl and SGLang, simply add ``actor_rollout_ref.rollout.name=sglang`` at startup script to seamlessly switch between the two inference frameworks. + +In addition, the SGLang team is actively working on supporting features such as Multi-Turn Agentic RL, VLM RLHF, Server-Based RLHF, and Partial Rollout. You can track the related development progress in the `Tracking Roadmap `_. + +Installation +------------ +Please always follow the following command to install SGLang with verl. + +.. code-block:: bash + + pip install --upgrade pip + # Currently 0.4.8, subject to updates at any time, please refer to the latest version specified in `setup.py` + pip install -e ".[sglang]" + +You can check the following dependencies are in your environment: + +.. note:: + + - **PyTorch**: 2.6.0+cu124 + - **CUDA**: 12.4 + - **flashinfer-python**: 0.2.5+cu124torch2.6 + - **SGLang**: 0.4.6.post5 + - **sgl-kernel**: 0.1.4 + +Using SGLang as the Inference Backend for PPO Training on a Single Machine +------------------------------------------------------------------------- +We use Qwen/Qwen2-7B-Instruct on the gsm8k dataset for a simple test. + +1. Run the following command to prepare the gsm8k dataset: + +.. code-block:: bash + + python3 examples/data_preprocess/gsm8k.py + +2. Run the following script to conduct a PPO experiment on a single machine with 4 GPUs: + +.. code-block:: bash + + export SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK=True + PYTHONUNBUFFERED=1 python3 -m verl.trainer.main_ppo \ + data.train_files=$HOME/data/gsm8k/train.parquet \ + data.val_files=$HOME/data/gsm8k/test.parquet \ + data.train_batch_size=4096 \ + data.max_prompt_length=4096 \ + data.max_response_length=4096 \ + actor_rollout_ref.rollout.name=sglang \ + actor_rollout_ref.model.path=Qwen/Qwen2-7B-Instruct \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.actor.ppo_mini_batch_size=64 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.fsdp_config.param_offload=True \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \ + actor_rollout_ref.rollout.tensor_model_parallel_size=2 \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \ + critic.optim.lr=1e-5 \ + critic.model.path=Qwen/Qwen2-7B-Instruct \ + critic.ppo_micro_batch_size_per_gpu=4 \ + critic.fsdp.param_offload=True \ + critic.fsdp.optimizer_offload=True \ + algorithm.kl_ctrl.kl_coef=0.001 \ + trainer.logger=console \ + trainer.val_before_train=False \ + trainer.n_gpus_per_node=4 \ + trainer.nnodes=1 \ + trainer.save_freq=-1 \ + trainer.test_freq=10 \ + trainer.total_epochs=15 2>&1 | tee verl_demo.log + +Why export SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK? +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +1. ``verl`` initializes a ``SGLangRollout`` module during rollout, which is used to evaluate/generate samples. + +2. ``SGLangRollout`` will initialize ``Engine``, and further initialize a ``torch.distributed.DeviceMesh``, used to support Tensor Parallel (TP). + +3. ``DeviceMesh.init()`` internally checks the free GPU memory of all participating devices. If the difference is too large (more than ~10%), it directly reports an error to avoid initialization failures or deadlocks. + +Why might there be inconsistent GPU memory? +""""""""""""""""""""""""""""""""""""""""""" + +**1. Ray Distributed Actor loads the model at different times** + +``verl`` uses Ray-based multi-process, multi-GPU concurrent training. Each ``WorkerDict`` may be called at different times: + +.. code-block:: python + + self.rollout = SGLangRollout(...) + +Different workers initialize the model at different times → different memory usage. + +**2. Delayed initialization causes memory bias** + +Some workers start model loading/inference (e.g., ``generate_sequences()``, ``compute_log_prob()``) earlier than others. +Early workers already use up GPU memory → late workers still have empty memory → memory difference appears. + +**3. SGLang's TP init uses "all-device broadcast", but there's no uniform release timing** + +Although ``SGLangRollout`` may only involve subset of GPUs, its ``Engine`` initialization calls ``torch.distributed.init_process_group()`` and broadcasts weights, so: + +- Non-rollout GPUs also join the communication. +- Later on, ``DeviceMesh`` init will fail due to "inconsistent memory". + +**4. Different FSDP/TP loading behaviors also lead to mismatch** + +If using: + +.. code-block:: bash + + actor.fsdp_config.param_offload=True + ref.fsdp_config.param_offload=True + +Then some workers keep params on CPU while others already sharded to GPU → leads to asymmetric memory layout. + +Using SGLang as the Inference Backend for PPO Training Across Multiple Machines +------------------------------------------------------------------------------ +SGLang also supports running verl's RAY-based cross-machine inference in IPv4 and IPv6 scenarios. In the script below, we use TP=16 for cross-machine inference. Suppose we have two interconnected machines: node0 with IP 10.94.16.4 and node1 with IP 10.94.16.5. + +1. Start Ray on node0: + +.. code-block:: bash + + ray start --head --dashboard-host=0.0.0.0 + +You will see the following prompt: + +.. code-block:: bash + + Usage stats collection is enabled. To disable this, add `--disable-usage-stats` to the command that starts the cluster, or run the following command: `ray disable-usage-stats` before starting the cluster. See https://docs.ray.io/en/master/cluster/usage-stats.html for more details. + + Local node IP: 10.94.16.4 + + -------------------- + Ray runtime started. + -------------------- + + Next steps + To add another node to this Ray cluster, run + ray start --address='10.94.16.4:6379' + +2. Have node1 join the Ray cluster: + +Run the following command on node1: + +.. code-block:: bash + + ray start --address='10.94.16.4:6379' + +Run the following command to confirm that the Ray cluster now has two nodes: + +.. code-block:: bash + + ray status + +You can see that the cluster has two nodes with 16 GPUs: + +.. code-block:: bash + + ======== Autoscaler status: 2025-04-09 09:25:37.694016 ======== + Node status + --------------------------------------------------------------- + Active: + 1 node_ef382ffd687d8f6b060c1b68e63ada7341b936fe5b1901dd04de1027 + 1 node_1eb4d7d07e793114c23a89d1a41f1f76acf6ef5b35af844a4ee8e4ba + Pending: + (no pending nodes) + Recent failures: + (no failures) + + Resources + --------------------------------------------------------------- + Usage: + 0.0/360.0 CPU + 0.0/16.0 GPU + 0B/3.39TiB memory + 0B/372.53GiB object_store_memory + +3. Run the following script to train meta-llama/Llama-3.1-8B-Instruct with TP=16 across 2 machines using 16 GPUs: + +.. code-block:: bash + + DATA_DIR=$HOME/data/gsm8k + + python3 -m verl.trainer.main_ppo \ + actor_rollout_ref.rollout.name=sglang \ + data.train_files=$DATA_DIR/train.parquet \ + data.val_files=$DATA_DIR/test.parquet \ + data.train_batch_size=4096 \ + data.max_prompt_length=4096 \ + data.max_response_length=4096 \ + actor_rollout_ref.model.path=meta-llama/Llama-3.1-8B-Instruct \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.model.use_remove_padding=True \ + actor_rollout_ref.actor.ppo_mini_batch_size=64 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.fsdp_config.param_offload=True \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=16 \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \ + actor_rollout_ref.rollout.free_cache_engine=True \ + actor_rollout_ref.ref.log_prob_micro_batch_size=16 \ + actor_rollout_ref.ref.fsdp_config.param_offload=True \ + critic.optim.lr=1e-5 \ + critic.model.use_remove_padding=True \ + critic.model.path=meta-llama/Llama-3.1-8B-Instruct \ + critic.model.enable_gradient_checkpointing=True \ + critic.ppo_micro_batch_size=16 \ + critic.fsdp.param_offload=True \ + critic.fsdp.optimizer_offload=True \ + algorithm.kl_ctrl.kl_coef=0.001 \ + trainer.critic_warmup=0 \ + trainer.logger=console \ + trainer.val_before_train=True \ + trainer.n_gpus_per_node=8 \ + trainer.nnodes=2 \ + trainer.save_freq=-1 \ + trainer.test_freq=10 \ + trainer.total_epochs=15 2>&1 | tee verl_demo.log diff --git a/verl/docs/workers/trtllm_worker.rst b/verl/docs/workers/trtllm_worker.rst new file mode 100644 index 0000000000000000000000000000000000000000..752b02bd1a01b0ab6da5aaba8ff2d9047e4d6ff9 --- /dev/null +++ b/verl/docs/workers/trtllm_worker.rst @@ -0,0 +1,74 @@ +TensorRT-LLM Backend +==================== + +Last updated: 5/6/2026. + +**Authored By NVIDIA TensorRT-LLM Team** + +Introduction +------------ +`TensorRT-LLM `_ is a high-performance LLM inference engine with state-of-the-art optimizations for NVIDIA GPUs. +The verl integration of TensorRT-LLM is based on TensorRT-LLM's `Ray orchestrator `_, with more features and performance optimizations to come. + +- For **synchronous training**, the TensorRT-LLM rollout adopts a mixed design combining aspects of the hybrid engine and colocated mode, instead of relying purely on standard colocated mode. +- For **asynchronous training**, the TensorRT-LLM rollout follows other rollout backends and uses standalone mode for trainer and rollout placement. + +TensorRT-LLM rollout supports the following key features, primarily tested on Qwen3 dense and MoE variants: + +- Synchronous training (GRPO, DAPO, etc.) +- Cross-node inference +- FP8 refit +- Asynchronous training (further optimizations planned) +- Preliminary support for VLM + +You can track our roadmap and share feedback at the `TensorRT-LLM rollout roadmap `_. + + +Installation +------------ +We recommend using `docker/Dockerfile.stable.trtllm `_ for building a docker image with TensorRT-LLM pre-installed. The verl integration is supported from ``nvcr.io/nvidia/tensorrt-llm/release:1.2.0rc6``, and you can choose other TensorRT-LLM versions via ``TRTLLM_BASE_IMAGE`` from the `NGC Catalog `_. The image is updated periodically to track TensorRT-LLM's weekly releases. + +Alternatively, refer to the `TensorRT-LLM installation guide `_ for compatible environments if you want to build your own. + +Install verl with TensorRT-LLM: + +.. code-block:: bash + + pip install --upgrade pip + pip install -e ".[trtllm]" + +.. note:: + + Using the TensorRT-LLM rollout requires setting the following environment variables before launching the Ray cluster. These have been included in all the example scripts: + + .. code-block:: bash + + # Clean all SLURM/MPI/PMIx env to avoid PMIx mismatch error. + for v in $(env | awk -F= '/^(PMI|PMIX|MPI|OMPI|SLURM)_/{print $1}'); do + unset "$v" + done + +Using TensorRT-LLM rollout for GRPO +------------------------------------ + +.. code-block:: bash + + ## For FSDP training engine + INFER_BACKEND=trtllm bash examples/grpo_trainer/run_qwen3_8b_fsdp.sh + ## For Megatron-Core training engine + INFER_BACKEND=trtllm bash examples/grpo_trainer/run_qwen3_8b_megatron.sh + +Using TensorRT-LLM rollout for DAPO with FP8 +--------------------------------------------- + +.. code-block:: bash + + # For Megatron-Core training engine with FP8 rollout + INFER_BACKEND=trtllm ROLLOUT_QUANTIZATION=fp8 bash examples/grpo_trainer/run_qwen3_30b_a3b_megatron.sh + +Using TensorRT-LLM rollout in fully async with GRPO +---------------------------------------------------- +.. code-block:: bash + + # Fully async policy with Megatron-Core training engine + bash verl/experimental/fully_async_policy/shell/grpo_30b_a3b_base_math_megatron_4_4_mis_trtllm.sh diff --git a/verl/examples/README.md b/verl/examples/README.md new file mode 100644 index 0000000000000000000000000000000000000000..a2d675a5d19fc54fa43826c1acc6a2fd7454c4f1 --- /dev/null +++ b/verl/examples/README.md @@ -0,0 +1,143 @@ +# verl Examples + +This directory hosts curated, minimal-dependency examples that drive +`verl.trainer.main_ppo` with the current Hydra API. Algorithm-specific +extensions, research baselines, and non-trivial entry points live under +`recipe/`; prefer that directory if you need a custom loss or reward beyond +what these examples show. + +## Conventions + +All run scripts follow the same shape: + +1. Canonical filename: + + ``` + run__.sh + ``` + + - ``: a single canonical size per model family. E.g. + `qwen3_8b`, `qwen3_30b_a3b`, `qwen3_235b_a22b`, `qwen3_vl_8b`, + `deepseek_v3`, `mimo_7b`, `nemotron_nano_v3`. + - ``: one of `fsdp`, `fsdp2`, `megatron`, `mindspeed`, + `automodel`, or `veomni`. **Must be the last underscore-separated + token before `.sh`**. + + Nothing follows ``. Per-example *features* — including + the inference backend (`vllm`/`sglang`/`trtllm`), the platform + (`DEVICE=gpu|npu`), the GPU machine type (`MACHINE=gb200`/`b200`/ + `blackwell`), Liger kernel, LoRA, FP8 quantization, sequence parallel + size, server vs sync rollout, etc. — do **not** show up in filenames. + They are exposed as env-var toggles inside the one canonical script. + Do not add `_npu`, `_amd`, `_vllm`, `_sglang`, `_trtllm`, or `_fp8` + script variants. For example, `sft/gsm8k/run_qwen2_5_0_5b_fsdp.sh` + covers plain SFT and its `USE_LIGER=1`, `SP_SIZE=2`, `USE_PEFT=1` + variants via env vars; `grpo_trainer/run_qwen3_8b_fsdp.sh` covers vLLM, + SGLang, and TRT-LLM rollouts, CUDA/NPU platforms, and `MACHINE=gb200` + (Blackwell) via toggles. + + This naming rule is enforced by the `check-example-naming` pre-commit + hook (see `tests/special_sanity/check_example_naming.py`). + +2. Every script exposes its important knobs in a user-adjustable region near + the top. Derived defaults and device/backend-specific details belong below + the "no user adjustment needed below" / "derived defaults" boundary. + Use uppercase env vars for user-facing knobs, e.g. + + ```bash + # ---- user-adjustable ---- + DEVICE=${DEVICE:-gpu} + INFER_BACKEND=${INFER_BACKEND:-vllm} + MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-8B} + NNODES=${NNODES:-1} + NDEVICES_PER_NODE=${NDEVICES_PER_NODE:-} + TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-1024} + ROLLOUT_TP=${ROLLOUT_TP:-2} + ROLLOUT_N=${ROLLOUT_N:-5} + PROJECT_NAME=${PROJECT_NAME:-verl_grpo_gsm8k_math} + EXPERIMENT_NAME=${EXPERIMENT_NAME:-qwen3_8b_grpo_vllm_fsdp} + # ---- end user-adjustable ---- + ... + ``` + + Override anything you care about on the command line: + + ```bash + DEVICE=npu MODEL_PATH=/my/local/qwen3-8b NDEVICES_PER_NODE=4 bash examples/grpo_trainer/run_qwen3_8b_fsdp.sh + ``` + + GPU and NPU paths should share the same `PROJECT_NAME` / + `EXPERIMENT_NAME` form. Do not append `_npu` to project or experiment + names just because `DEVICE=npu` is selected. + +3. Defaults (unless a directory explicitly documents otherwise): + + - `data.train_files` + `data.val_files` = GSM8K + MATH for text LLMs + (`geo3k` for vision, `dapo-math-17k` / `aime-2024` for scale-demo 235B / + 671B scripts). + - `actor_rollout_ref.actor.use_dynamic_bsz=True` + - `trainer.balance_batch=True` + - `trainer.logger=["console","wandb"]`. + +4. No deprecated Hydra knobs: + + - `ppo_megatron_trainer.yaml` → use `actor_rollout_ref.actor.model_engine=megatron`. + - `actor_rollout_ref.rollout.mode=async` → removed; async rollout is no + longer selected this way in example scripts. + - `actor_rollout_ref.hybrid_engine=True` → removed; the trainer now + enforces the supported hybrid-engine path internally. + - `ppo_micro_batch_size` / `log_prob_micro_batch_size` → use the + `_per_gpu` suffix. + - `data.val_batch_size` → removed. + - Top-level `reward_model.*` → use `reward_model.reward_model.*` / + `reward.reward_model.*` as applicable. + - `actor.ulysses_sequence_parallel_size` → use + `actor_rollout_ref.actor.fsdp_config.ulysses_sequence_parallel_size`. + +## Directory layout + +### Algorithm trainers + +Each directory holds a canonical recipe for one training algorithm. Adding a +new algorithm? If it needs its own trainer entry point or reward code, put it +under `recipe/` instead. + +| Dir | Algorithm | `algorithm.adv_estimator` / `policy_loss.loss_mode` | +|------------------------------------|--------------------------------|-----------------------------------------------------| +| `ppo_trainer/` | PPO (actor + critic) | `adv_estimator=gae` | +| `grpo_trainer/` | GRPO | `adv_estimator=grpo` | +| `rloo_trainer/` | RLOO | `adv_estimator=rloo` | +| `remax_trainer/` | ReMax | `adv_estimator=remax` | +| `reinforce_plus_plus_trainer/` | REINFORCE++ / baseline | `adv_estimator=reinforce_plus_plus[_baseline]` | +| `cispo_trainer/` | CISPO | `loss_mode=cispo` | +| `dppo_trainer/` | DPPO (TV / KL variants) | `loss_mode=dppo_tv \| dppo_kl` | +| `gdpo_trainer/` | GDPO | `adv_estimator=gdpo` | +| `gmpo_trainer/` | GMPO | `loss_mode=geo_mean` | +| `gpg_trainer/` | GPG | `adv_estimator=gpg`, `loss_mode=gpg` | +| `gspo_trainer/` | GSPO | `loss_mode=gspo` | +| `sapo_trainer/` | SAPO | `loss_mode=sapo` | +| `otb_trainer/` | OTB | `adv_estimator=optimal_token_baseline` | +| `mtp_trainer/` | DAPO + MTP (MiMo-7B) | `adv_estimator=grpo`, MTP flags | +| `on_policy_distillation_trainer/` | on-policy distillation | GRPO + distillation loss | +| `flowgrpo_trainer/` | Flow-GRPO (diffusion) | image-gen specific | + +### Feature / infra + +| Dir | Purpose | +|----------------------|------------------------------------------------------------------------------------------| +| `tuning/` | LoRA (`tuning/lora/`) and scaling demos (`tuning/scaling/`). | +| `profile/` | NPU profiler / torch-memory profiler runs. | +| `sft/` | Supervised fine-tuning examples. | +| `generation/` | Rollout-only inference launches. | +| `vllm_omni/` | vLLM omni backend examples. | +| `data_preprocess/` | Scripts that produce the `$HOME/data//*.parquet` layout the run scripts expect. | +| `prefix_grouper/` | Prefix-grouped rollout examples. | +| `rollout_correction/`| Rollout correction examples. | +| `router_replay/` | Router replay examples. | +| `tutorial/` | Tutorials and cluster launchers (`ray/`, `slurm/`, `skypilot/`, `agent_loop_get_started/`). | + +### Where are the algorithm research variants? + +`recipe/` — e.g. `recipe/dapo`, `recipe/prime`, `recipe/retool`, +`recipe/r1`, `recipe/spin`, `recipe/gvpo`, `recipe/flowrl`, ... +They ship their own trainer entry points and reward code. diff --git a/verl/examples/ascend_extras/grpo_trainer/run_qwen3_235b_256k_megatron.sh b/verl/examples/ascend_extras/grpo_trainer/run_qwen3_235b_256k_megatron.sh new file mode 100644 index 0000000000000000000000000000000000000000..24c94cd55622eb1abb0ba00a592e8343f1c4e6f6 --- /dev/null +++ b/verl/examples/ascend_extras/grpo_trainer/run_qwen3_235b_256k_megatron.sh @@ -0,0 +1,207 @@ +#!/bin/bash +# set -xeuo pipefail + +## !!!!!!!supplement!!!!!! +## This script can be used for inference in 256 K and 128 K + +ulimit -n 32768 + +# Project Configuration +project_name='GRPO-Qwen3-235B-A22B-Instruct-MATH' +exp_name='GRPO-Qwen3-235B-A22B-Instruct-Megatron-vLLM' + +# Node Info +NNODES=${NNODES:-16} +NPUS_PER_NODE=${NPUS_PER_NODE:-16} + +# Model Weights Paths +# MODEL_PATH=/mnt/weight/Qwen3-235B-A22B +MODEL_PATH=${WORK_DIR}/Qwen3-235B-A22B-Instruct-2507 +MCORE_MODEL_PATH=${WORK_DIR}/Qwen3-235B-A22B-Instruct-2507-Mcore +RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"} +CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"} + +# File System Paths +TRAIN_FILE=${WORK_DIR}/gsm8k/train.parquet +TEST_FILE=${WORK_DIR}/gsm8k/test.parquet + +# Data Configuration +max_prompt_length=$((1024 * 1)) +max_response_length=$((1024 * 255)) + +# Algorithm Configuration +adv_estimator=grpo +use_kl_in_reward=False +kl_coef=0.0 +use_kl_loss=True +kl_loss_coef=0.001 + +# Training Batch Configuration +train_prompt_bsz=4 +n_resp_per_prompt=4 +train_prompt_mini_bsz=4 + +# Performance and Memory Related Configuration +all_offload=True +use_dynamic_bsz=False +actor_ppo_max_token_len=$((max_prompt_length + max_response_length)) +infer_ppo_max_token_len=$((max_prompt_length + max_response_length)) +optimizer_offload_fraction=1 + +# Megatron Configuration +train_tp=2 +train_ep=16 +train_etp=1 +train_pp=16 +train_cp=8 + +# vLLM Configuration +gen_tp=4 +gen_dp=32 +gen_ep=128 +gpu_memory_utilization=0.7 +max_model_len=$((max_prompt_length + max_response_length)) +max_num_batched_tokens=2048 + +# Pipeline Layer Configuration +first_layer=5 +last_layer=5 + +# Data Configuration +DATA_ARGS=( + data.train_files="${TRAIN_FILE}" + data.val_files="${TEST_FILE}" + data.prompt_key=prompt + data.train_batch_size=${train_prompt_bsz} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=False + data.truncation='left' +) + +# Model Configuration +MODEL_ARGS=( + actor_rollout_ref.model.path="${MODEL_PATH}" + actor_rollout_ref.model.use_remove_padding=True +) + +# RL Algorithm Configuration +ALGORITHM_ARGS=( + algorithm.adv_estimator=${adv_estimator} + algorithm.use_kl_in_reward=${use_kl_in_reward} + algorithm.kl_ctrl.kl_coef=${kl_coef} +) + +# Actor Model Configuration +ACTOR_ARGS=( + actor_rollout_ref.actor.use_torch_compile=False + actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} + actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.entropy_coeff=0 + actor_rollout_ref.actor.ppo_epochs=1 + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} + actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.optim.clip_grad=1.0 + actor_rollout_ref.actor.optim.lr_warmup_steps=10 + actor_rollout_ref.actor.optim.weight_decay=0.1 + actor_rollout_ref.actor.optim.lr=1e-6 + +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=${optimizer_offload_fraction} + +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=True + +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${train_tp} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${train_pp} + actor_rollout_ref.actor.megatron.context_parallel_size=${train_cp} + actor_rollout_ref.actor.megatron.expert_model_parallel_size=${train_ep} + actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${train_etp} + actor_rollout_ref.actor.megatron.param_offload=${all_offload} + actor_rollout_ref.actor.megatron.optimizer_offload=${all_offload} + actor_rollout_ref.actor.megatron.grad_offload=${all_offload} + actor_rollout_ref.actor.megatron.use_dist_checkpointing=True + actor_rollout_ref.actor.megatron.use_mbridge=False + actor_rollout_ref.actor.megatron.dist_checkpointing_path=${MCORE_MODEL_PATH} + +actor_rollout_ref.actor.megatron.override_transformer_config.use_flash_attn=True + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 + +actor_rollout_ref.actor.megatron.override_transformer_config.context_parallel_size=${train_cp} + +actor_rollout_ref.actor.megatron.override_transformer_config.num_layers_in_first_pipeline_stage=${first_layer} + +actor_rollout_ref.actor.megatron.override_transformer_config.num_layers_in_last_pipeline_stage=${last_layer} + +actor_rollout_ref.actor.megatron.override_transformer_config.normalization=RMSNorm + +actor_rollout_ref.actor.megatron.override_transformer_config.use_fused_rmsnorm=True + +actor_rollout_ref.actor.megatron.override_transformer_config.swiglu=True + +actor_rollout_ref.actor.megatron.override_transformer_config.use_fused_swiglu=True + +actor_rollout_ref.actor.megatron.override_transformer_config.use_distributed_optimizer=True + +actor_rollout_ref.actor.megatron.override_transformer_config.sequence_parallel=True +) + +# Reference Model Configuration +REF_ARGS=( + actor_rollout_ref.ref.use_torch_compile=False + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${train_tp} + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${train_pp} + actor_rollout_ref.ref.megatron.context_parallel_size=${train_cp} + actor_rollout_ref.ref.megatron.expert_model_parallel_size=${train_ep} + actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=${train_etp} + actor_rollout_ref.ref.megatron.param_offload=${all_offload} + actor_rollout_ref.ref.megatron.use_mbridge=False + actor_rollout_ref.ref.megatron.use_dist_checkpointing=True + actor_rollout_ref.ref.megatron.dist_checkpointing_path=${MCORE_MODEL_PATH} + ++actor_rollout_ref.ref.megatron.override_transformer_config.use_flash_attn=True + +actor_rollout_ref.ref.megatron.override_transformer_config.sequence_parallel=True +) + +# Rollout Configuration +ROLLOUT_ARGS=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.n=${n_resp_per_prompt} + actor_rollout_ref.rollout.max_num_seqs=16 + actor_rollout_ref.rollout.top_p=1.0 + actor_rollout_ref.rollout.top_k=-1 + actor_rollout_ref.rollout.temperature=1.0 + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} + actor_rollout_ref.rollout.gpu_memory_utilization=${gpu_memory_utilization} + actor_rollout_ref.rollout.max_num_batched_tokens=${max_num_batched_tokens} + actor_rollout_ref.rollout.max_model_len=${max_model_len} + actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} + actor_rollout_ref.rollout.data_parallel_size=${gen_dp} + actor_rollout_ref.rollout.expert_parallel_size=${gen_ep} + actor_rollout_ref.rollout.enable_chunked_prefill=True + actor_rollout_ref.rollout.enable_prefix_caching=True + actor_rollout_ref.rollout.enforce_eager=True + actor_rollout_ref.rollout.free_cache_engine=True +) + +# Trainer Configuration +TRAINER_ARGS=( + trainer.logger='["console","tensorboard"]' + trainer.project_name="${project_name}" + trainer.experiment_name="${exp_name}" + trainer.nnodes="${NNODES}" + trainer.n_gpus_per_node="${NPUS_PER_NODE}" + trainer.device='npu' + trainer.total_epochs=15 + trainer.val_before_train=False + trainer.test_freq=-1 + trainer.save_freq=-1 + trainer.default_local_dir="${CKPTS_DIR}" +) + +python3 -m verl.trainer.main_ppo \ + --config-path=config \ + --config-name='ppo_megatron_trainer.yaml' \ + "${DATA_ARGS[@]}" \ + "${MODEL_ARGS[@]}" \ + "${ACTOR_ARGS[@]}" \ + "${REF_ARGS[@]}" \ + "${ROLLOUT_ARGS[@]}" \ + "${ALGORITHM_ARGS[@]}" \ + "${TRAINER_ARGS[@]}" \ + "$@" | tee logs/run_qwen3moe-wy_235b_grpo_megatron_vllm_npu_$(date +%Y%m%d_%H%M%S).log \ No newline at end of file diff --git a/verl/examples/ascend_extras/grpo_trainer/run_qwen3_30b_a3b_megatron.sh b/verl/examples/ascend_extras/grpo_trainer/run_qwen3_30b_a3b_megatron.sh new file mode 100644 index 0000000000000000000000000000000000000000..c94de6a386f5abdf7cc7b2b99fe82fe32d70fe5f --- /dev/null +++ b/verl/examples/ascend_extras/grpo_trainer/run_qwen3_30b_a3b_megatron.sh @@ -0,0 +1,236 @@ +#!/bin/bash +set -xeuo pipefail +# Project Configuration +project_name='DAPO-Qwen3-30b-A3B-BASE-MATH' +exp_name='DAPO-Qwen3-30B-A3B-BASE-Megatron-SGLang' + +# Necessary env +export HCCL_CONNECT_TIMEOUT=1500 +export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050 +export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050 + +export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 +export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15 + +export DISABLE_L2_CACHE=1 +export TASK_QUEUE_ENABLE=1 + +# Node Info +NNODES=${NNODES:-1} +NPUS_PER_NODE=${NPUS_PER_NODE:-16} + +# Model Weights Paths +MODEL_PATH=Qwen/Qwen3-30B-A3B +MCORE_MODEL_PATH=Qwen/Qwen3-30B-A3B-mcore +RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"} +CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"} + +# File System Paths +TRAIN_FILE=$RAY_DATA_HOME/dataset/dapo-math-17k.parquet +TEST_FILE=$RAY_DATA_HOME/dataset/aime-2024.parquet +# Data Length Configuration +max_prompt_length=$((1024 * 2)) +max_response_length=$((1024 * 8)) + +# Training Batch Configuration +train_prompt_bsz=16 +train_prompt_mini_bsz=16 +n_resp_per_prompt=8 + +# Algorithm Configuration +adv_estimator=grpo +use_kl_in_reward=False +kl_coef=0.0 +use_kl_loss=True +kl_loss_coef=0.001 + +# Performance and Memory Management Configuration +all_offload=True +use_dynamic_bsz=False +actor_ppo_max_token_len=$(((max_prompt_length + max_response_length))) +infer_ppo_max_token_len=$(((max_prompt_length + max_response_length))) + +# Megatron Parallelism Configuration +train_tp=4 +train_ep=4 +train_etp=4 +train_pp=1 +train_cp=1 + +# SGLang Generation Configuration +gen_tp=4 +gen_dp=1 +gen_ep=1 +gpu_memory_utilization=0.5 +max_model_len=$((max_prompt_length + max_response_length)) +max_num_batched_tokens=$(((max_prompt_length + max_response_length) * 1)) + +# Data Configuration +DATA_CONFIG=( + # File Paths + data.train_files="${TRAIN_FILE}" + data.val_files="${TEST_FILE}" + # Data Structure + data.prompt_key=prompt + # Batch and Length Configuration + data.train_batch_size=${train_prompt_bsz} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + # Preprocessing + data.filter_overlong_prompts=False + data.truncation='left' +) + +# Model Configuration +MODEL_CONFIG=( + # Model Path + actor_rollout_ref.model.path="${MODEL_PATH}" + # Model Processing + actor_rollout_ref.model.use_remove_padding=True +) + +# Reinforcement Learning Algorithm Configuration +ALGORITHM_CONFIG=( + # Advantage Estimation + algorithm.adv_estimator=${adv_estimator} + # KL Divergence Control + algorithm.use_kl_in_reward=${use_kl_in_reward} + algorithm.kl_ctrl.kl_coef=${kl_coef} +) + +ACTOR_CONFIG=( + # Core Runtime Settings + actor_rollout_ref.actor.use_torch_compile=False + actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} + # Loss Function Configuration + actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.entropy_coeff=0 + # PPO Training Parameters + actor_rollout_ref.actor.ppo_epochs=1 + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} + actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} + # Optimizer Settings + actor_rollout_ref.actor.optim.lr=1e-6 + # Megatron Parallelism Strategy + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${train_tp} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${train_pp} + actor_rollout_ref.actor.megatron.context_parallel_size=${train_cp} + actor_rollout_ref.actor.megatron.expert_model_parallel_size=${train_ep} + actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${train_etp} + # Memory Optimization + actor_rollout_ref.actor.megatron.param_offload=${all_offload} + actor_rollout_ref.actor.megatron.optimizer_offload=${all_offload} + actor_rollout_ref.actor.megatron.grad_offload=${all_offload} + # Model Weights Management + actor_rollout_ref.actor.megatron.dist_checkpointing_path=${MCORE_MODEL_PATH} + actor_rollout_ref.actor.megatron.use_dist_checkpointing=True + actor_rollout_ref.actor.megatron.use_mbridge=False + # Transformer Architecture Optimizations + +actor_rollout_ref.actor.megatron.override_transformer_config.use_flash_attn=True + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 +) + +REF_CONFIG=( + # Core Runtime Settings + actor_rollout_ref.ref.use_torch_compile=False + # Log Probability Inference + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} + # Megatron Parallelism Strategy + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${train_tp} + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${train_pp} + actor_rollout_ref.ref.megatron.context_parallel_size=${train_cp} + actor_rollout_ref.ref.megatron.expert_model_parallel_size=${train_ep} + actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=${train_etp} + # Memory Optimization + actor_rollout_ref.ref.megatron.param_offload=${all_offload} + # Model Weights Management + actor_rollout_ref.ref.megatron.dist_checkpointing_path=${MCORE_MODEL_PATH} + actor_rollout_ref.ref.megatron.use_dist_checkpointing=True + actor_rollout_ref.ref.megatron.use_mbridge=False +) + +ROLLOUT_CONFIG=( + # Rollout Engine + actor_rollout_ref.rollout.name=sglang + +actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend="ascend" + # Generation Parameters + actor_rollout_ref.rollout.n=${n_resp_per_prompt} + actor_rollout_ref.rollout.top_p=1.0 + actor_rollout_ref.rollout.top_k=-1 + actor_rollout_ref.rollout.temperature=1.0 + # Log Probability Inference + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} + # Memory Management + actor_rollout_ref.rollout.gpu_memory_utilization=${gpu_memory_utilization} + # Parallelism Strategy + actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} + actor_rollout_ref.rollout.data_parallel_size=${gen_dp} + actor_rollout_ref.rollout.expert_parallel_size=${gen_ep} + +actor_rollout_ref.rollout.engine_kwargs.sglang.enable_dp_attention=False + # Performance Optimization + +actor_rollout_ref.rollout.engine_kwargs.sglang.chunked_prefill_size=-1 + actor_rollout_ref.rollout.enforce_eager=False + # Validation Generation + actor_rollout_ref.rollout.val_kwargs.n=1 + actor_rollout_ref.rollout.val_kwargs.do_sample=True + actor_rollout_ref.rollout.val_kwargs.top_p=1.0 + actor_rollout_ref.rollout.val_kwargs.top_k=-1 + actor_rollout_ref.rollout.val_kwargs.temperature=1.0 +) + +TRAINER_CONFIG=( + # Logger Configuration + trainer.logger='["console"]' + # Project Settings + trainer.project_name="${project_name}" + trainer.experiment_name="${exp_name}" + # Hardware Configuration + trainer.nnodes="${NNODES}" + trainer.n_gpus_per_node="${NPUS_PER_NODE}" + trainer.device='npu' + # Training Schedule + trainer.total_epochs=15 + trainer.val_before_train=False + trainer.test_freq=-1 + trainer.save_freq=-1 + # Checkpoint Directory + trainer.default_local_dir="${CKPTS_DIR}" +) + +# profiling configuration +PROF_CONFIG=( + global_profiler.tool=npu + global_profiler.steps=null + global_profiler.save_path=/profpath + actor_rollout_ref.actor.profiler.enable=True + actor_rollout_ref.actor.profiler.ranks="[0]" + actor_rollout_ref.actor.profiler.all_ranks=False + actor_rollout_ref.actor.profiler.tool_config.npu.discrete=True + actor_rollout_ref.actor.profiler.tool_config.npu.contents=['npu','cpu'] + actor_rollout_ref.actor.profiler.tool_config.npu.level=level0 + actor_rollout_ref.actor.profiler.tool_config.npu.analysis=True + actor_rollout_ref.rollout.profiler.enable=True + actor_rollout_ref.rollout.profiler.ranks="[0]" + actor_rollout_ref.rollout.profiler.all_ranks=False +) + +python3 -m verl.trainer.main_ppo \ + --config-path=config \ + --config-name='ppo_megatron_trainer.yaml' \ + "${DATA_CONFIG[@]}" \ + "${MODEL_CONFIG[@]}" \ + "${ACTOR_CONFIG[@]}" \ + "${REF_CONFIG[@]}" \ + "${ROLLOUT_CONFIG[@]}" \ + "${ALGORITHM_CONFIG[@]}" \ + "${TRAINER_CONFIG[@]}" \ + "${PROF_CONFIG[@]}" \ + "$@" \ No newline at end of file diff --git a/verl/examples/ascend_extras/grpo_trainer/run_qwen3_30b_a3b_mindspeed.sh b/verl/examples/ascend_extras/grpo_trainer/run_qwen3_30b_a3b_mindspeed.sh new file mode 100644 index 0000000000000000000000000000000000000000..fd22d7472d9c96a114dca1525609d937348e5406 --- /dev/null +++ b/verl/examples/ascend_extras/grpo_trainer/run_qwen3_30b_a3b_mindspeed.sh @@ -0,0 +1,247 @@ +#!/bin/bash +set -xeuo pipefail +# Project Configuration +project_name='GRPO-Qwen3-30b-A3B-BASE-MATH' +exp_name='GRPO-Qwen3-30B-A3B-BASE-MindSpeedLLM-SGLang' + +# Necessary env +export HCCL_CONNECT_TIMEOUT=1500 +export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050 +export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050 + +export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 +export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15 + +export DISABLE_L2_CACHE=1 +export TASK_QUEUE_ENABLE=1 + +#For CANN versions 8.5.0 and above, using mbridge, set this ENV +export HCCL_OP_EXPANSION_MODE="AIV" + +# Node Info +NNODES=${NNODES:-1} +NPUS_PER_NODE=${NPUS_PER_NODE:-16} + +# Model Weights Paths +MODEL_PATH=Qwen/Qwen3-30B-A3B +RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"} +CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"} + +# File System Paths +TRAIN_FILE=$RAY_DATA_HOME/gsm8k/train.parquet +TEST_FILE=$RAY_DATA_HOME/gsm8k/test.parquet +# Data Length Configuration +max_prompt_length=$((1024 * 2)) +max_response_length=$((1024 * 2)) + +# Training Batch Configuration +train_prompt_bsz=16 +train_prompt_mini_bsz=16 +n_resp_per_prompt=8 +micro_batch_size=1 + +# Algorithm Configuration +adv_estimator=grpo +use_kl_in_reward=False +kl_coef=0.0 +use_kl_loss=True +kl_loss_coef=0.001 + +# Performance and Memory Management Configuration +all_offload=True +use_dynamic_bsz=False +actor_ppo_max_token_len=$(((max_prompt_length + max_response_length))) +infer_ppo_max_token_len=$(((max_prompt_length + max_response_length))) + +# Megatron Parallelism Configuration +train_tp=4 +train_ep=4 +train_etp=1 +train_pp=4 +train_cp=1 + +# SGLang Generation Configuration +gen_tp=4 +gen_dp=1 +gen_ep=1 +gpu_memory_utilization=0.5 +max_model_len=$((max_prompt_length + max_response_length)) +max_num_batched_tokens=$(((max_prompt_length + max_response_length) * 1)) + +# Data Configuration +DATA_CONFIG=( + # File Paths + data.train_files="${TRAIN_FILE}" + data.val_files="${TEST_FILE}" + # Data Structure + data.prompt_key=prompt + # Batch and Length Configuration + data.train_batch_size=${train_prompt_bsz} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + # Preprocessing + data.filter_overlong_prompts=False + data.truncation='left' +) + +# Model Configuration +MODEL_CONFIG=( + # Model Path + actor_rollout_ref.model.path="${MODEL_PATH}" + # Model Processing + actor_rollout_ref.model.use_remove_padding=True +) + +# Reinforcement Learning Algorithm Configuration +ALGORITHM_CONFIG=( + # Advantage Estimation + algorithm.adv_estimator=${adv_estimator} + # KL Divergence Control + algorithm.use_kl_in_reward=${use_kl_in_reward} + algorithm.kl_ctrl.kl_coef=${kl_coef} +) + +ACTOR_CONFIG=( + # Core Runtime Settings + actor_rollout_ref.actor.use_torch_compile=False + actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} + # Loss Function Configuration + actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.entropy_coeff=0 + # PPO Training Parameters + actor_rollout_ref.actor.ppo_epochs=1 + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${micro_batch_size} + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} + actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} + # Optimizer Settings + actor_rollout_ref.actor.optim.lr=1e-6 + # Megatron Parallelism Strategy + actor_rollout_ref.actor.mindspeed.tensor_model_parallel_size=${train_tp} + actor_rollout_ref.actor.mindspeed.pipeline_model_parallel_size=${train_pp} + actor_rollout_ref.actor.mindspeed.context_parallel_size=${train_cp} + actor_rollout_ref.actor.mindspeed.expert_model_parallel_size=${train_ep} + actor_rollout_ref.actor.mindspeed.expert_tensor_parallel_size=${train_etp} + # Memory Optimization + actor_rollout_ref.actor.mindspeed.param_offload=${all_offload} + actor_rollout_ref.actor.mindspeed.optimizer_offload=${all_offload} + actor_rollout_ref.actor.mindspeed.grad_offload=${all_offload} + # Model Weights Management + actor_rollout_ref.actor.mindspeed.use_mbridge=True + actor_rollout_ref.actor.mindspeed.vanilla_mbridge=True + # Transformer Architecture Optimizations + actor_rollout_ref.actor.mindspeed.llm_kwargs.spec='[mindspeed_llm.tasks.models.spec.qwen3_spec, layer_spec]' + actor_rollout_ref.actor.mindspeed.llm_kwargs.seq_length=${max_model_len} + actor_rollout_ref.actor.mindspeed.llm_kwargs.micro_batch_size=${micro_batch_size} + +actor_rollout_ref.actor.mindspeed.llm_kwargs.num_query_groups=4 + +actor_rollout_ref.actor.mindspeed.llm_kwargs.recompute_method=uniform + +actor_rollout_ref.actor.mindspeed.llm_kwargs.recompute_granularity=full + +actor_rollout_ref.actor.mindspeed.llm_kwargs.recompute_num_layers=1 + # MOE + +actor_rollout_ref.actor.mindspeed.llm_kwargs.moe_router_load_balancing_type=aux_loss + +actor_rollout_ref.actor.mindspeed.llm_kwargs.moe_permutation_async_comm=True + +actor_rollout_ref.actor.mindspeed.llm_kwargs.moe_token_dispatcher_type=alltoall + +actor_rollout_ref.actor.mindspeed.llm_kwargs.moe_aux_loss_coeff=0.001 + +actor_rollout_ref.actor.mindspeed.llm_kwargs.moe_grouped_gemm=True +) + +REF_CONFIG=( + # Core Runtime Settings + actor_rollout_ref.ref.use_torch_compile=False + # Log Probability Inference + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${micro_batch_size} + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} + # Megatron Parallelism Strategy + actor_rollout_ref.ref.mindspeed.tensor_model_parallel_size=${train_tp} + actor_rollout_ref.ref.mindspeed.pipeline_model_parallel_size=${train_pp} + actor_rollout_ref.ref.mindspeed.context_parallel_size=${train_cp} + actor_rollout_ref.ref.mindspeed.expert_model_parallel_size=${train_ep} + actor_rollout_ref.ref.mindspeed.expert_tensor_parallel_size=${train_etp} + # Memory Optimization + actor_rollout_ref.ref.mindspeed.param_offload=${all_offload} + # Model Weights Management + actor_rollout_ref.ref.mindspeed.use_mbridge=True + actor_rollout_ref.ref.mindspeed.vanilla_mbridge=True +) + +ROLLOUT_CONFIG=( + # Rollout Engine + actor_rollout_ref.rollout.name=sglang + +actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend="ascend" + # Generation Parameters + actor_rollout_ref.rollout.n=${n_resp_per_prompt} + actor_rollout_ref.rollout.top_p=1.0 + actor_rollout_ref.rollout.top_k=-1 + actor_rollout_ref.rollout.temperature=1.0 + # Log Probability Inference + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${micro_batch_size} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} + # Memory Management + actor_rollout_ref.rollout.gpu_memory_utilization=${gpu_memory_utilization} + # Parallelism Strategy + actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} + actor_rollout_ref.rollout.data_parallel_size=${gen_dp} + actor_rollout_ref.rollout.expert_parallel_size=${gen_ep} + +actor_rollout_ref.rollout.engine_kwargs.sglang.enable_dp_attention=False + # Performance Optimization + +actor_rollout_ref.rollout.engine_kwargs.sglang.chunked_prefill_size=-1 + actor_rollout_ref.rollout.enforce_eager=False + # Validation Generation + actor_rollout_ref.rollout.val_kwargs.n=1 + actor_rollout_ref.rollout.val_kwargs.do_sample=True + actor_rollout_ref.rollout.val_kwargs.top_p=1.0 + actor_rollout_ref.rollout.val_kwargs.top_k=-1 + actor_rollout_ref.rollout.val_kwargs.temperature=1.0 +) + +TRAINER_CONFIG=( + # Logger Configuration + trainer.logger='["console"]' + # Project Settings + trainer.project_name="${project_name}" + trainer.experiment_name="${exp_name}" + # Hardware Configuration + trainer.nnodes="${NNODES}" + trainer.n_gpus_per_node="${NPUS_PER_NODE}" + trainer.device='npu' + # Training Schedule + trainer.total_epochs=15 + trainer.val_before_train=False + trainer.test_freq=-1 + trainer.save_freq=-1 + # Checkpoint Directory + trainer.default_local_dir="${CKPTS_DIR}" +) + +# profiling configuration +PROF_CONFIG=( + global_profiler.tool=npu + global_profiler.steps=null + global_profiler.save_path=/profpath + actor_rollout_ref.actor.profiler.enable=True + actor_rollout_ref.actor.profiler.ranks="[0]" + actor_rollout_ref.actor.profiler.all_ranks=False + actor_rollout_ref.actor.profiler.tool_config.npu.discrete=True + actor_rollout_ref.actor.profiler.tool_config.npu.contents=['npu','cpu'] + actor_rollout_ref.actor.profiler.tool_config.npu.level=level0 + actor_rollout_ref.actor.profiler.tool_config.npu.analysis=True + actor_rollout_ref.rollout.profiler.enable=True + actor_rollout_ref.rollout.profiler.ranks="[0]" + actor_rollout_ref.rollout.profiler.all_ranks=False +) + +python3 -m verl.trainer.main_ppo \ + --config-path=config \ + --config-name='ppo_trainer.yaml' \ + model_engine=mindspeed \ + "${DATA_CONFIG[@]}" \ + "${MODEL_CONFIG[@]}" \ + "${ACTOR_CONFIG[@]}" \ + "${REF_CONFIG[@]}" \ + "${ROLLOUT_CONFIG[@]}" \ + "${ALGORITHM_CONFIG[@]}" \ + "${TRAINER_CONFIG[@]}" \ + "${PROF_CONFIG[@]}" \ + "$@" \ No newline at end of file diff --git a/verl/examples/ascend_extras/grpo_trainer/run_qwen3_32b_fsdp.sh b/verl/examples/ascend_extras/grpo_trainer/run_qwen3_32b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..6384a564a6afd19dcd690cea86f8104306096412 --- /dev/null +++ b/verl/examples/ascend_extras/grpo_trainer/run_qwen3_32b_fsdp.sh @@ -0,0 +1,59 @@ +set -x + +project_name='GRPO-Qwen3' +exp_name='GRPO-Qwen3-32b-npu' +gen_tp=4 +RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"} +MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen3-32B"} +TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/gsm8k/train.parquet"} +TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/gsm8k/test.parquet"} + +python3 -m verl.trainer.main_ppo \ + algorithm.adv_estimator=grpo \ + data.train_files="${TRAIN_FILE}" \ + data.val_files="${TEST_FILE}" \ + data.train_batch_size=1024 \ + data.max_prompt_length=2048 \ + data.max_response_length=2048 \ + data.filter_overlong_prompts=True \ + data.truncation='error' \ + data.shuffle=False \ + actor_rollout_ref.model.path=${MODEL_PATH} \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.model.use_remove_padding=True \ + actor_rollout_ref.actor.ulysses_sequence_parallel_size=4 \ + +actor_rollout_ref.actor.fsdp_config.mixed_precision.param_dtype=bf16 \ + +actor_rollout_ref.actor.fsdp_config.mixed_precision.reduce_dtype=bf16 \ + +actor_rollout_ref.actor.fsdp_config.mixed_precision.buffer_dtype=fp32 \ + actor_rollout_ref.actor.ppo_mini_batch_size=64 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=8 \ + actor_rollout_ref.actor.use_kl_loss=True \ + actor_rollout_ref.actor.entropy_coeff=0 \ + actor_rollout_ref.actor.kl_loss_coef=0.001 \ + actor_rollout_ref.actor.kl_loss_type=low_var_kl \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.fsdp_config.param_offload=True \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \ + actor_rollout_ref.rollout.name=vllm \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.7 \ + actor_rollout_ref.rollout.n=4 \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=8 \ + actor_rollout_ref.ref.fsdp_config.param_offload=True \ + actor_rollout_ref.actor.use_torch_compile=False \ + actor_rollout_ref.ref.use_torch_compile=False \ + actor_rollout_ref.rollout.enable_chunked_prefill=True \ + actor_rollout_ref.rollout.max_num_batched_tokens=32768 \ + actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=4096 \ + algorithm.use_kl_in_reward=False \ + trainer.critic_warmup=0 \ + trainer.logger=['console','tensorboard'] \ + trainer.project_name="${project_name}" \ + trainer.experiment_name="${exp_name}" \ + trainer.n_gpus_per_node=8 \ + trainer.nnodes=4 \ + trainer.resume_from_path=checkpoints/ \ + trainer.save_freq=500 \ + trainer.test_freq=50 \ + trainer.total_epochs=50 $@ \ No newline at end of file diff --git a/verl/examples/ascend_extras/grpo_trainer/run_qwen3_32b_mindspeed.sh b/verl/examples/ascend_extras/grpo_trainer/run_qwen3_32b_mindspeed.sh new file mode 100644 index 0000000000000000000000000000000000000000..90b8e2f276c9193ab7c073bf55136cad16892a28 --- /dev/null +++ b/verl/examples/ascend_extras/grpo_trainer/run_qwen3_32b_mindspeed.sh @@ -0,0 +1,235 @@ +#!/bin/bash +set -xeuo pipefail +# Project Configuration +project_name='GRPO-Qwen3-32B-BASE-MATH' +exp_name='GRPO-Qwen3-32B-BASE-MindSpeedLLM-SGLang' + +# Necessary env +export HCCL_CONNECT_TIMEOUT=1500 +export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050 +export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050 + +export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 +export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15 + +export DISABLE_L2_CACHE=1 +export TASK_QUEUE_ENABLE=1 + +#For CANN versions 8.5.0 and above, using mbridge, set this ENV +export HCCL_OP_EXPANSION_MODE="AIV" + +# Node Info +NNODES=${NNODES:-1} +NPUS_PER_NODE=${NPUS_PER_NODE:-16} + +# Model Weights Paths +MODEL_PATH=Qwen/Qwen3-32B +RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"} +CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"} + +# File System Paths +TRAIN_FILE=$RAY_DATA_HOME/gsm8k/train.parquet +TEST_FILE=$RAY_DATA_HOME/gsm8k/test.parquet +# Data Length Configuration +max_prompt_length=$((1024 * 2)) +max_response_length=$((1024 * 2)) + +# Training Batch Configuration +train_prompt_bsz=16 +train_prompt_mini_bsz=16 +n_resp_per_prompt=8 +micro_batch_size=1 + +# Algorithm Configuration +adv_estimator=grpo +use_kl_in_reward=False +kl_coef=0.0 +use_kl_loss=True +kl_loss_coef=0.001 + +# Performance and Memory Management Configuration +all_offload=True +use_dynamic_bsz=False +actor_ppo_max_token_len=$(((max_prompt_length + max_response_length))) +infer_ppo_max_token_len=$(((max_prompt_length + max_response_length))) + +# Megatron Parallelism Configuration +train_tp=4 +train_pp=4 +train_cp=1 + +# SGLang Generation Configuration +gen_tp=4 +gen_dp=1 +gpu_memory_utilization=0.5 +max_model_len=$((max_prompt_length + max_response_length)) +max_num_batched_tokens=$(((max_prompt_length + max_response_length) * 1)) + +# Data Configuration +DATA_CONFIG=( + # File Paths + data.train_files="${TRAIN_FILE}" + data.val_files="${TEST_FILE}" + # Data Structure + data.prompt_key=prompt + # Batch and Length Configuration + data.train_batch_size=${train_prompt_bsz} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + # Preprocessing + data.filter_overlong_prompts=False + data.truncation='left' +) + +# Model Configuration +MODEL_CONFIG=( + # Model Path + actor_rollout_ref.model.path="${MODEL_PATH}" + # Model Processing + actor_rollout_ref.model.use_remove_padding=True +) + +# Reinforcement Learning Algorithm Configuration +ALGORITHM_CONFIG=( + # Advantage Estimation + algorithm.adv_estimator=${adv_estimator} + # KL Divergence Control + algorithm.use_kl_in_reward=${use_kl_in_reward} + algorithm.kl_ctrl.kl_coef=${kl_coef} +) + +ACTOR_CONFIG=( + # Core Runtime Settings + actor_rollout_ref.actor.use_torch_compile=False + actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} + # Loss Function Configuration + actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.entropy_coeff=0 + # PPO Training Parameters + actor_rollout_ref.actor.ppo_epochs=1 + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${micro_batch_size} + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} + actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} + # Optimizer Settings + actor_rollout_ref.actor.optim.lr=1e-6 + # Megatron Parallelism Strategy + actor_rollout_ref.actor.mindspeed.tensor_model_parallel_size=${train_tp} + actor_rollout_ref.actor.mindspeed.pipeline_model_parallel_size=${train_pp} + actor_rollout_ref.actor.mindspeed.context_parallel_size=${train_cp} + # Memory Optimization + actor_rollout_ref.actor.mindspeed.param_offload=${all_offload} + actor_rollout_ref.actor.mindspeed.optimizer_offload=${all_offload} + actor_rollout_ref.actor.mindspeed.grad_offload=${all_offload} + # Model Weights Management + actor_rollout_ref.actor.mindspeed.use_mbridge=True + actor_rollout_ref.actor.mindspeed.vanilla_mbridge=True + # Transformer Architecture Optimizations + actor_rollout_ref.actor.mindspeed.llm_kwargs.spec='[mindspeed_llm.tasks.models.spec.qwen3_spec, layer_spec]' + actor_rollout_ref.actor.mindspeed.llm_kwargs.seq_length=${max_model_len} + actor_rollout_ref.actor.mindspeed.llm_kwargs.micro_batch_size=${micro_batch_size} + +actor_rollout_ref.actor.mindspeed.llm_kwargs.num_query_groups=8 + +actor_rollout_ref.actor.mindspeed.llm_kwargs.recompute_method=uniform + +actor_rollout_ref.actor.mindspeed.llm_kwargs.recompute_granularity=full + +actor_rollout_ref.actor.mindspeed.llm_kwargs.recompute_num_layers=1 + +actor_rollout_ref.actor.mindspeed.llm_kwargs.overlap_grad_reduce=True + +actor_rollout_ref.actor.mindspeed.llm_kwargs.overlap_param_gather=True +) + +REF_CONFIG=( + # Core Runtime Settings + actor_rollout_ref.ref.use_torch_compile=False + # Log Probability Inference + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${micro_batch_size} + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} + # Megatron Parallelism Strategy + actor_rollout_ref.ref.mindspeed.tensor_model_parallel_size=${train_tp} + actor_rollout_ref.ref.mindspeed.pipeline_model_parallel_size=${train_pp} + actor_rollout_ref.ref.mindspeed.context_parallel_size=${train_cp} + # Memory Optimization + actor_rollout_ref.ref.mindspeed.param_offload=${all_offload} + # Model Weights Management + actor_rollout_ref.ref.mindspeed.use_mbridge=True + actor_rollout_ref.ref.mindspeed.vanilla_mbridge=True +) + +ROLLOUT_CONFIG=( + # Rollout Engine + actor_rollout_ref.rollout.name=sglang + +actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend="ascend" + # Generation Parameters + actor_rollout_ref.rollout.n=${n_resp_per_prompt} + actor_rollout_ref.rollout.top_p=1.0 + actor_rollout_ref.rollout.top_k=-1 + actor_rollout_ref.rollout.temperature=1.0 + # Log Probability Inference + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${micro_batch_size} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} + # Memory Management + actor_rollout_ref.rollout.gpu_memory_utilization=${gpu_memory_utilization} + # Parallelism Strategy + actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} + actor_rollout_ref.rollout.data_parallel_size=${gen_dp} + +actor_rollout_ref.rollout.engine_kwargs.sglang.enable_dp_attention=False + # Performance Optimization + +actor_rollout_ref.rollout.engine_kwargs.sglang.chunked_prefill_size=-1 + actor_rollout_ref.rollout.enforce_eager=False + # Validation Generation + actor_rollout_ref.rollout.val_kwargs.n=1 + actor_rollout_ref.rollout.val_kwargs.do_sample=True + actor_rollout_ref.rollout.val_kwargs.top_p=1.0 + actor_rollout_ref.rollout.val_kwargs.top_k=-1 + actor_rollout_ref.rollout.val_kwargs.temperature=1.0 +) + +TRAINER_CONFIG=( + # Logger Configuration + trainer.logger='["console"]' + # Project Settings + trainer.project_name="${project_name}" + trainer.experiment_name="${exp_name}" + # Hardware Configuration + trainer.nnodes="${NNODES}" + trainer.n_gpus_per_node="${NPUS_PER_NODE}" + trainer.device='npu' + # Training Schedule + trainer.total_epochs=15 + trainer.val_before_train=False + trainer.test_freq=-1 + trainer.save_freq=-1 + # Checkpoint Directory + trainer.default_local_dir="${CKPTS_DIR}" +) + +# profiling configuration +PROF_CONFIG=( + global_profiler.tool=npu + global_profiler.steps=null + global_profiler.save_path=/profpath + actor_rollout_ref.actor.profiler.enable=True + actor_rollout_ref.actor.profiler.ranks="[0]" + actor_rollout_ref.actor.profiler.all_ranks=False + actor_rollout_ref.actor.profiler.tool_config.npu.discrete=True + actor_rollout_ref.actor.profiler.tool_config.npu.contents=['npu','cpu'] + actor_rollout_ref.actor.profiler.tool_config.npu.level=level0 + actor_rollout_ref.actor.profiler.tool_config.npu.analysis=True + actor_rollout_ref.rollout.profiler.enable=True + actor_rollout_ref.rollout.profiler.ranks="[0]" + actor_rollout_ref.rollout.profiler.all_ranks=False +) + +python3 -m verl.trainer.main_ppo \ + --config-path=config \ + --config-name='ppo_trainer.yaml' \ + model_engine=mindspeed \ + "${DATA_CONFIG[@]}" \ + "${MODEL_CONFIG[@]}" \ + "${ACTOR_CONFIG[@]}" \ + "${REF_CONFIG[@]}" \ + "${ROLLOUT_CONFIG[@]}" \ + "${ALGORITHM_CONFIG[@]}" \ + "${TRAINER_CONFIG[@]}" \ + "${PROF_CONFIG[@]}" \ + "$@" \ No newline at end of file diff --git a/verl/examples/ascend_extras/grpo_trainer/run_qwen3_next_80b_fsdp.sh b/verl/examples/ascend_extras/grpo_trainer/run_qwen3_next_80b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..b18bb18e28104f1ccf1c41e95d6df2499e19ffd4 --- /dev/null +++ b/verl/examples/ascend_extras/grpo_trainer/run_qwen3_next_80b_fsdp.sh @@ -0,0 +1,175 @@ +#!/usr/bin/env bash +set -xeuo pipefail + +project_name="verl_grpo_qwen3-next-80b" +experiment_name="Qwen3_Next_80B_Instruct" + +# Paths +WORK_DIR=${WORK_DIR:-"${HOME}/verl"} +MODEL_PATH=${WORK_DIR}/Qwen3-Next-80B-A3B-Instruct +TRAIN_FILE=${WORK_DIR}/datasets/dapo-math-17k/dapo-math-17k.parquet +TEST_FILE=${WORK_DIR}/datasets/aime/aime-2024.parquet + +# algorithm +adv_estimator=grpo + +use_kl_in_reward=False +kl_coef=0.0 +use_kl_loss=True +kl_loss_coef=0.001 + +clip_ratio_low=0.2 +clip_ratio_high=0.28 + +temperature=1.0 +top_p=1.0 +top_k=-1 # 0 for HF rollout, -1 for vLLM rollout +val_top_p=0.7 + +# batch +train_batch_size=16 +rollout_n=16 +ppo_mini_batch_size=8 + +# length +max_prompt_length=$((1024 * 2)) +max_response_length=$((1024 * 20)) + +# algorithm +learning_rate=1e-6 +warmup_steps=0 +# enable_filter_groups=True + +# performance +sp_size=8 +gen_tp=4 +use_dynamic_bsz=True +actor_ppo_max_token_len=$(((max_prompt_length + max_response_length) / sp_size)) +infer_ppo_max_token_len=$(((max_prompt_length + max_response_length) / sp_size)) +offload=True + +DATA=( + data.train_files="${TRAIN_FILE}" + data.val_files="${TEST_FILE}" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.truncation='error' +) + +ACTOR=( + actor_rollout_ref.actor.strategy=fsdp2 + actor_rollout_ref.nccl_timeout=14400 + + # fsdp + actor_rollout_ref.actor.fsdp_config.use_orig_params=True + actor_rollout_ref.actor.fsdp_config.model_dtype=bfloat16 + actor_rollout_ref.actor.fsdp_config.param_offload=${offload} + actor_rollout_ref.actor.fsdp_config.optimizer_offload=${offload} + actor_rollout_ref.actor.fsdp_config.forward_prefetch=False + actor_rollout_ref.actor.fsdp_config.fsdp_size=-1 + +actor_rollout_ref.actor.fsdp_config.mixed_precision.reduce_dtype=bf16 + + # optimizer + actor_rollout_ref.actor.optim.lr=${learning_rate} + actor_rollout_ref.actor.optim.lr_warmup_steps=${warmup_steps} + + # ppo config + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} + actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} + + # entropy + actor_rollout_ref.actor.entropy_checkpointing=True + actor_rollout_ref.actor.entropy_from_logits_with_chunking=True + + actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} + actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} + actor_rollout_ref.actor.clip_ratio_c=10.0 + + actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} + actor_rollout_ref.actor.use_torch_compile=False +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=0.8 + actor_rollout_ref.rollout.load_format=auto + actor_rollout_ref.rollout.enforce_eager=True + actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) + actor_rollout_ref.rollout.calculate_log_probs=True + + actor_rollout_ref.rollout.temperature=${temperature} + actor_rollout_ref.rollout.top_p=${top_p} + actor_rollout_ref.rollout.top_k=${top_k} + actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} + actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} + actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} + actor_rollout_ref.rollout.val_kwargs.do_sample=True + actor_rollout_ref.rollout.val_kwargs.n=1 + + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} +) + +REF=( + actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} + actor_rollout_ref.ref.use_torch_compile=False + actor_rollout_ref.ref.fsdp_config.param_offload=${offload} + actor_rollout_ref.ref.fsdp_config.optimizer_offload=${offload} + actor_rollout_ref.ref.fsdp_config.forward_prefetch=False + + actor_rollout_ref.ref.entropy_checkpointing=True + actor_rollout_ref.ref.entropy_from_logits_with_chunking=True + + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} +) + +TRAINER=( + trainer.logger='["console"]' + trainer.project_name="${project_name}" + trainer.experiment_name="${experiment_name}" + trainer.n_gpus_per_node=16 + trainer.nnodes=4 + trainer.val_before_train=False + trainer.save_freq=5 + trainer.test_freq=-1 + trainer.total_epochs=1 + trainer.device=npu +) + +MODEL=( + actor_rollout_ref.model.path=${MODEL_PATH} + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_activation_offload=${offload} +) + +ALGORITHM=( + algorithm.adv_estimator=${adv_estimator} + algorithm.use_kl_in_reward=${use_kl_in_reward} + algorithm.kl_ctrl.kl_coef=${kl_coef} +) + +# ========================================================= +echo "Starting Training with:" +echo "Project: ${project_name}, Exp: ${experiment_name}" +echo "Rollout N: ${rollout_n}, Batch Size: ${train_batch_size}, LR: ${learning_rate}" + + +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${ALGORITHM[@]}" \ + "${MODEL[@]}" \ No newline at end of file diff --git a/verl/examples/cispo_trainer/README.md b/verl/examples/cispo_trainer/README.md new file mode 100644 index 0000000000000000000000000000000000000000..5ed4afd45e95c233287ec5e467aff4970d4cdfed --- /dev/null +++ b/verl/examples/cispo_trainer/README.md @@ -0,0 +1,17 @@ +# CISPO + +CISPO (Clipped IS-weight Policy Optimization) is a policy-loss variant that decouples the lower/upper clip ratios to stabilize IS-ratio-weighted updates, used in MiniMax-M1. + +Reference: [MiniMax-M1: Scaling Test-Time Compute Efficiently with Lightning Attention](https://arxiv.org/abs/2506.13585). + +## Canonical Scripts + +| Script | Infer | Train | Platform | +|--------------------------------------|-------|-------|----------| +| `run_qwen3_8b_fsdp.sh` | vLLM | FSDP | NVIDIA | + +## Key Flags + +- `actor_rollout_ref.actor.policy_loss.loss_mode=cispo` +- `actor_rollout_ref.actor.clip_ratio_low=10` (effectively unclamped on lower side) +- `actor_rollout_ref.actor.clip_ratio_high=0.2` diff --git a/verl/examples/cispo_trainer/run_qwen3_8b_fsdp.sh b/verl/examples/cispo_trainer/run_qwen3_8b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..5d03fd5f92e35d8664b34b8b6abe4c4fc2279deb --- /dev/null +++ b/verl/examples/cispo_trainer/run_qwen3_8b_fsdp.sh @@ -0,0 +1,120 @@ +#!/usr/bin/env bash +# CISPO | text | vLLM rollout | FSDP training | NVIDIA GPUs +# CISPO is a policy-loss variant on top of GRPO that decouples high/low clip ratios. + +set -xeuo pipefail + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-8B} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +train_batch_size=${TRAIN_BATCH_SIZE:-1024} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-256} +max_prompt_length=${MAX_PROMPT_LENGTH:-1024} +max_response_length=${MAX_RESPONSE_LENGTH:-2048} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-24576} + +actor_lr=${ACTOR_LR:-1e-6} +kl_loss_coef=${KL_LOSS_COEF:-0.001} +entropy_coeff=${ENTROPY_COEFF:-0} + +clip_ratio_low=${CLIP_RATIO_LOW:-10} +clip_ratio_high=${CLIP_RATIO_HIGH:-0.2} + +rollout_tp=${ROLLOUT_TP:-2} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.6} +rollout_n=${ROLLOUT_N:-5} + +total_epochs=${TOTAL_EPOCHS:-15} +save_freq=${SAVE_FREQ:-20} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_cispo_gsm8k_math} +experiment_name=${EXPERIMENT_NAME:-qwen3_8b_vllm_fsdp} +# ---- end user-adjustable ---- + +gsm8k_train=$HOME/data/gsm8k/train.parquet +gsm8k_test=$HOME/data/gsm8k/test.parquet +math_train=$HOME/data/math/train.parquet +math_test=$HOME/data/math/test.parquet + +train_files="['$gsm8k_train', '$math_train']" +val_files="['$gsm8k_test', '$math_test']" +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files="$train_files" + data.val_files="$val_files" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.policy_loss.loss_mode=cispo + actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} + actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.fsdp_config.param_offload=False + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.fsdp_config.param_offload=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/data_preprocess/aime2024_multiturn_w_tool.py b/verl/examples/data_preprocess/aime2024_multiturn_w_tool.py new file mode 100644 index 0000000000000000000000000000000000000000..76cdd0576d3801118b160b850bcbd8d2fe6723b1 --- /dev/null +++ b/verl/examples/data_preprocess/aime2024_multiturn_w_tool.py @@ -0,0 +1,79 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2024 SGLang Team +# Copyright 2025 ModelBest Inc. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Preprocess the DAPO-Math-17k dataset to multiturn format +""" + +import argparse +import os + +import datasets + +from verl.utils.hdfs_io import copy, makedirs + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--local_dir", default=None, help="The save directory for the preprocessed dataset.") + parser.add_argument("--hdfs_dir", default=None) + parser.add_argument("--local_dataset_path", default=None, help="The local path to the raw dataset, if it exists.") + parser.add_argument( + "--local_save_dir", default="~/data/retool_aime2024", help="The save directory for the preprocessed dataset." + ) + + args = parser.parse_args() + local_dataset_path = args.local_dataset_path + + data_path = "BytedTsinghua-SIA/AIME-2024" + + if local_dataset_path is not None: + dataset = datasets.load_dataset(local_dataset_path, "default") + else: + dataset = datasets.load_dataset(data_path, "default") + + train_dataset = dataset["train"] + + # add a row to each data item that represents a unique id + def make_map_fn(split): + def process_fn(example, idx): + orig_extra_info = example.pop("extra_info") + extra_info = orig_extra_info.copy() + extra_info["need_tools_kwargs"] = True + extra_info["tools_kwargs"] = { + "code_interpreter": { + "create_kwargs": { + "ground_truth": example["reward_model"]["ground_truth"], + }, + }, + } + example["extra_info"] = extra_info + return example + + return process_fn + + train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True) + + hdfs_dir = args.hdfs_dir + local_save_dir = args.local_dir + if local_save_dir is not None: + print("Warning: Argument 'local_dir' is deprecated. Please use 'local_save_dir' instead.") + else: + local_save_dir = args.local_save_dir + + train_dataset.to_parquet(os.path.join(local_save_dir, "train.parquet")) + + if hdfs_dir is not None: + makedirs(hdfs_dir) + copy(src=local_save_dir, dst=hdfs_dir) diff --git a/verl/examples/data_preprocess/dapo_multiturn_w_tool.py b/verl/examples/data_preprocess/dapo_multiturn_w_tool.py new file mode 100644 index 0000000000000000000000000000000000000000..aab356f41bf38e789a31f1ee879ce9beb8b0aa40 --- /dev/null +++ b/verl/examples/data_preprocess/dapo_multiturn_w_tool.py @@ -0,0 +1,79 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2024 SGLang Team +# Copyright 2025 ModelBest Inc. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Preprocess the DAPO-Math-17k dataset to multiturn format +""" + +import argparse +import os + +import datasets + +from verl.utils.hdfs_io import copy, makedirs + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--local_dir", default=None, help="The save directory for the preprocessed dataset.") + parser.add_argument("--hdfs_dir", default=None) + parser.add_argument("--local_dataset_path", default=None, help="The local path to the raw dataset, if it exists.") + parser.add_argument( + "--local_save_dir", default="~/data/retool_dapo", help="The save directory for the preprocessed dataset." + ) + + args = parser.parse_args() + local_dataset_path = args.local_dataset_path + + data_path = "BytedTsinghua-SIA/DAPO-Math-17k" + + if local_dataset_path is not None: + dataset = datasets.load_dataset(local_dataset_path, "default") + else: + dataset = datasets.load_dataset(data_path, "default") + + train_dataset = dataset["train"] + + # add a row to each data item that represents a unique id + def make_map_fn(split): + def process_fn(example, idx): + orig_extra_info = example.pop("extra_info") + extra_info = orig_extra_info.copy() + extra_info["need_tools_kwargs"] = True + extra_info["tools_kwargs"] = { + "code_interpreter": { + "create_kwargs": { + "ground_truth": example["reward_model"]["ground_truth"], + }, + }, + } + example["extra_info"] = extra_info + return example + + return process_fn + + train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True) + + hdfs_dir = args.hdfs_dir + local_save_dir = args.local_dir + if local_save_dir is not None: + print("Warning: Argument 'local_dir' is deprecated. Please use 'local_save_dir' instead.") + else: + local_save_dir = args.local_save_dir + + train_dataset.to_parquet(os.path.join(local_save_dir, "train.parquet")) + + if hdfs_dir is not None: + makedirs(hdfs_dir) + copy(src=local_save_dir, dst=hdfs_dir) diff --git a/verl/examples/data_preprocess/full_hh_rlhf.py b/verl/examples/data_preprocess/full_hh_rlhf.py new file mode 100644 index 0000000000000000000000000000000000000000..4e8a148df1e322f476cedffe4eadc5ae6ee9b6f1 --- /dev/null +++ b/verl/examples/data_preprocess/full_hh_rlhf.py @@ -0,0 +1,161 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +- Preprocess data and split the training set into 75% for training RM and 25% for validting RM. +- All the training data is used to train SFT and RL. +- Both chosen and rejected is used to train SFT +""" + +import argparse +import os + +import pandas as pd +from datasets import load_dataset +from tqdm.auto import tqdm + +from verl.utils.fs import copy, makedirs + + +def generate_sft_dataset(target_hdfs_path_dir, local_dir="~/data/full_hh_rlh/sft", local_dataset_path=None): + if local_dataset_path is not None: + dataset = load_dataset(local_dataset_path) + else: + dataset = load_dataset("Dahoas/full-hh-rlhf") + output = {"prompt": [], "response": []} + for data in tqdm(dataset["train"]): + # add chosen + output["prompt"].append(data["prompt"]) + output["response"].append(data["chosen"]) + + # add rejection + output["prompt"].append(data["prompt"]) + output["response"].append(data["rejected"]) + + df = pd.DataFrame(output) + + local_dir = os.path.expanduser(local_dir) + os.makedirs(local_dir, exist_ok=True) + + local_path = os.path.join(local_dir, "train.parquet") + + df.to_parquet(path=local_path) + + if target_hdfs_path_dir is not None: + hdfs_dir = target_hdfs_path_dir + "/" + "train.parquet" + makedirs(hdfs_dir) + + copy(local_path, hdfs_dir) + + +def generate_rm_dataset(target_hdfs_path_dir, local_dir="~/data/full_hh_rlh/rm", local_dataset_path=None): + if local_dataset_path is not None: + train_dataset = load_dataset(local_dataset_path, split="train[:75%]") + test_dataset = load_dataset(local_dataset_path, split="train[-25%:]") + else: + train_dataset = load_dataset("Dahoas/full-hh-rlhf", split="train[:75%]") + test_dataset = load_dataset("Dahoas/full-hh-rlhf", split="train[-25%:]") + + local_dir = os.path.expanduser(local_dir) + os.makedirs(local_dir, exist_ok=True) + + for dataset, name in zip([train_dataset, test_dataset], ["train", "test"], strict=True): + output = {"prompt": [], "chosen": [], "rejected": []} + for data in tqdm(dataset): + # add chosen + output["prompt"].append(data["prompt"]) + output["chosen"].append(data["chosen"]) + output["rejected"].append(data["rejected"]) + + df = pd.DataFrame(output) + + local_path = os.path.join(local_dir, name + ".parquet") + + df.to_parquet(path=local_path) + + if target_hdfs_path_dir is not None: + hdfs_dir = target_hdfs_path_dir + "/" + name + ".parquet" + makedirs(hdfs_dir) + + copy(local_path, hdfs_dir) + + +def generate_rl_dataset(target_hdfs_path_dir, local_dir="~/data/full_hh_rlhf/rl", local_dataset_path=None): + if local_dataset_path is not None: + dataset = load_dataset(local_dataset_path) + else: + dataset = load_dataset("Dahoas/full-hh-rlhf") + train_dataset = dataset["train"] + + data_source = "Dahoas/full-hh-rlhf" + + # add a row to each data item that represents a unique id + def make_map_fn(split): + def process_fn(example, idx): + prompt = example.pop("prompt") + response = example.pop("response") + + data = { + "data_source": data_source, + "prompt": [{"role": "user", "content": prompt}], + "ability": "alignment", + "reward_model": { + "style": "model", + "ground_truth": response, # should not be used + }, + "extra_info": {"split": split, "index": idx}, + } + return data + + return process_fn + + train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True) + local_dir = os.path.expanduser(local_dir) + local_path = os.path.join(local_dir, "train.parquet") + train_dataset.to_parquet(local_path) + + if target_hdfs_path_dir is not None: + hdfs_dir = target_hdfs_path_dir + "/" + "train.parquet" + makedirs(hdfs_dir) + + copy(local_path, hdfs_dir) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--split", type=str, choices=["sft", "rm", "rl"], required=True) + parser.add_argument("--local_dir", default=None, help="The save directory for the preprocessed dataset.") + parser.add_argument("--hdfs_dir", type=str, required=False, default=None) + parser.add_argument("--local_dataset_path", default=None, help="The local path to the raw dataset, if it exists.") + parser.add_argument( + "--local_save_dir", + type=str, + default="~/data/full_hh_rlhf", + help="The save directory for the preprocessed dataset.", + ) + + args = parser.parse_args() + local_save_dir = args.local_dir + if local_save_dir is not None: + print("Warning: Argument 'local_dir' is deprecated. Please use 'local_save_dir' instead.") + else: + local_save_dir = args.local_save_dir + + if args.split == "sft": + generate_sft_dataset(args.hdfs_dir, os.path.join(local_save_dir, args.split), args.local_dataset_path) + elif args.split == "rm": + generate_rm_dataset(args.hdfs_dir, os.path.join(local_save_dir, args.split), args.local_dataset_path) + elif args.split == "rl": + generate_rl_dataset(args.hdfs_dir, os.path.join(local_save_dir, args.split), args.local_dataset_path) + else: + raise NotImplementedError diff --git a/verl/examples/data_preprocess/geo3k.py b/verl/examples/data_preprocess/geo3k.py new file mode 100644 index 0000000000000000000000000000000000000000..ba84fd3fc440761a200d0fbdea1535bfe9889b45 --- /dev/null +++ b/verl/examples/data_preprocess/geo3k.py @@ -0,0 +1,102 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Preprocess the Geometry3k dataset to parquet format +""" + +import argparse +import os + +import datasets + +from verl.utils.hdfs_io import copy, makedirs + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--local_dir", default=None) + parser.add_argument("--hdfs_dir", default=None) + parser.add_argument("--local_dataset_path", default=None, help="The local path to the raw dataset, if it exists.") + parser.add_argument( + "--local_save_dir", default="~/data/geo3k", help="The save directory for the preprocessed dataset." + ) + + args = parser.parse_args() + local_dataset_path = args.local_dataset_path + + data_source = "hiyouga/geometry3k" + + if local_dataset_path is not None: + dataset = datasets.load_dataset( + local_dataset_path, + ) + else: + dataset = datasets.load_dataset( + data_source, + ) + + train_dataset = dataset["train"] + test_dataset = dataset["test"] + + instruction_following = ( + r"You FIRST think about the reasoning process as an internal monologue and then provide the final answer. " + r"The reasoning process MUST BE enclosed within tags. " + r"The final answer MUST BE put in \boxed{}." + ) + + # add a row to each data item that represents a unique id + def make_map_fn(split): + def process_fn(example, idx): + problem = example.pop("problem") + prompt = problem + " " + instruction_following + answer = example.pop("answer") + images = example.pop("images") + + data = { + "data_source": data_source, + "prompt": [ + { + "role": "user", + "content": prompt, + } + ], + "images": images, + "ability": "math", + "reward_model": {"style": "rule", "ground_truth": answer}, + "extra_info": { + "split": split, + "index": idx, + "answer": answer, + "question": problem, + }, + } + return data + + return process_fn + + train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True, num_proc=8) + test_dataset = test_dataset.map(function=make_map_fn("test"), with_indices=True, num_proc=8) + + hdfs_dir = args.hdfs_dir + local_save_dir = args.local_dir + if local_save_dir is not None: + print("Warning: Argument 'local_dir' is deprecated. Please use 'local_save_dir' instead.") + else: + local_save_dir = args.local_save_dir + + train_dataset.to_parquet(os.path.join(local_save_dir, "train.parquet")) + test_dataset.to_parquet(os.path.join(local_save_dir, "test.parquet")) + + if hdfs_dir is not None: + makedirs(hdfs_dir) + copy(src=local_save_dir, dst=hdfs_dir) diff --git a/verl/examples/data_preprocess/geo3k_multiturn_w_tool.py b/verl/examples/data_preprocess/geo3k_multiturn_w_tool.py new file mode 100644 index 0000000000000000000000000000000000000000..53c7197f9d2d00ccf256aee57e2de1847a926725 --- /dev/null +++ b/verl/examples/data_preprocess/geo3k_multiturn_w_tool.py @@ -0,0 +1,120 @@ +# Copyright 2023-2025 SGLang Team +# Copyright Amazon.com, Inc. or its affiliates. +# Copyright 2025 Reallm Labs Ltd. or its affiliates +# Copyright 2025 ModelBest Inc. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Preprocess the Geometry3k dataset to parquet format +""" + +import argparse +import os + +import datasets + +from verl.utils.hdfs_io import copy, makedirs + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--local_dir", default=None, help="The save directory for the preprocessed dataset.") + parser.add_argument("--hdfs_dir", default=None) + parser.add_argument("--local_dataset_path", default=None, help="The local path to the raw dataset, if it exists.") + parser.add_argument( + "--local_save_dir", + default="~/data/geo3k_multiturn_w_tool", + help="The save directory for the preprocessed dataset.", + ) + + args = parser.parse_args() + local_dataset_path = args.local_dataset_path + + data_source = "hiyouga/geometry3k" + + if local_dataset_path is not None: + dataset = datasets.load_dataset(local_dataset_path) + else: + dataset = datasets.load_dataset(data_source) + + train_dataset = dataset["train"] + test_dataset = dataset["test"] + + instruction_following = ( + r"You FIRST think about the reasoning process as an internal monologue and then provide the final answer. " + r"The reasoning process MUST BE enclosed within tags. " + r"The final answer MUST BE put in \boxed{}." + ) + + # add a row to each data item that represents a unique id + def make_map_fn(split): + def process_fn(example, idx): + problem = example.pop("problem") + prompt = problem + " " + instruction_following + answer = example.pop("answer") + images = example.pop("images") + data = { + "data_source": data_source, + "prompt": [ + { + "role": "system", + "content": ( + "You are a math expert. You are given a question and you need to solve it step by step. " + "Reasoning step by step before any tool call. " + "You should use the `calc_geo3k_reward` tool after step by step solving the question, " + "before generate final answer at least once and refine your answer if necessary. " + ), + }, + { + "role": "user", + "content": prompt, + }, + ], + "images": images, + "ability": "math", + "reward_model": {"style": "rule", "ground_truth": answer}, + "extra_info": { + "split": split, + "index": idx, + "answer": answer, + "question": problem, + "need_tools_kwargs": True, + "tools_kwargs": { + "calc_geo3k_reward": { + "create_kwargs": {"ground_truth": answer}, + # "execute_kwargs": {}, + # "calc_reward_kwargs": {}, + # "release_kwargs": {}, + }, + }, + }, + } + return data + + return process_fn + + train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True, num_proc=8) + test_dataset = test_dataset.map(function=make_map_fn("test"), with_indices=True, num_proc=8) + + hdfs_dir = args.hdfs_dir + local_save_dir = args.local_dir + if local_save_dir is not None: + print("Warning: Argument 'local_dir' is deprecated. Please use 'local_save_dir' instead.") + else: + local_save_dir = args.local_save_dir + + train_dataset.to_parquet(os.path.join(local_save_dir, "train.parquet")) + test_dataset.to_parquet(os.path.join(local_save_dir, "test.parquet")) + if hdfs_dir is not None: + makedirs(hdfs_dir) + copy(src=local_save_dir, dst=hdfs_dir) diff --git a/verl/examples/data_preprocess/gsm8k.py b/verl/examples/data_preprocess/gsm8k.py new file mode 100644 index 0000000000000000000000000000000000000000..1656cdbc896a8f14fc7e09705d36335f52165533 --- /dev/null +++ b/verl/examples/data_preprocess/gsm8k.py @@ -0,0 +1,105 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Preprocess the GSM8k dataset to parquet format +""" + +import argparse +import os +import re + +import datasets + +from verl.utils.hdfs_io import copy, makedirs + + +def extract_solution(solution_str): + solution = re.search("#### (\\-?[0-9\\.\\,]+)", solution_str) + assert solution is not None + final_solution = solution.group(0) + final_solution = final_solution.split("#### ")[1].replace(",", "") + return final_solution + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--local_dir", default=None, help="The save directory for the preprocessed dataset.") + parser.add_argument("--hdfs_dir", default=None) + parser.add_argument("--local_dataset_path", default=None, help="The local path to the raw dataset, if it exists.") + parser.add_argument( + "--local_save_dir", default="~/data/gsm8k", help="The save directory for the preprocessed dataset." + ) + + args = parser.parse_args() + local_dataset_path = args.local_dataset_path + + data_source = "openai/gsm8k" + + if local_dataset_path is not None: + dataset = datasets.load_dataset(local_dataset_path, "main") + else: + dataset = datasets.load_dataset(data_source, "main") + + train_dataset = dataset["train"] + test_dataset = dataset["test"] + + instruction_following = 'Let\'s think step by step and output the final answer after "####".' + + # add a row to each data item that represents a unique id + def make_map_fn(split): + def process_fn(example, idx): + question_raw = example.pop("question") + + question = question_raw + " " + instruction_following + + answer_raw = example.pop("answer") + solution = extract_solution(answer_raw) + data = { + "data_source": data_source, + "prompt": [ + { + "role": "user", + "content": question, + } + ], + "ability": "math", + "reward_model": {"style": "rule", "ground_truth": solution}, + "extra_info": { + "split": split, + "index": idx, + "answer": answer_raw, + "question": question_raw, + }, + } + return data + + return process_fn + + train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True) + test_dataset = test_dataset.map(function=make_map_fn("test"), with_indices=True) + + hdfs_dir = args.hdfs_dir + local_save_dir = args.local_dir + if local_save_dir is not None: + print("Warning: Argument 'local_dir' is deprecated. Please use 'local_save_dir' instead.") + else: + local_save_dir = args.local_save_dir + + train_dataset.to_parquet(os.path.join(local_save_dir, "train.parquet")) + test_dataset.to_parquet(os.path.join(local_save_dir, "test.parquet")) + + if hdfs_dir is not None: + makedirs(hdfs_dir) + + copy(src=local_save_dir, dst=hdfs_dir) diff --git a/verl/examples/data_preprocess/gsm8k_multiturn_sft.py b/verl/examples/data_preprocess/gsm8k_multiturn_sft.py new file mode 100644 index 0000000000000000000000000000000000000000..4589362f933aa95493fdd98ce965eb810180c98a --- /dev/null +++ b/verl/examples/data_preprocess/gsm8k_multiturn_sft.py @@ -0,0 +1,102 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Preprocess the GSM8k dataset to parquet format +""" + +import argparse +import os +import re + +import datasets + +from verl.utils.hdfs_io import copy, makedirs + + +def extract_solution(solution_str): + solution = re.search("#### (\\-?[0-9\\.\\,]+)", solution_str) + assert solution is not None + final_solution = solution.group(0) + final_solution = final_solution.split("#### ")[1].replace(",", "") + return final_solution + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--local_dir", default=None) + parser.add_argument("--local_dataset_path", default=None, help="The local path to the raw dataset, if it exists.") + parser.add_argument( + "--local_save_dir", default="~/data/gsm8k_sft", help="The save directory for the preprocessed dataset." + ) + parser.add_argument("--hdfs_dir", default=None) + + args = parser.parse_args() + local_dataset_path = args.local_dataset_path + + data_source = "openai/gsm8k" + + if local_dataset_path is not None: + dataset = datasets.load_dataset(local_dataset_path, "main") + else: + dataset = datasets.load_dataset(data_source, "main") + + train_dataset = dataset["train"] + test_dataset = dataset["test"] + + instruction_following = 'Let\'s think step by step and output the final answer after "####".' + + # add a row to each data item that represents a unique id + def make_map_fn(split): + def process_fn(example, idx): + question_raw = example.pop("question") + + question = question_raw + " " + instruction_following + + answer_raw = example.pop("answer") + data = { + "messages": [ + { + "role": "user", + "content": question, + }, + { + "role": "assistant", + "content": answer_raw, + }, + ], + } + return data + + return process_fn + + train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True) + test_dataset = test_dataset.map(function=make_map_fn("test"), with_indices=True) + + hdfs_dir = args.hdfs_dir + + local_save_dir = args.local_dir + if local_save_dir is not None: + print("Warning: Argument 'local_dir' is deprecated. Please use 'local_save_dir' instead.") + else: + local_save_dir = args.local_save_dir + + local_save_dir = os.path.expanduser(local_save_dir) + + train_dataset.to_parquet(os.path.join(local_save_dir, "train.parquet")) + test_dataset.to_parquet(os.path.join(local_save_dir, "test.parquet")) + + if hdfs_dir is not None: + makedirs(hdfs_dir) + + copy(src=local_save_dir, dst=hdfs_dir) diff --git a/verl/examples/data_preprocess/gsm8k_multiturn_w_tool.py b/verl/examples/data_preprocess/gsm8k_multiturn_w_tool.py new file mode 100644 index 0000000000000000000000000000000000000000..8ce02127dfd3dace9bbd3bd4cb12dc12e18db289 --- /dev/null +++ b/verl/examples/data_preprocess/gsm8k_multiturn_w_tool.py @@ -0,0 +1,125 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2024 SGLang Team +# Copyright 2025 ModelBest Inc. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Preprocess the GSM8k dataset to parquet format +""" + +import argparse +import os +import re + +import datasets + +from verl.utils.hdfs_io import copy, makedirs + + +def extract_solution(solution_str): + solution = re.search("#### (\\-?[0-9\\.\\,]+)", solution_str) + assert solution is not None + final_solution = solution.group(0) + final_solution = final_solution.split("#### ")[1].replace(",", "") + return final_solution + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--local_dir", default=None, help="The save directory for the preprocessed dataset.") + parser.add_argument("--hdfs_dir", default=None) + parser.add_argument("--local_dataset_path", default=None, help="The local path to the raw dataset, if it exists.") + parser.add_argument( + "--local_save_dir", default="~/data/gsm8k", help="The save directory for the preprocessed dataset." + ) + + args = parser.parse_args() + local_dataset_path = args.local_dataset_path + + data_source = "openai/gsm8k" + + if local_dataset_path is not None: + dataset = datasets.load_dataset(local_dataset_path, "main") + else: + dataset = datasets.load_dataset(data_source, "main") + + train_dataset = dataset["train"] + test_dataset = dataset["test"] + + instruction_following = "Let's think step by step and output the final answer after `####`." + + # add a row to each data item that represents a unique id + def make_map_fn(split): + def process_fn(example, idx): + question_raw = example.pop("question") + + question = question_raw + " " + instruction_following + + answer_raw = example.pop("answer") + solution = extract_solution(answer_raw) + data = { + "data_source": data_source, + "prompt": [ + { + "role": "system", + "content": ( + "You are a math expert. You are given a question and you need to solve it step by step. " + "Reasoning step by step before any tool call. " + "You should use the `calc_gsm8k_reward` tool after step by step solving the question, " + "before generate final answer at least once and refine your answer if necessary. " + "Put your final answer in the format of `#### `." + ), + }, + { + "role": "user", + "content": question, + }, + ], + "ability": "math", + "reward_model": {"style": "rule", "ground_truth": solution}, + "extra_info": { + "split": split, + "index": idx, + "answer": answer_raw, + "question": question_raw, + "need_tools_kwargs": True, + "tools_kwargs": { + "calc_gsm8k_reward": { + "create_kwargs": {"ground_truth": solution}, + # "execute_kwargs": {}, + # "calc_reward_kwargs": {}, + # "release_kwargs": {}, + }, + }, + }, + } + return data + + return process_fn + + train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True) + test_dataset = test_dataset.map(function=make_map_fn("test"), with_indices=True) + + hdfs_dir = args.hdfs_dir + local_save_dir = args.local_dir + if local_save_dir is not None: + print("Warning: Argument 'local_dir' is deprecated. Please use 'local_save_dir' instead.") + else: + local_save_dir = args.local_save_dir + + train_dataset.to_parquet(os.path.join(local_save_dir, "train.parquet")) + test_dataset.to_parquet(os.path.join(local_save_dir, "test.parquet")) + + if hdfs_dir is not None: + makedirs(hdfs_dir) + copy(src=local_save_dir, dst=hdfs_dir) diff --git a/verl/examples/data_preprocess/gsm8k_tool_agent_loop.py b/verl/examples/data_preprocess/gsm8k_tool_agent_loop.py new file mode 100644 index 0000000000000000000000000000000000000000..0bb1d9ef08d055fce3047f9b502c1e6cb84884d6 --- /dev/null +++ b/verl/examples/data_preprocess/gsm8k_tool_agent_loop.py @@ -0,0 +1,126 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2024 SGLang Team +# Copyright 2025 ModelBest Inc. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Preprocess the GSM8k dataset to parquet format +""" + +import argparse +import os +import re + +import datasets + +from verl.utils.hdfs_io import copy, makedirs + + +def extract_solution(solution_str): + solution = re.search("#### (\\-?[0-9\\.\\,]+)", solution_str) + assert solution is not None + final_solution = solution.group(0) + final_solution = final_solution.split("#### ")[1].replace(",", "") + return final_solution + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--local_dir", default=None, help="The save directory for the preprocessed dataset.") + parser.add_argument("--hdfs_dir", default=None) + parser.add_argument("--local_dataset_path", default=None, help="The local path to the raw dataset, if it exists.") + parser.add_argument( + "--local_save_dir", default="~/data/gsm8k", help="The save directory for the preprocessed dataset." + ) + + args = parser.parse_args() + local_dataset_path = args.local_dataset_path + + data_source = "openai/gsm8k" + + if local_dataset_path is not None: + dataset = datasets.load_dataset(local_dataset_path, "main") + else: + dataset = datasets.load_dataset(data_source, "main") + + train_dataset = dataset["train"] + test_dataset = dataset["test"] + + instruction_following = "Let's think step by step and output the final answer after `####`." + + # add a row to each data item that represents a unique id + def make_map_fn(split): + def process_fn(example, idx): + question_raw = example.pop("question") + + question = question_raw + " " + instruction_following + + answer_raw = example.pop("answer") + solution = extract_solution(answer_raw) + data = { + "data_source": data_source, + "agent_name": "tool_agent", + "prompt": [ + { + "role": "system", + "content": ( + "You are a math expert. You are given a question and you need to solve it step by step. " + "Reasoning step by step before any tool call. " + "You should use the `calc_gsm8k_reward` tool after step by step solving the question, " + "before generate final answer at least once and refine your answer if necessary. " + "Put your final answer in the format of `#### `." + ), + }, + { + "role": "user", + "content": question, + }, + ], + "ability": "math", + "reward_model": {"style": "rule", "ground_truth": solution}, + "extra_info": { + "split": split, + "index": idx, + "answer": answer_raw, + "question": question_raw, + "need_tools_kwargs": True, + "tools_kwargs": { + "calc_gsm8k_reward": { + "create_kwargs": {"ground_truth": solution}, + # "execute_kwargs": {}, + # "calc_reward_kwargs": {}, + # "release_kwargs": {}, + }, + }, + }, + } + return data + + return process_fn + + train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True) + test_dataset = test_dataset.map(function=make_map_fn("test"), with_indices=True) + + hdfs_dir = args.hdfs_dir + local_save_dir = args.local_dir + if local_save_dir is not None: + print("Warning: Argument 'local_dir' is deprecated. Please use 'local_save_dir' instead.") + else: + local_save_dir = args.local_save_dir + + train_dataset.to_parquet(os.path.join(local_save_dir, "train.parquet")) + test_dataset.to_parquet(os.path.join(local_save_dir, "test.parquet")) + + if hdfs_dir is not None: + makedirs(hdfs_dir) + copy(src=local_save_dir, dst=hdfs_dir) diff --git a/verl/examples/data_preprocess/hellaswag.py b/verl/examples/data_preprocess/hellaswag.py new file mode 100644 index 0000000000000000000000000000000000000000..dc73a810a80570d406bb727099f5524037be2370 --- /dev/null +++ b/verl/examples/data_preprocess/hellaswag.py @@ -0,0 +1,108 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Preprocess Hellaswag dataset. + +""" + +import argparse +import os +import re + +import datasets + +from verl.utils.hdfs_io import copy, makedirs + + +def preprocess(text): + text = text.strip() + # NOTE: Brackets are artifacts of the WikiHow dataset portion of HellaSwag. + text = text.replace(" [title]", ". ") + text = re.sub("\\[.*?\\]", "", text) + text = text.replace(" ", " ") + return text + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--local_dir", default=None, help="The save directory for the preprocessed dataset.") + parser.add_argument("--hdfs_dir", default=None) + parser.add_argument("--local_dataset_path", default=None, help="The local path to the raw dataset, if it exists.") + parser.add_argument( + "--local_save_dir", default="~/data/hellaswag", help="The save directory for the preprocessed dataset." + ) + + args = parser.parse_args() + local_dataset_path = args.local_dataset_path + + data_source = "Rowan/hellaswag" + + if local_dataset_path is not None: + dataset = datasets.load_dataset(local_dataset_path) + else: + dataset = datasets.load_dataset(data_source, trust_remote_code=True) + + train_dataset = dataset["train"] + val_dataset = dataset["validation"] + test_dataset = dataset["test"] + + instruction = "Please complete the following sentence.\n" + + def make_map_fn(split): + def process_fn(doc, idx): + ctx = doc["ctx_a"] + " " + doc["ctx_b"].capitalize() + query = preprocess(doc["activity_label"] + ": " + ctx) + choices = [preprocess(ending) for ending in doc["endings"]] + gold = int(doc["label"]) + + data = { + "data_source": data_source, + "prompt": [{"role": "user", "content": query}], + "ability": "nlp", + "reward_model": { + "style": "model", + "eval": "multiple_choice", # using loglikelihood + "ground_truth": gold, + "choices": choices, + }, + "extra_info": {"split": split, "index": idx}, + } + return data + + return process_fn + + # filter data that doesn't have a label + train_dataset = train_dataset.filter(lambda x: len(x["label"]) > 0) + val_dataset = val_dataset.filter(lambda x: len(x["label"]) > 0) + test_dataset = test_dataset.filter(lambda x: len(x["label"]) > 0) + + train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True) + val_dataset = val_dataset.map(function=make_map_fn("validation"), with_indices=True) + test_dataset = test_dataset.map(function=make_map_fn("test"), with_indices=True) + + hdfs_dir = args.hdfs_dir + local_save_dir = args.local_dir + if local_save_dir is not None: + print("Warning: Argument 'local_dir' is deprecated. Please use 'local_save_dir' instead.") + else: + local_save_dir = args.local_save_dir + + train_dataset.to_parquet(os.path.join(local_save_dir, "train.parquet")) + val_dataset.to_parquet(os.path.join(local_save_dir, "validation.parquet")) + test_dataset.to_parquet(os.path.join(local_save_dir, "test.parquet")) + + if hdfs_dir is not None: + makedirs(hdfs_dir) + + copy(src=local_save_dir, dst=hdfs_dir) diff --git a/verl/examples/data_preprocess/math_dataset.py b/verl/examples/data_preprocess/math_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..b23a032fb1207a47dcd1bc77194a7c1a124aad55 --- /dev/null +++ b/verl/examples/data_preprocess/math_dataset.py @@ -0,0 +1,106 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Preprocess the MATH-lighteval dataset to parquet format +""" + +import argparse +import json +import os + +import datasets + +from verl.utils.hdfs_io import copy, makedirs +from verl.utils.reward_score.math_reward import last_boxed_only_string, remove_boxed + + +def extract_solution(solution_str): + return remove_boxed(last_boxed_only_string(solution_str)) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--local_dir", default=None) + parser.add_argument("--hdfs_dir", default=None) + parser.add_argument("--local_dataset_path", default=None, help="The local path to the raw dataset, if it exists.") + parser.add_argument( + "--local_save_dir", default="~/data/math", help="The save directory for the preprocessed dataset." + ) + + args = parser.parse_args() + local_dataset_path = args.local_dataset_path + + # 'lighteval/MATH' is no longer available on huggingface. + # Use mirror repo: DigitalLearningGmbH/MATH-lighteval + data_source = "DigitalLearningGmbH/MATH-lighteval" + print(f"Loading the {data_source} dataset from huggingface...", flush=True) + if local_dataset_path is not None: + dataset = datasets.load_dataset( + local_dataset_path, + ) + else: + dataset = datasets.load_dataset( + data_source, + ) + + train_dataset = dataset["train"] + test_dataset = dataset["test"] + + instruction_following = "Let's think step by step and output the final answer within \\boxed{}." + + # add a row to each data item that represents a unique id + def make_map_fn(split): + def process_fn(example, idx): + question = example.pop("problem") + + question = question + " " + instruction_following + + answer = example.pop("solution") + solution = extract_solution(answer) + data = { + "data_source": data_source, + "prompt": [{"role": "user", "content": question}], + "ability": "math", + "reward_model": {"style": "rule", "ground_truth": solution}, + "extra_info": {"split": split, "index": idx}, + } + return data + + return process_fn + + train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True) + test_dataset = test_dataset.map(function=make_map_fn("test"), with_indices=True) + + local_save_dir = args.local_dir + if local_save_dir is not None: + print("Warning: Argument 'local_dir' is deprecated. Please use 'local_save_dir' instead.") + else: + local_save_dir = args.local_save_dir + + local_dir = os.path.expanduser(local_save_dir) + hdfs_dir = args.hdfs_dir + + train_dataset.to_parquet(os.path.join(local_dir, "train.parquet")) + test_dataset.to_parquet(os.path.join(local_dir, "test.parquet")) + # Save one example as JSON for reference + example = train_dataset[0] + with open(os.path.join(local_dir, "train_example.json"), "w") as f: + json.dump(example, f, indent=2) + example = test_dataset[0] + with open(os.path.join(local_dir, "test_example.json"), "w") as f: + json.dump(example, f, indent=2) + if hdfs_dir is not None: + makedirs(hdfs_dir) + + copy(src=local_dir, dst=hdfs_dir) diff --git a/verl/examples/data_preprocess/multiturn.py b/verl/examples/data_preprocess/multiturn.py new file mode 100644 index 0000000000000000000000000000000000000000..4bf0192b817d3acfc41f7e4a74a4a891d3ae61d6 --- /dev/null +++ b/verl/examples/data_preprocess/multiturn.py @@ -0,0 +1,125 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Create a simple multi-turn dataset for testing +""" + +import argparse +import os + +import pandas as pd + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--local_dir", default="~/data/multiturn") + parser.add_argument("--hdfs_dir", default=None) + args = parser.parse_args() + + # Create example conversations + conversations = [] + + # Conversation 1 + conversations.append( + { + "messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "What is the capital of France?"}, + {"role": "assistant", "content": "The capital of France is Paris."}, + {"role": "user", "content": "And what about Germany?"}, + {"role": "assistant", "content": "The capital of Germany is Berlin."}, + ] + } + ) + + # Conversation 2 + conversations.append( + { + "messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Can you explain quantum computing?"}, + { + "role": "assistant", + "content": "Quantum computing is a type of computing that uses quantum-mechanical phenomena, " + "such as superposition and entanglement, to perform operations on data.", + }, + {"role": "user", "content": "How is it different from classical computing?"}, + { + "role": "assistant", + "content": "Classical computing uses bits that are either 0 or 1, while quantum computing uses " + "quantum bits or qubits that can exist in multiple states simultaneously due to superposition.", + }, + ] + } + ) + + # Conversation 3 + conversations.append( + { + "messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Write a simple Python function to calculate factorial."}, + { + "role": "assistant", + "content": ( + "```python\ndef factorial(n):\n if n == 0 or n == 1:\n return 1\n else:\n " + "return n * factorial(n-1)\n```\n\nThis is a recursive function to calculate the " + "factorial of a number." + ), + }, + {"role": "user", "content": "Can you make it iterative instead?"}, + { + "role": "assistant", + "content": ( + "```python\ndef factorial(n):\n result = 1\n for i in range(1, n+1):\n " + "result *= i\n return result\n```\n\nThis is an iterative version of the factorial function." + ), + }, + ] + } + ) + + # Create train and test datasets + train_data = conversations[:2] # First 2 conversations for training + test_data = conversations[2:] # Last conversation for testing + + # Create output directory + local_dir = os.path.expanduser(args.local_dir) + os.makedirs(local_dir, exist_ok=True) + + # Save to parquet files + train_df = pd.DataFrame(train_data) + test_df = pd.DataFrame(test_data) + + train_df.to_parquet(os.path.join(local_dir, "train.parquet")) + test_df.to_parquet(os.path.join(local_dir, "test.parquet")) + + # Handle HDFS if specified + if args.hdfs_dir is not None: + try: + from verl.utils.hdfs_io import copy, makedirs + + makedirs(args.hdfs_dir) + copy(src=local_dir, dst=args.hdfs_dir) + except ImportError: + print("Warning: HDFS support not available. Skipping HDFS copy.") + + # Print statistics + print(f"Train dataset size: {len(train_df)}") + print(f"Test dataset size: {len(test_df)}") + print(f"Data saved to {local_dir}") + + +if __name__ == "__main__": + main() diff --git a/verl/examples/data_preprocess/pokemon.py b/verl/examples/data_preprocess/pokemon.py new file mode 100644 index 0000000000000000000000000000000000000000..3bbf4d4b46ee98669eaa40a9a9084f918791f50d --- /dev/null +++ b/verl/examples/data_preprocess/pokemon.py @@ -0,0 +1,75 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +""" +Preprocess the llamafactory/pokemon-gpt4o-captions dataset to parquet format +""" + +import argparse +import os + +import datasets + +from verl.utils.hdfs_io import copy, makedirs + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--local_dir", default=None) + parser.add_argument("--hdfs_dir", default=None) + parser.add_argument("--local_dataset_path", default=None, help="The local path to the raw dataset, if it exists.") + parser.add_argument( + "--local_save_dir", + default="~/data/pokemon-gpt4o-captions", + help="The save directory for the preprocessed dataset.", + ) + + args = parser.parse_args() + local_dataset_path = args.local_dataset_path + + data_source = "llamafactory/pokemon-gpt4o-captions" + + if local_dataset_path is not None: + dataset = datasets.load_dataset( + local_dataset_path, + ) + else: + dataset = datasets.load_dataset( + data_source, + ) + + def map_fn(row: dict): + messages = [] + conversation = row.pop("conversations") + for conv in conversation: + if conv["from"] == "gpt": + role = "assistant" + elif conv["from"] == "human": + role = "user" + else: + raise ValueError(f"Unknown role: {conv['from']}") + messages.append( + { + "role": role, + "content": conv["value"], + } + ) + + row["messages"] = messages + return row + + dataset = dataset["train"].map(map_fn, num_proc=16) + dataset = dataset.train_test_split(test_size=0.1) + train_dataset = dataset["train"] + test_dataset = dataset["test"] + + hdfs_dir = args.hdfs_dir + local_save_dir = args.local_dir + if local_save_dir is not None: + print("Warning: Argument 'local_dir' is deprecated. Please use 'local_save_dir' instead.") + else: + local_save_dir = args.local_save_dir + + train_dataset.to_parquet(os.path.join(local_save_dir, "train.parquet")) + test_dataset.to_parquet(os.path.join(local_save_dir, "test.parquet")) + + if hdfs_dir is not None: + makedirs(hdfs_dir) + copy(src=local_save_dir, dst=hdfs_dir) diff --git a/verl/examples/data_preprocess/preprocess_search_r1_dataset.py b/verl/examples/data_preprocess/preprocess_search_r1_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..a0c10d59b9c006ae7234ce21f7bdb25562259b23 --- /dev/null +++ b/verl/examples/data_preprocess/preprocess_search_r1_dataset.py @@ -0,0 +1,178 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2024 SGLang Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import logging +import os +import tempfile + +import pandas as pd +from huggingface_hub import hf_hub_download +from huggingface_hub.utils import EntryNotFoundError + +from verl.utils.hdfs_io import copy, makedirs + +# Setup logging +logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") +logger = logging.getLogger(__name__) + +# Configuration constants +DEFAULT_SYSTEM_CONTENT = "You are a helpful and harmless assistant." +DEFAULT_USER_CONTENT_PREFIX = ( + "Answer the given question. You must conduct reasoning inside and " + "first every time you get new information. After reasoning, if you find you lack " + "some knowledge, you can call a search engine by query " + "and it will return the top searched results between and " + ". You can search as many times as your want. If you find no " + "further external knowledge needed, you can directly provide the answer inside " + " and , without detailed illustrations. For example, " + " Beijing . Question: " +) + + +def process_single_row(row, current_split_name, row_index): + """ + Process a single row of data for SearchR1-like format. + + Args: + row: DataFrame row containing the original data + current_split_name: Name of the current split (train/test) + row_index: Index of the row in the DataFrame + + Returns: + pd.Series: Processed row data in the required format + """ + question = row.get("question", "") + + # Build prompt structure + user_content = user_content_prefix.rstrip("\n") + question + prompt = [{"role": "system", "content": system_content}, {"role": "user", "content": user_content}] + + # Extract ground truth from reward_model or fallback to golden_answers + reward_model_data = row.get("reward_model") + if isinstance(reward_model_data, dict) and "ground_truth" in reward_model_data: + ground_truth = reward_model_data.get("ground_truth") + else: + ground_truth = row.get("golden_answers", []) + + # Process data source + data_source_tagged = "searchR1_" + str(row.get("data_source", "")) + + # Build tools kwargs structure + tools_kwargs = { + "search": { + "create_kwargs": {"ground_truth": ground_truth, "question": question, "data_source": data_source_tagged} + } + } + + # Build complete extra_info structure + extra_info = { + "index": row_index, + "need_tools_kwargs": True, + "question": question, + "split": current_split_name, + "tools_kwargs": tools_kwargs, + } + + return pd.Series( + { + "data_source": data_source_tagged, + "prompt": prompt, + "ability": row.get("ability"), + "reward_model": reward_model_data, + "extra_info": extra_info, + "metadata": row.get("metadata"), + } + ) + + +def main(): + local_save_dir = os.path.expanduser(args.local_dir) + os.makedirs(local_save_dir, exist_ok=True) + + processed_files = [] + + # Download and process files using temporary directory + with tempfile.TemporaryDirectory() as tmp_download_dir: + for split in ["train", "test"]: + parquet_filename = f"{split}.parquet" + logger.info(f"Processing {split} split...") + + try: + # Download Parquet file from HuggingFace + logger.info(f"Downloading {parquet_filename} from {args.hf_repo_id}") + local_parquet_filepath = hf_hub_download( + repo_id=args.hf_repo_id, + filename=parquet_filename, + repo_type="dataset", + local_dir=tmp_download_dir, + local_dir_use_symlinks=False, + ) + + # Load and process Parquet file + df_raw = pd.read_parquet(local_parquet_filepath) + logger.info(f"Loaded {len(df_raw)} rows from {parquet_filename}") + + def apply_process_row(row, split_name=split): + return process_single_row(row, current_split_name=split_name, row_index=row.name) + + df_processed = df_raw.apply(apply_process_row, axis=1) + + # Save processed DataFrame + output_file_path = os.path.join(local_save_dir, f"{split}.parquet") + df_processed.to_parquet(output_file_path, index=False) + logger.info(f"Saved {len(df_processed)} processed rows to {output_file_path}") + processed_files.append(output_file_path) + + except EntryNotFoundError: + logger.warning(f"{parquet_filename} not found in repository {args.hf_repo_id}") + except Exception as e: + logger.error(f"Error processing {split} split: {e}") + + if not processed_files: + logger.warning("No data was processed or saved") + return + + logger.info(f"Successfully processed {len(processed_files)} files to {local_save_dir}") + + # Copy to HDFS if specified + if args.hdfs_dir: + try: + makedirs(args.hdfs_dir) + copy(src=local_save_dir, dst=args.hdfs_dir) + logger.info(f"Successfully copied files to HDFS: {args.hdfs_dir}") + except Exception as e: + logger.error(f"Error copying files to HDFS: {e}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Download Search-R1 from HuggingFace, process, and save to Parquet.") + parser.add_argument( + "--hf_repo_id", default="PeterJinGo/nq_hotpotqa_train", help="HuggingFace dataset repository ID." + ) + parser.add_argument( + "--local_dir", + default="~/data/searchR1_processed_direct", + help="Local directory to save the processed Parquet files.", + ) + parser.add_argument("--hdfs_dir", default=None, help="Optional HDFS directory to copy the Parquet files to.") + + args = parser.parse_args() + + # System and user content configuration + system_content = DEFAULT_SYSTEM_CONTENT + user_content_prefix = DEFAULT_USER_CONTENT_PREFIX + + main() diff --git a/verl/examples/dppo_trainer/README.md b/verl/examples/dppo_trainer/README.md new file mode 100644 index 0000000000000000000000000000000000000000..4de94dc563cc4086a0c05863e376a10c8af38f30 --- /dev/null +++ b/verl/examples/dppo_trainer/README.md @@ -0,0 +1,94 @@ +# Divergence Proximal Policy Optimization (DPPO) + + +
+ +## Rethinking the Trust Region in LLM Reinforcement Learning + +[![Paper](https://img.shields.io/badge/paper-A42C25?style=for-the-badge&logo=arxiv&logoColor=white )](https://arxiv.org/pdf/2602.04879) +[![Github](https://img.shields.io/badge/Stable_RL-000000?style=for-the-badge&logo=github&logoColor=000&logoColor=white)](https://github.com/sail-sg/Stable-RL) +[![Twitter](https://img.shields.io/badge/Twitter-%23000000.svg?style=for-the-badge&logo=twitter&logoColor=white)](https://x.com/QPHutu/status/2019435642539897303) + +
+ + +## ✨Getting started + +1. Prepare the datasets by running [prepare_dapo_data.sh](https://github.com/verl-project/verl-recipe/blob/3490a22a0a3adeb7e4787fe70b1060b642efbae4/dapo/prepare_dapo_data.sh): + +```bash +bash prepare_dapo_data.sh # This downloads the datasets to ${HOME}/verl/data by default +``` + +2. Prepare the model: + +```bash +hf download Qwen/Qwen3-30B-A3B-Base --local-dir ${HOME}/verl/models/Qwen3-30B-A3B-Base +``` + +3. Run the script: + +```bash +# run DPPO-Binary-KL +LOSS_MODE=dppo_kl bash examples/dppo_trainer/run_qwen3_30b_a3b_megatron.sh + +# run DPPO-Binary-TV +LOSS_MODE=dppo_tv bash examples/dppo_trainer/run_qwen3_30b_a3b_megatron.sh + +# run GRPO baseline +LOSS_MODE=vanilla CLIP_LOW=0.2 CLIP_HIGH=0.2 bash examples/dppo_trainer/run_qwen3_30b_a3b_megatron.sh +# or GRPO with clip higher +LOSS_MODE=vanilla CLIP_LOW=0.2 CLIP_HIGH=0.28 bash examples/dppo_trainer/run_qwen3_30b_a3b_megatron.sh +``` + +## 📖Introduction + +
+ issue +
+ +Comparison of **PPO** and the proposed **DPPO** (the Binary-TV variant). **(Left)** The surrogate objective and corresponding masks for PPO and DPPO. PPO (and variants like GRPO) employs a heuristic mask based on the probability ratio. In contrast, DPPO utilizes a more principled mask based on a direct approximation of policy divergence (e.g., Total Variation), ensuring updates stay within a theoretically grounded trust region. **(Right)** Experimental results on the AIME24 using Qwen3-30B-A3B-Base. DPPO significantly outperforms GRPO baselines, achieving superior training stability and final performance even without rollout routing replay (R3). + +
+ issue +
+ +DPPO variants achieve stable training while controlling the training-inference mismatch at a low level. In contrast, methods without a trust region (PG-IS, CISPO) or with a misspecified one (MiniRL) suffer from growing mismatch and eventual collapse. + +
+ issue +
+ +The plots show numerical differences between a training and an inference engine for Qwen3-30B-A3B-Base with identical parameters. **(Left)** The probability ratio (used in PPO) is highly volatile for low-probability tokens. **(Right)** In contrast, the TV divergence is more stable. This highlights a key flaw of PPO's clipping mechanism: it **over-penalizes low-probability tokens**, which can slow down learning; and **under-penalizes high-probability tokens**, which can permit large, destabilizing updates. + + +
+ issue +
+ +The most frequently clipped tokens (by GRPO) are important to the reasoning task! +They are dominated by: +- numbers, like 1, 4 +- mathematical symbols, like +, -, = +- reasoning and structural Words: Wait, Thus, Next + +## Top-K divergence approximation + +We only implement the DPPO-Binary-TV/DPPO-Binary-KL here due to their simplicity. + +For the TopK divergence approximation, please refer to the [the original repo](https://github.com/sail-sg/Stable-RL) for a complete implementation. + +## Citation +If you find our works useful for your research, please consider citing: + +```bibtex +@article{qi2026dppo, + title={Rethinking the Trust Region in LLM Reinforcement Learning}, + author={Qi, Penghui and Zhou, Xiangxin and Liu, Zichen and Pang, Tianyu and Du, Chao and Lin, Min and Lee, Wee Sun}, + journal={arXiv preprint arXiv:2602.04879}, + year={2026} +} +``` + +## 🌻Acknowledgement +We implement our reinforcement learning algorithm extending from [verl](https://github.com/verl-project/verl). We utilize [vLLM](https://github.com/vllm-project/vllm) and [sglang](https://github.com/sgl-project/sglang) for inference. Our models are trained primarily on [Qwen3 family](https://huggingface.co/collections/Qwen/qwen3). Our training data is built from [DAPO-MATH](https://huggingface.co/datasets/BytedTsinghua-SIA/DAPO-Math-17k). Thanks for their great contributions! diff --git a/verl/examples/dppo_trainer/run_qwen3_30b_a3b_megatron.sh b/verl/examples/dppo_trainer/run_qwen3_30b_a3b_megatron.sh new file mode 100644 index 0000000000000000000000000000000000000000..4e21a2cd5c455b912cda9052e049767fa5c83bad --- /dev/null +++ b/verl/examples/dppo_trainer/run_qwen3_30b_a3b_megatron.sh @@ -0,0 +1,148 @@ +#!/usr/bin/env bash +# DPPO | MoE | vLLM rollout | Megatron training | NVIDIA GPUs +# DPPO replaces PPO's ratio clip with a TV/KL-divergence clip (paper: arXiv:2602.04879). + +set -xeuo pipefail + +export CUDA_DEVICE_MAX_CONNECTIONS=1 +export VLLM_USE_V1=1 + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-30B-A3B-Base} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +# LOSS_MODE selects DPPO variant: dppo_tv | dppo_kl (or vanilla for GRPO baseline) +LOSS_MODE=${LOSS_MODE:-dppo_tv} +case $LOSS_MODE in + dppo_tv) CLIP_DEFAULT=0.15 ;; + dppo_kl) CLIP_DEFAULT=0.05 ;; + vanilla) CLIP_DEFAULT=0.20 ;; + *) echo "Unknown LOSS_MODE: $LOSS_MODE"; exit 1 ;; +esac +clip_ratio_low=${CLIP_LOW:-$CLIP_DEFAULT} +clip_ratio_high=${CLIP_HIGH:-$CLIP_DEFAULT} + +train_batch_size=${TRAIN_BATCH_SIZE:-256} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-32} +max_prompt_length=${MAX_PROMPT_LENGTH:-2048} +max_response_length=${MAX_RESPONSE_LENGTH:-8192} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-30720} + +actor_lr=${ACTOR_LR:-1e-6} +entropy_coeff=${ENTROPY_COEFF:-0} + +actor_tp=${ACTOR_TP:-2} +actor_pp=${ACTOR_PP:-1} +actor_ep=${ACTOR_EP:-8} +actor_etp=${ACTOR_ETP:-1} + +rollout_tp=${ROLLOUT_TP:-4} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.8} +rollout_n=${ROLLOUT_N:-16} + +total_epochs=${TOTAL_EPOCHS:-10} +save_freq=${SAVE_FREQ:-50} +test_freq=${TEST_FREQ:-10} + +project_name=${PROJECT_NAME:-verl_dppo_qwen3_moe} +experiment_name=${EXPERIMENT_NAME:-qwen3_30b_a3b_${LOSS_MODE}_vllm_megatron} +# ---- end user-adjustable ---- + +train_file=${TRAIN_FILE:-$HOME/data/dapo-math-17k/train.parquet} +val_file=${VAL_FILE:-$HOME/data/aime-2024/test.parquet} +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + algorithm.norm_adv_by_std_in_grpo=False + data.train_files="['$train_file']" + data.val_files="['$val_file']" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True +) + +ACTOR=( + actor_rollout_ref.actor.policy_loss.loss_mode=${LOSS_MODE} + actor_rollout_ref.actor.loss_agg_mode=seq-mean-token-sum-norm + actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} + actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} + actor_rollout_ref.actor.clip_ratio_c=10000.0 + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=False + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${actor_pp} + actor_rollout_ref.actor.megatron.expert_model_parallel_size=${actor_ep} + actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${actor_etp} + actor_rollout_ref.actor.megatron.param_offload=True + actor_rollout_ref.actor.megatron.grad_offload=True + actor_rollout_ref.actor.megatron.optimizer_offload=True + actor_rollout_ref.actor.megatron.use_mbridge=True + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32 + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=True + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${actor_pp} + actor_rollout_ref.ref.megatron.expert_model_parallel_size=${actor_ep} + actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=${actor_etp} + actor_rollout_ref.ref.megatron.param_offload=True + actor_rollout_ref.ref.megatron.use_mbridge=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.val_before_train=False + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( + model_engine=megatron +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/gdpo_trainer/README.md b/verl/examples/gdpo_trainer/README.md new file mode 100644 index 0000000000000000000000000000000000000000..88d38454cd8ce242e6f7df35427d21f73dd81ad2 --- /dev/null +++ b/verl/examples/gdpo_trainer/README.md @@ -0,0 +1,18 @@ +# GDPO + +GDPO is a multi-reward, rubric-style variant whose advantage estimator aggregates several reward signals (accuracy, format, etc.). It uses a custom reward manager and a custom scoring function. + +## Canonical Scripts + +| Script | Infer | Train | Platform | +|--------------------------------------|-------|-------|----------| +| `run_qwen3_8b_fsdp.sh` | vLLM | FSDP | NVIDIA | + +Prepare a rubric-style dataset (e.g. `rlla_4k`) and point `DATA_DIR` to it. + +## Key Flags + +- `algorithm.adv_estimator=gdpo` +- `+algorithm.gdpo_reward_keys='["accuracy_reward", "format_reward"]'` +- `reward.reward_manager.name=gdpo` +- `reward.custom_reward_function.path=$REPO_ROOT/verl/utils/reward_score/rlla.py` diff --git a/verl/examples/gdpo_trainer/run_qwen3_8b_fsdp.sh b/verl/examples/gdpo_trainer/run_qwen3_8b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..f0b314a3a9dd096f849effe3146e0b52edbce23b --- /dev/null +++ b/verl/examples/gdpo_trainer/run_qwen3_8b_fsdp.sh @@ -0,0 +1,117 @@ +#!/usr/bin/env bash +# GDPO | text | vLLM rollout | FSDP training | NVIDIA GPUs +# GDPO uses its own advantage estimator and a custom reward manager/reward function. + +set -xeuo pipefail + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-8B} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +DATA_DIR=${DATA_DIR:-$HOME/data/rlla_4k} + +train_batch_size=${TRAIN_BATCH_SIZE:-128} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-32} +max_prompt_length=${MAX_PROMPT_LENGTH:-2048} +max_response_length=${MAX_RESPONSE_LENGTH:-2048} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-24576} + +actor_lr=${ACTOR_LR:-1e-6} +kl_loss_coef=${KL_LOSS_COEF:-0.001} + +rollout_tp=${ROLLOUT_TP:-2} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.6} +rollout_n=${ROLLOUT_N:-4} + +total_epochs=${TOTAL_EPOCHS:-15} +save_freq=${SAVE_FREQ:-20} +test_freq=${TEST_FREQ:-10} + +project_name=${PROJECT_NAME:-verl_gdpo} +experiment_name=${EXPERIMENT_NAME:-qwen3_8b_vllm_fsdp} +# ---- end user-adjustable ---- + +REPO_ROOT=$(cd -- "$(dirname "${BASH_SOURCE[0]}")/../.." >/dev/null 2>&1 && pwd) +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=gdpo + +algorithm.gdpo_reward_keys='["accuracy_reward", "format_reward"]' + algorithm.kl_ctrl.kl_coef=${kl_loss_coef} + data.train_files=$DATA_DIR/train.parquet + data.val_files=$DATA_DIR/test.parquet + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=False + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.fsdp_config.param_offload=False + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.fsdp_config.param_offload=True +) + +REWARD=( + reward.custom_reward_function.path="$REPO_ROOT/verl/utils/reward_score/rlla.py" + reward.custom_reward_function.name=compute_score + reward.reward_manager.name=gdpo +) + +TRAINER=( + trainer.balance_batch=True + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.val_before_train=False + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${REWARD[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/generation/run_deepseek_llm_7b.sh b/verl/examples/generation/run_deepseek_llm_7b.sh new file mode 100644 index 0000000000000000000000000000000000000000..a28c6c1559fdc302bb7c5e4be3ae4a30b0d58b22 --- /dev/null +++ b/verl/examples/generation/run_deepseek_llm_7b.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +# main_generation_server | text | vLLM rollout | NVIDIA GPUs +# Rollout-only generation with DeepSeek-LLM-7B-Chat. +# +# Single-node (default): +# bash run_deepseek_llm_7b.sh +# +# Multi-node (e.g. 2 nodes, rollout_tp=16): +# NNODES=2 ROLLOUT_TP=16 bash run_deepseek_llm_7b.sh + +set -xeuo pipefail + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-deepseek-ai/deepseek-llm-7b-chat} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +DATA_PATH=${DATA_PATH:-$HOME/data/gsm8k/test.parquet} +OUTPUT_PATH=${OUTPUT_PATH:-$HOME/data/gsm8k/deepseek_llm_7b_gen_test.parquet} + +PROMPT_LENGTH=${PROMPT_LENGTH:-2048} +RESPONSE_LENGTH=${RESPONSE_LENGTH:-1024} +ROLLOUT_TP=${ROLLOUT_TP:-2} +ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.8} +N_SAMPLES=${N_SAMPLES:-1} +# ---- end user-adjustable ---- + +python3 -m verl.trainer.main_generation_server \ + trainer.nnodes="${NNODES}" \ + trainer.n_gpus_per_node="${NGPUS_PER_NODE}" \ + data.train_files="${DATA_PATH}" \ + data.prompt_key=prompt \ + +data.output_path="${OUTPUT_PATH}" \ + actor_rollout_ref.model.path="${MODEL_PATH}" \ + actor_rollout_ref.model.trust_remote_code=True \ + actor_rollout_ref.rollout.name=vllm \ + actor_rollout_ref.rollout.temperature=1.0 \ + actor_rollout_ref.rollout.top_k=50 \ + actor_rollout_ref.rollout.top_p=0.7 \ + actor_rollout_ref.rollout.prompt_length="${PROMPT_LENGTH}" \ + actor_rollout_ref.rollout.response_length="${RESPONSE_LENGTH}" \ + actor_rollout_ref.rollout.tensor_model_parallel_size="${ROLLOUT_TP}" \ + actor_rollout_ref.rollout.gpu_memory_utilization="${ROLLOUT_GPU_MEM_UTIL}" \ + actor_rollout_ref.rollout.n="${N_SAMPLES}" "$@" diff --git a/verl/examples/gmpo_trainer/README.md b/verl/examples/gmpo_trainer/README.md new file mode 100644 index 0000000000000000000000000000000000000000..c81764554a990792e191c978999aa333b91b9b1e --- /dev/null +++ b/verl/examples/gmpo_trainer/README.md @@ -0,0 +1,56 @@ +
+ +# Geometric-Mean Policy Optimization +
+ +This is the official implementaion of paper [***Geometric-Mean Policy Optimization***](https://arxiv.org/abs/2507.20673). + +
+image +
+ +## 1. Contents +- Geometric-Mean Policy Optimization + - [1. Contents](#1-contents) + - [2. Introduction](#2-introduction) + - [3. Code Usage](#3-code-usage) + - [4. Contacts](#4-contacts) + - [5. Citation](#5-citation) + +## 2. Introduction + +Group Relative Policy Optimization (GRPO) has significantly enhanced the reasoning capability of large language models by optimizing the arithmetic mean of token-level rewards. Unfortunately, GRPO is observed to suffer from unstable policy updates when facing tokens with outlier importance-weighted rewards, which manifest as extreme importance sampling ratios during training. In this study, we propose Geometric-Mean Policy Optimization (GMPO), with the aim to improve the stability of GRPO through suppressing token reward outliers. Instead of optimizing the arithmetic mean, GMPO maximizes the geometric mean of token-level rewards, which is inherently less sensitive to outliers and maintains a more stable range of importance sampling ratio. GMPO is plug-and-play—simply replacing GRPO's arithmetic mean with the geometric mean of token-level rewards, as the latter is inherently less sensitive to outliers. GMPO is theoretically plausible—analysis reveals that both GMPO and GRPO are weighted forms of the policy gradient while the former enjoys more stable weights, which consequently benefits policy optimization and performance. Experiments on multiple mathematical reasoning benchmarks show that GMPO-7B improves the average Pass@1 of GRPO by up to 4.1%, outperforming many state-of-the-art approaches. + +## 3. Code Usage + +The key configurations are: +``` +clip_ratio_low=0.4 +clip_ratio_high=0.4 +loss_mode=geo_mean +``` +We observed that using a large clip ratio during Mixture-of-Experts (MoE) model training often leads to optimization instability. When training MoE models, consider lowering the clip ratio to achieve more stable convergence. + +To get started quickly: + +```bash +bash examples/gmpo_trainer/run_qwen3_8b_fsdp.sh +# override any of MODEL_PATH, CLIP_RATIO, ROLLOUT_N, etc. via env vars +``` + +## 4. Contacts +If you have any question about our work or this repository, please don't hesitate to contact us by emails or open an issue under this project. +- [zhaoyuzhong20@mails.ucas.ac.cn](zhaoyuzhong20@mails.ucas.ac.cn) +- [liuyue171@mails.ucas.ac.cn](liuyue171@mails.ucas.ac.cn) +- [lecu@microsoft.com](lecu@microsoft.com) +- [wanfang@ucas.ac.cn](wanfang@ucas.ac.cn) + +## 5. Citation +``` +@article{zhao2025geometric, + title={Geometric-mean policy optimization}, + author={Zhao, Yuzhong and Liu, Yue and Liu, Junpeng and Chen, Jingye and Wu, Xun and Hao, Yaru and Lv, Tengchao and Huang, Shaohan and Cui, Lei and Ye, Qixiang and others}, + journal={arXiv preprint arXiv:2507.20673}, + year={2025} +} +``` diff --git a/verl/examples/gmpo_trainer/run_qwen3_8b_fsdp.sh b/verl/examples/gmpo_trainer/run_qwen3_8b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..ca362f1e69bff239f7506bfc3ddbd0e5a904f5a0 --- /dev/null +++ b/verl/examples/gmpo_trainer/run_qwen3_8b_fsdp.sh @@ -0,0 +1,119 @@ +#!/usr/bin/env bash +# GMPO | text | vLLM rollout | FSDP training | NVIDIA GPUs +# GMPO uses a geometric-mean aggregated ratio in the policy loss. + +set -xeuo pipefail + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-8B} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +train_batch_size=${TRAIN_BATCH_SIZE:-1024} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-256} +max_prompt_length=${MAX_PROMPT_LENGTH:-1024} +max_response_length=${MAX_RESPONSE_LENGTH:-2048} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-24576} + +actor_lr=${ACTOR_LR:-1e-6} +kl_loss_coef=${KL_LOSS_COEF:-0.001} +entropy_coeff=${ENTROPY_COEFF:-0} + +clip_ratio=${CLIP_RATIO:-0.4} + +rollout_tp=${ROLLOUT_TP:-2} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.6} +rollout_n=${ROLLOUT_N:-5} + +total_epochs=${TOTAL_EPOCHS:-15} +save_freq=${SAVE_FREQ:-20} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_gmpo_gsm8k_math} +experiment_name=${EXPERIMENT_NAME:-qwen3_8b_vllm_fsdp} +# ---- end user-adjustable ---- + +gsm8k_train=$HOME/data/gsm8k/train.parquet +gsm8k_test=$HOME/data/gsm8k/test.parquet +math_train=$HOME/data/math/train.parquet +math_test=$HOME/data/math/test.parquet + +train_files="['$gsm8k_train', '$math_train']" +val_files="['$gsm8k_test', '$math_test']" +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files="$train_files" + data.val_files="$val_files" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.policy_loss.loss_mode=geo_mean + actor_rollout_ref.actor.clip_ratio_low=${clip_ratio} + actor_rollout_ref.actor.clip_ratio_high=${clip_ratio} + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.fsdp_config.param_offload=False + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.fsdp_config.param_offload=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/gpg_trainer/README.md b/verl/examples/gpg_trainer/README.md new file mode 100644 index 0000000000000000000000000000000000000000..b40cc83bcd7aeaaef43622df7659fc03b394138d --- /dev/null +++ b/verl/examples/gpg_trainer/README.md @@ -0,0 +1,34 @@ +# GPG: Group Policy Gradient + +Group Policy Gradient (GPG) is a minimalist reinforcement learning (RL) method that enhances the reasoning ability of large language models without relying on supervised fine-tuning or complex tricks. GPG revisits traditional policy gradients and directly optimizes the RL objective—no surrogate losses, no KL penalties, no critic, and no reference model. Compared to GRPO, GPG is simpler, more efficient, and achieves better results on many tasks. For more details, please refer to the original paper [GPG: A Simple and Strong Reinforcement Learning Baseline for Model Reasoning +](https://arxiv.org/abs/2504.02546). + +## Key Components +- Use a corrected advantage function to improve policy gradient accuracy and training efficiency. +- By eliminating the critic and reference models, avoiding KL divergence constraints, significantly simplifies the training process compared to Group Relative Policy Optimization (GRPO) + +## Configuration +To configure GPG within the framework, use the following YAML settings. + +```yaml +algorithm: + adv_estimator: gpg +actor_rollout_ref: + actor: + policy_loss: + loss_mode: "gpg" +``` + +## Advanced Extensions +GPG is a simple and strong baseline for model reasoning. Although it avoids using KL loss in its original form, you can still use KL loss to further improve the performance. + +```yaml +algorithm: + adv_estimator: gpg +actor_rollout_ref: + actor: + use_kl_loss: True # enable kl regularization + kl_loss_coef: 0.01 + policy_loss: + loss_mode: "gpg" +``` \ No newline at end of file diff --git a/verl/examples/gpg_trainer/run_qwen3_8b_fsdp.sh b/verl/examples/gpg_trainer/run_qwen3_8b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..d9e018614202c3a7d8a2f04eb6e6fc513fc14c5b --- /dev/null +++ b/verl/examples/gpg_trainer/run_qwen3_8b_fsdp.sh @@ -0,0 +1,111 @@ +#!/usr/bin/env bash +# GPG | text | vLLM rollout | FSDP training | NVIDIA GPUs + +set -xeuo pipefail + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-8B} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +train_batch_size=${TRAIN_BATCH_SIZE:-1024} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-256} +max_prompt_length=${MAX_PROMPT_LENGTH:-1024} +max_response_length=${MAX_RESPONSE_LENGTH:-2048} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-24576} + +actor_lr=${ACTOR_LR:-1e-6} +entropy_coeff=${ENTROPY_COEFF:-0} + +rollout_tp=${ROLLOUT_TP:-2} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.6} +rollout_n=${ROLLOUT_N:-5} + +total_epochs=${TOTAL_EPOCHS:-15} +save_freq=${SAVE_FREQ:-20} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_gpg_gsm8k_math} +experiment_name=${EXPERIMENT_NAME:-qwen3_8b_vllm_fsdp} +# ---- end user-adjustable ---- + +gsm8k_train=$HOME/data/gsm8k/train.parquet +gsm8k_test=$HOME/data/gsm8k/test.parquet +math_train=$HOME/data/math/train.parquet +math_test=$HOME/data/math/test.parquet + +train_files="['$gsm8k_train', '$math_train']" +val_files="['$gsm8k_test', '$math_test']" +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=gpg + algorithm.use_kl_in_reward=False + data.train_files="$train_files" + data.val_files="$val_files" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.policy_loss.loss_mode=gpg + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=False + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.fsdp_config.param_offload=False + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.fsdp_config.param_offload=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/gpg_trainer/run_qwen3_8b_megatron.sh b/verl/examples/gpg_trainer/run_qwen3_8b_megatron.sh new file mode 100644 index 0000000000000000000000000000000000000000..b3b7cf09f6b7877ebd1d714b910bd843872322e6 --- /dev/null +++ b/verl/examples/gpg_trainer/run_qwen3_8b_megatron.sh @@ -0,0 +1,117 @@ +#!/usr/bin/env bash +# GPG | text | vLLM rollout | Megatron training | NVIDIA GPUs + +set -xeuo pipefail + +export CUDA_DEVICE_MAX_CONNECTIONS=1 + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-8B} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +train_batch_size=${TRAIN_BATCH_SIZE:-1024} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-256} +max_prompt_length=${MAX_PROMPT_LENGTH:-1024} +max_response_length=${MAX_RESPONSE_LENGTH:-2048} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-24576} + +actor_lr=${ACTOR_LR:-1e-6} +entropy_coeff=${ENTROPY_COEFF:-0} + +actor_tp=${ACTOR_TP:-2} +actor_pp=${ACTOR_PP:-2} + +rollout_tp=${ROLLOUT_TP:-2} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.6} +rollout_n=${ROLLOUT_N:-5} + +total_epochs=${TOTAL_EPOCHS:-15} +save_freq=${SAVE_FREQ:-20} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_gpg_gsm8k_math} +experiment_name=${EXPERIMENT_NAME:-qwen3_8b_vllm_megatron} +# ---- end user-adjustable ---- + +gsm8k_train=$HOME/data/gsm8k/train.parquet +gsm8k_test=$HOME/data/gsm8k/test.parquet +math_train=$HOME/data/math/train.parquet +math_test=$HOME/data/math/test.parquet + +train_files="['$gsm8k_train', '$math_train']" +val_files="['$gsm8k_test', '$math_test']" +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=gpg + algorithm.use_kl_in_reward=False + data.train_files="$train_files" + data.val_files="$val_files" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True +) + +ACTOR=( + actor_rollout_ref.actor.policy_loss.loss_mode=gpg + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=False + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${actor_pp} +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${actor_pp} +) + +TRAINER=( + trainer.balance_batch=True + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( + model_engine=megatron +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/grpo_trainer/README.md b/verl/examples/grpo_trainer/README.md new file mode 100644 index 0000000000000000000000000000000000000000..7ad8f4c2436013516a2ea7b883c809e132e7a4f8 --- /dev/null +++ b/verl/examples/grpo_trainer/README.md @@ -0,0 +1,105 @@ +# Group Relative Policy Optimization (GRPO) + +In reinforcement learning, classic algorithms like PPO rely on a "critic" model to estimate the value of actions, guiding the learning process. However, training this critic model can be resource-intensive. + +GRPO simplifies this process by eliminating the need for a separate critic model. Instead, it operates as follows: +- Group Sampling: for a given problem, the model generates multiple possible solutions, forming a "group" of outputs. +- Reward Assignment: each solution is evaluated and assigned a reward based on its correctness or quality. +- Baseline Calculation: the average reward of the group serves as a baseline. +- Policy Update: the model updates its parameters by comparing each solution's reward to the group baseline, reinforcing better-than-average solutions and discouraging worse-than-average ones. + +For more details, refer to the original paper [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://arxiv.org/pdf/2402.03300). + +## Key Components + +- No Value Function (Critic-less): unlike PPO, GRPO does not train a separate value network (critic). +- Group Sampling (Grouped Rollouts): instead of evaluating one rollout per input, GRPO generates multiple completions (responses) from the current policy for each prompt. This set of completions is referred to as a group. +- Relative Rewards: within each group, completions are scored (e.g., based on correctness), and rewards are normalized relative to the group. + +## Important knobs + +- `actor_rollout_ref.rollout.n`: per-prompt sample count (required >= 2 for GRPO). +- `data.train_batch_size`: prompts per global step. Total trajectories = `train_batch_size * rollout.n`. +- `actor_rollout_ref.actor.ppo_mini_batch_size`: global mini-batch for actor updates (must divide `train_batch_size * n`). +- `actor_rollout_ref.actor.ppo_epochs`: inner-loop epochs over the sampled trajectories. +- `actor_rollout_ref.actor.clip_ratio`: PPO clip range, default `0.2`. +- `actor_rollout_ref.actor.loss_agg_mode`: `token-mean` (default), `seq-mean-token-sum`, or `seq-mean-token-mean`. +- `actor_rollout_ref.actor.use_kl_loss=True` + `actor_rollout_ref.actor.kl_loss_coef` / `kl_loss_type`: regularise toward the reference policy via KL loss on the actor. +- `algorithm.adv_estimator=grpo`. + +## Dr. GRPO + +To enable Dr. GRPO (see [Understanding R1-Zero-Like Training](https://arxiv.org/pdf/2503.20783)), set on top of the canonical GRPO overrides: + +``` +actor_rollout_ref.actor.loss_agg_mode=seq-mean-token-sum-norm +actor_rollout_ref.actor.use_kl_loss=False +algorithm.norm_adv_by_std_in_grpo=False +``` + +## Canonical scripts + +All scripts in this directory follow the naming convention: + +``` +run__[_].sh +``` + +Where: +- `` is the canonical size for a model family + (`qwen3_8b` for dense text, `qwen3_30b_a3b` for MoE, `qwen2_5_vl_7b` / `qwen3_vl_8b` for vision, + `qwen3_235b_a22b` / `deepseek_v3_671b` for scale demos). +- `` ∈ {`fsdp`, `megatron`, `mindspeed`}. +- `` is used only for hardware-specific variants such as `gb200`, `fp8`, `veomni`, + or MindSpeed NPU scripts. +- `INFER_BACKEND` selects rollout backend inside scripts that support multiple choices + (`vllm`, `sglang`, or `trtllm`). +- `DEVICE` selects GPU/NPU paths inside scripts that support both platforms. + +Every script exposes the commonly tuned knobs as environment variables at the top, so you can run: + +```bash +MODEL_PATH=Qwen/Qwen3-14B \ +NNODES=2 NGPUS_PER_NODE=8 \ +INFER_BACKEND=sglang ROLLOUT_N=8 TRAIN_BATCH_SIZE=2048 \ +bash examples/grpo_trainer/run_qwen3_8b_fsdp.sh +``` + +### Defaults + +- `dynamic batch size` and `sequence balancing` are enabled by default on all scripts. +- Text LLM scripts train on `gsm8k + math` by default; vision scripts train on `geo3k`. +- Scale-demo scripts (235B, 671B) train on `dapo-math-17k` / `aime-2024`. + +### Matrix + +| Model family | `vllm` | `sglang` | `trtllm` | Train backend | Platforms | +| --------------------- | :----: | :------: | :------: | --------------- | --------- | +| Qwen3-8B (dense) | ✓ | ✓ | ✓ | FSDP, Megatron | nvidia, npu (FSDP + MindSpeed), `_gb200` variant | +| Qwen2.5-VL-7B | ✓ | ✓ | ✓ | FSDP, Megatron | nvidia | +| Qwen3-VL-8B | ✓ | | | FSDP, Megatron | nvidia, npu (FSDP) | +| Qwen3-VL-30B-A3B | ✓ | | | FSDP, Megatron | nvidia, npu (FSDP, VeOmni) | +| Qwen3-VL-235B-A22B | ✓ | | | Megatron | nvidia | +| Qwen3-30B-A3B (MoE) | ✓ | ✓ | ✓ | FSDP, Megatron | nvidia, npu (MindSpeed, VeOmni) | +| Qwen3-235B-A22B | ✓ | | ✓ | Megatron | nvidia, npu | +| Qwen3-Next-80B-A3B | ✓ | | | FSDP | npu | +| Qwen3.5-27B (dense) | ✓ | | | FSDP2 | nvidia, npu | +| Qwen3.5-35B (dense) | ✓ | | | FSDP2, Megatron | nvidia, npu | +| Qwen3.5-35B-A3B (MoE) | | ✓ | | VeOmni | nvidia | +| Qwen3.5-122B-A10B | ✓ | | | Megatron | nvidia | +| DeepSeek-V3 671B | ✓ | | | Megatron | nvidia | +| GLM-4.1V-9B | ✓ | | | FSDP | nvidia | +| MiniCPM-o-2.6 | ✓ | | | FSDP | nvidia | +| Moonlight-16B-A3B | ✓ | | | Megatron | nvidia | +| Nemotron-Nano-v3-30B-A3B | ✓ | | | Megatron | nvidia | +| Seed-OSS-36B | ✓ | | | FSDP2 | nvidia | +| GPT-OSS-20B | | ✓ | | FSDP | nvidia | +| Mistral-Nemo-12B (RM demo) | ✓ | | | FSDP | nvidia | + +LoRA variants live in `examples/tuning/lora/`, profiling variants in `examples/profile/`. +Scale / hardware-specific demos (e.g. `run_qwen3_8b_fsdp_gb200.sh`, FP8 variants, VeOmni) keep a trailing suffix to stay discoverable. + +## Reference + +- See [verl baselines](https://verl.readthedocs.io/en/latest/algo/baseline.html) for reference metrics. +- Qwen2.5 GRPO training log: [experiments/gsm8k/qwen2-7b-fsdp2.log](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/qwen2-7b-fsdp2.log). diff --git a/verl/examples/grpo_trainer/run_deepseek_v3_671b_megatron.sh b/verl/examples/grpo_trainer/run_deepseek_v3_671b_megatron.sh new file mode 100644 index 0000000000000000000000000000000000000000..c7075b6ac71a079d1497e6d2c20519ac5f000a1e --- /dev/null +++ b/verl/examples/grpo_trainer/run_deepseek_v3_671b_megatron.sh @@ -0,0 +1,198 @@ +#!/usr/bin/env bash +# GRPO scale demo | DeepSeek-V3 671B | vLLM rollout | Megatron training | NVIDIA GPUs +# +# Prerequisites on every node: +# CUDA_DEVICE_MAX_CONNECTIONS=1 +# NCCL_NVLS_ENABLE=0 +# VLLM_USE_V1=1 +# pip install git+https://github.com/ISEEKYAN/mbridge +# Also: remove `quantization_config` from DeepSeek-V3 config.json and set +# `num_nextn_predict_layers=0` (MTP not yet supported). +# Minimum 12 nodes x 8x 80GB+ GPUs recommended. + +set -xeuo pipefail + +export CUDA_DEVICE_MAX_CONNECTIONS=1 +export NCCL_NVLS_ENABLE=0 +export VLLM_USE_V1=1 + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-deepseek-ai/DeepSeek-V3} +NNODES=${NNODES:-12} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +train_batch_size=${TRAIN_BATCH_SIZE:-96} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-32} +max_prompt_length=${MAX_PROMPT_LENGTH:-2048} +max_response_length=${MAX_RESPONSE_LENGTH:-8192} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-$((max_prompt_length + max_response_length))} + +actor_lr=${ACTOR_LR:-1e-6} +kl_loss_coef=${KL_LOSS_COEF:-0.001} +entropy_coeff=${ENTROPY_COEFF:-0} +clip_ratio_low=${CLIP_RATIO_LOW:-0.2} +clip_ratio_high=${CLIP_RATIO_HIGH:-0.28} + +actor_tp=${ACTOR_TP:-8} +actor_pp=${ACTOR_PP:-12} +actor_ep=${ACTOR_EP:-8} +actor_etp=${ACTOR_ETP:-1} +actor_cp=${ACTOR_CP:-1} +last_layer=${LAST_LAYER:-6} +offload=${OFFLOAD:-True} +optim_offload=${OPTIM_OFFLOAD:-True} +optimizer_offload_fraction=${OPTIMIZER_OFFLOAD_FRACTION:-1.0} + +rollout_tp=${ROLLOUT_TP:-32} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.6} +rollout_n=${ROLLOUT_N:-8} + +overlong_buffer_len=${OVERLONG_BUFFER_LEN:-4096} +overlong_penalty_factor=${OVERLONG_PENALTY_FACTOR:-1.0} + +total_epochs=${TOTAL_EPOCHS:-10} +save_freq=${SAVE_FREQ:-100} +test_freq=${TEST_FREQ:-10} + +project_name=${PROJECT_NAME:-verl_grpo_scale_demo} +experiment_name=${EXPERIMENT_NAME:-deepseek_v3_671b_vllm_megatron} +CKPTS_DIR=${CKPTS_DIR:-"${HOME}/verl/ckpts/${project_name}/${experiment_name}"} +# ---- end user-adjustable ---- + +train_files=$HOME/data/dapo-math-17k.parquet +val_files=$HOME/data/aime-2024.parquet +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + algorithm.kl_ctrl.kl_coef=0.0 + data.train_files="$train_files" + data.val_files="$val_files" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.truncation=left +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_fused_kernels=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.optim.lr_warmup_steps=10 + actor_rollout_ref.actor.optim.weight_decay=0.1 + actor_rollout_ref.actor.optim.clip_grad=1.0 + +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=${optimizer_offload_fraction} + +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True + +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=True + +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} + actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} + actor_rollout_ref.actor.clip_ratio_c=10.0 + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.loss_agg_mode=token-mean + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${actor_pp} + actor_rollout_ref.actor.megatron.expert_model_parallel_size=${actor_ep} + actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${actor_etp} + actor_rollout_ref.actor.megatron.context_parallel_size=${actor_cp} + actor_rollout_ref.actor.megatron.param_offload=${offload} + actor_rollout_ref.actor.megatron.optimizer_offload=${optim_offload} + actor_rollout_ref.actor.megatron.grad_offload=${offload} + actor_rollout_ref.actor.megatron.use_mbridge=True + +actor_rollout_ref.actor.megatron.override_transformer_config.apply_rope_fusion=False + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32 + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_shared_expert_overlap=False + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_enable_deepep=True + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_token_dispatcher_type=flex + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=True + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 + +actor_rollout_ref.actor.megatron.override_transformer_config.gradient_accumulation_fusion=True + +actor_rollout_ref.actor.megatron.override_transformer_config.account_for_embedding_in_pipeline_split=False + +actor_rollout_ref.actor.megatron.override_transformer_config.account_for_loss_in_pipeline_split=False + +actor_rollout_ref.actor.megatron.override_transformer_config.num_layers_in_last_pipeline_stage=${last_layer} +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.rollout.enable_chunked_prefill=True + actor_rollout_ref.rollout.enforce_eager=True + actor_rollout_ref.rollout.temperature=1.0 + actor_rollout_ref.rollout.top_p=1.0 + actor_rollout_ref.rollout.top_k=-1 + actor_rollout_ref.rollout.val_kwargs.temperature=1.0 + actor_rollout_ref.rollout.val_kwargs.top_p=0.7 + actor_rollout_ref.rollout.val_kwargs.top_k=-1 + actor_rollout_ref.rollout.val_kwargs.do_sample=True + actor_rollout_ref.rollout.val_kwargs.n=1 +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${actor_pp} + actor_rollout_ref.ref.megatron.expert_model_parallel_size=${actor_ep} + actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=${actor_etp} + actor_rollout_ref.ref.megatron.context_parallel_size=${actor_cp} + actor_rollout_ref.ref.megatron.param_offload=${offload} +) + +REWARD=( + reward.reward_manager.name=dapo + +reward.reward_kwargs.overlong_buffer_cfg.enable=True + +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} + +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} + +reward.reward_kwargs.overlong_buffer_cfg.log=False + +reward.reward_kwargs.max_resp_len=${max_response_length} +) + +TRAINER=( + trainer.balance_batch=True + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.val_before_train=False + trainer.resume_mode=auto + trainer.log_val_generations=10 + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} + trainer.default_local_dir="${CKPTS_DIR}" +) + +EXTRA=( + model_engine=megatron + actor_rollout_ref.nccl_timeout=1200 +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${REWARD[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/grpo_trainer/run_glm4_1v_9b_fsdp.sh b/verl/examples/grpo_trainer/run_glm4_1v_9b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..4996afb74453526406caf2a0167aff1e5c9d1af0 --- /dev/null +++ b/verl/examples/grpo_trainer/run_glm4_1v_9b_fsdp.sh @@ -0,0 +1,105 @@ +set -x + +# ---- user-adjustable ---- +INFER_BACKEND=${INFER_BACKEND:-vllm} +MODEL_PATH=${MODEL_PATH:-zai-org/GLM-4.1V-9B-Thinking} +TRAIN_FILE=${TRAIN_FILE:-$HOME/data/geo3k/train.parquet} +TEST_FILE=${TEST_FILE:-$HOME/data/geo3k/test.parquet} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-512} +PPO_MINI_BATCH_SIZE=${PPO_MINI_BATCH_SIZE:-128} +PPO_MICRO_BATCH_SIZE_PER_GPU=${PPO_MICRO_BATCH_SIZE_PER_GPU:-10} +LOG_PROB_MICRO_BATCH_SIZE_PER_GPU=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU:-20} +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-1024} +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-2048} + +ACTOR_LR=${ACTOR_LR:-1e-6} +KL_LOSS_COEF=${KL_LOSS_COEF:-0.01} +ENTROPY_COEFF=${ENTROPY_COEFF:-0} + +ROLLOUT_TP=${ROLLOUT_TP:-2} +ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.6} +ROLLOUT_N=${ROLLOUT_N:-5} + +PROJECT_NAME=${PROJECT_NAME:-verl_grpo_example_geo3k} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-glm41v_9b_function_rm} +SAVE_FREQ=${SAVE_FREQ:-20} +TEST_FREQ=${TEST_FREQ:-5} +TOTAL_EPOCHS=${TOTAL_EPOCHS:-15} +# ---- end user-adjustable ---- +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + data.train_files=${TRAIN_FILE} + data.val_files=${TEST_FILE} + data.train_batch_size=${TRAIN_BATCH_SIZE} + data.max_prompt_length=${MAX_PROMPT_LENGTH} + data.max_response_length=${MAX_RESPONSE_LENGTH} + data.filter_overlong_prompts=True + data.truncation='error' + data.image_key=images + algorithm.use_kl_in_reward=False +) + +MODEL=( + actor_rollout_ref.model.path=${MODEL_PATH} + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${ACTOR_LR} + actor_rollout_ref.actor.ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE} + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${PPO_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${KL_LOSS_COEF} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${ENTROPY_COEFF} + actor_rollout_ref.actor.fsdp_config.param_offload=False + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False +) + +ROLLOUT=( + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.rollout.tensor_model_parallel_size=${ROLLOUT_TP} + actor_rollout_ref.rollout.name=${INFER_BACKEND} + actor_rollout_ref.rollout.gpu_memory_utilization=${ROLLOUT_GPU_MEM_UTIL} + actor_rollout_ref.rollout.enable_chunked_prefill=False + actor_rollout_ref.rollout.enforce_eager=False + actor_rollout_ref.rollout.free_cache_engine=True + actor_rollout_ref.rollout.n=${ROLLOUT_N} +) + +REF=( + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.ref.fsdp_config.param_offload=True +) + +TRAINER=( + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${PROJECT_NAME} + trainer.experiment_name=${EXPERIMENT_NAME} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${SAVE_FREQ} + trainer.test_freq=${TEST_FREQ} + trainer.total_epochs=${TOTAL_EPOCHS} +) + +EXTRA=( +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/grpo_trainer/run_gpt_oss_20b_fsdp.sh b/verl/examples/grpo_trainer/run_gpt_oss_20b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..e005a0769a9de6513d32b7c60658dcb94fe2380b --- /dev/null +++ b/verl/examples/grpo_trainer/run_gpt_oss_20b_fsdp.sh @@ -0,0 +1,140 @@ +#!/bin/bash + +# ---- user-adjustable ---- +SOURCE_MODEL_ID=${SOURCE_MODEL_ID:-openai/gpt-oss-20b} +MODEL_DIR=${MODEL_DIR:-$HOME/models/gpt-oss-20b-bf16} +TRAIN_FILE=${TRAIN_FILE:-$HOME/data/gsm8k/train.parquet} +TEST_FILE=${TEST_FILE:-$HOME/data/gsm8k/test.parquet} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-256} +PPO_MINI_BATCH_SIZE=${PPO_MINI_BATCH_SIZE:-256} +PPO_MICRO_BATCH_SIZE_PER_GPU=${PPO_MICRO_BATCH_SIZE_PER_GPU:-32} +LOG_PROB_MICRO_BATCH_SIZE_PER_GPU=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU:-32} +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-512} +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-8192} +REASONING_EFFORT=${REASONING_EFFORT:-medium} + +ACTOR_LR=${ACTOR_LR:-1e-6} +KL_LOSS_COEF=${KL_LOSS_COEF:-0.001} +ENTROPY_COEFF=${ENTROPY_COEFF:-0} + +ROLLOUT_TP=${ROLLOUT_TP:-2} +ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.7} +ROLLOUT_N=${ROLLOUT_N:-5} + +PROJECT_NAME=${PROJECT_NAME:-verl_grpo_example_gsm8k_math} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-oai_oss_20b_function_rm} +SAVE_FREQ=${SAVE_FREQ:-50} +TEST_FREQ=${TEST_FREQ:-10} +TOTAL_EPOCHS=${TOTAL_EPOCHS:-15} +# ---- end user-adjustable ---- + +# ---- no user adjustment needed below ---- + +cat > get_model.py << EOF +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer, Mxfp4Config + +model_id = "${SOURCE_MODEL_ID}" +output_dir = "${MODEL_DIR}" + +quantization_config = Mxfp4Config(dequantize=True) +model_kwargs = dict( + attn_implementation="eager", + torch_dtype=torch.bfloat16, + quantization_config=quantization_config, + use_cache=False, + device_map="auto", +) + +model = AutoModelForCausalLM.from_pretrained(model_id, **model_kwargs) + +# Patch config with custom attribute before saving +model.config.attn_implementation = "eager" + +model.save_pretrained(output_dir) +tokenizer = AutoTokenizer.from_pretrained(model_id) +tokenizer.save_pretrained(output_dir) +EOF + +python get_model.py +# or you can use lmsys/gpt-oss-20b-bf16 +# recommend to use same value for train_batch_size and ppo_mini_batch_size +# to avoid MOE training instability +# use large value for max_response_length if you want to use reasoning effort high. +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + data.train_files="${TRAIN_FILE}" + data.val_files="${TEST_FILE}" + data.train_batch_size=${TRAIN_BATCH_SIZE} + data.max_prompt_length=${MAX_PROMPT_LENGTH} + data.max_response_length=${MAX_RESPONSE_LENGTH} + data.filter_overlong_prompts=True + data.truncation='error' + +data.apply_chat_template_kwargs.reasoning_effort=${REASONING_EFFORT} + algorithm.use_kl_in_reward=False +) + +MODEL=( + actor_rollout_ref.model.path=${MODEL_DIR} + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${ACTOR_LR} + actor_rollout_ref.actor.ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE} + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${PPO_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${KL_LOSS_COEF} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${ENTROPY_COEFF} + actor_rollout_ref.actor.fsdp_config.param_offload=False + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False + +actor_rollout_ref.actor.fsdp_config.model_dtype=bfloat16 +) + +ROLLOUT=( + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.rollout.tensor_model_parallel_size=${ROLLOUT_TP} + actor_rollout_ref.rollout.name=sglang + actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend=triton + actor_rollout_ref.rollout.gpu_memory_utilization=${ROLLOUT_GPU_MEM_UTIL} + actor_rollout_ref.rollout.n=${ROLLOUT_N} + actor_rollout_ref.rollout.load_format=safetensors +) + +REF=( + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.ref.fsdp_config.param_offload=True +) + +TRAINER=( + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${PROJECT_NAME} + trainer.experiment_name=${EXPERIMENT_NAME} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${SAVE_FREQ} + trainer.test_freq=${TEST_FREQ} + trainer.total_epochs=${TOTAL_EPOCHS} +) + +EXTRA=( +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/grpo_trainer/run_minicpm_o_2_6_fsdp.sh b/verl/examples/grpo_trainer/run_minicpm_o_2_6_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..b439375f69cdfd8434fcc67f2341c87b4a45e006 --- /dev/null +++ b/verl/examples/grpo_trainer/run_minicpm_o_2_6_fsdp.sh @@ -0,0 +1,110 @@ +set -x + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-openbmb/MiniCPM-o-2_6} +TRAIN_FILE=${TRAIN_FILE:-$HOME/data/geo3k/train.parquet} +TEST_FILE=${TEST_FILE:-$HOME/data/geo3k/test.parquet} +CUSTOM_CLS_PATH=${CUSTOM_CLS_PATH:-recipe/minicpmo/rl_dataset.py} +CUSTOM_CLS_NAME=${CUSTOM_CLS_NAME:-RLHFDataset} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-128} +PPO_MINI_BATCH_SIZE=${PPO_MINI_BATCH_SIZE:-32} +PPO_MICRO_BATCH_SIZE_PER_GPU=${PPO_MICRO_BATCH_SIZE_PER_GPU:-16} +LOG_PROB_MICRO_BATCH_SIZE_PER_GPU=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU:-16} +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-1024} +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-2048} + +ACTOR_LR=${ACTOR_LR:-1e-6} +KL_LOSS_COEF=${KL_LOSS_COEF:-0.001} + +ROLLOUT_TP=${ROLLOUT_TP:-2} +ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.6} +ROLLOUT_N=${ROLLOUT_N:-8} + +PROJECT_NAME=${PROJECT_NAME:-verl_grpo_example_geo3k} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-minicpmo2_6_function_rm} +SAVE_FREQ=${SAVE_FREQ:--1} +TEST_FREQ=${TEST_FREQ:-5} +TOTAL_EPOCHS=${TOTAL_EPOCHS:-15} +# ---- end user-adjustable ---- +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + data.train_files=${TRAIN_FILE} + data.val_files=${TEST_FILE} + data.train_batch_size=${TRAIN_BATCH_SIZE} + data.max_prompt_length=${MAX_PROMPT_LENGTH} + data.max_response_length=${MAX_RESPONSE_LENGTH} + data.filter_overlong_prompts=False + data.truncation='error' + data.image_key=images + data.trust_remote_code=True + data.custom_cls.path=${CUSTOM_CLS_PATH} + data.custom_cls.name=${CUSTOM_CLS_NAME} + algorithm.kl_ctrl.kl_coef=${KL_LOSS_COEF} +) + +MODEL=( + actor_rollout_ref.model.path=${MODEL_PATH} + actor_rollout_ref.model.trust_remote_code=True + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${ACTOR_LR} + actor_rollout_ref.actor.ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE} + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${PPO_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.use_dynamic_bsz=False + actor_rollout_ref.actor.kl_loss_coef=${KL_LOSS_COEF} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.fsdp_config.param_offload=False + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False + actor_rollout_ref.actor.fsdp_config.use_orig_params=True +) + +ROLLOUT=( + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.rollout.tensor_model_parallel_size=${ROLLOUT_TP} + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.gpu_memory_utilization=${ROLLOUT_GPU_MEM_UTIL} + actor_rollout_ref.rollout.enable_chunked_prefill=False + actor_rollout_ref.rollout.enforce_eager=False + actor_rollout_ref.rollout.free_cache_engine=False + actor_rollout_ref.rollout.n=${ROLLOUT_N} +) + +REF=( + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.ref.fsdp_config.param_offload=True +) + +TRAINER=( + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${PROJECT_NAME} + trainer.experiment_name=${EXPERIMENT_NAME} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${SAVE_FREQ} + trainer.test_freq=${TEST_FREQ} + trainer.total_epochs=${TOTAL_EPOCHS} +) + +EXTRA=( +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/grpo_trainer/run_mistral_nemo_12b_skyworkrm_fsdp.sh b/verl/examples/grpo_trainer/run_mistral_nemo_12b_skyworkrm_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..dd75353308ccd4a3dbefc6e2135dd3de373d9491 --- /dev/null +++ b/verl/examples/grpo_trainer/run_mistral_nemo_12b_skyworkrm_fsdp.sh @@ -0,0 +1,104 @@ +# ---- user-adjustable ---- +TRAIN_FILE=${TRAIN_FILE:-data/full_hh_rlhf/rl/train.parquet} +TEST_FILE=${TEST_FILE:-data/full_hh_rlhf/rl/train.parquet} # no use +MODEL_PATH=${MODEL_PATH:-mistralai/Mistral-Nemo-Instruct-2407} +REWARD_MODEL_PATH=${REWARD_MODEL_PATH:-Skywork/Skywork-Reward-Llama-3.1-8B} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-512} +PPO_MINI_BATCH_SIZE=${PPO_MINI_BATCH_SIZE:-128} +PPO_MICRO_BATCH_SIZE_PER_GPU=${PPO_MICRO_BATCH_SIZE_PER_GPU:-10} +LOG_PROB_MICRO_BATCH_SIZE_PER_GPU=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU:-10} +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-4096} +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-2048} + +ACTOR_LR=${ACTOR_LR:-1e-6} +ROLLOUT_TP=${ROLLOUT_TP:-4} +ROLLOUT_N=${ROLLOUT_N:-5} +ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.7} +REWARD_NUM_WORKERS=${REWARD_NUM_WORKERS:-8} +REWARD_ROLLOUT_GPU_MEM_UTIL=${REWARD_ROLLOUT_GPU_MEM_UTIL:-0.8} +REWARD_ROLLOUT_TP=${REWARD_ROLLOUT_TP:-1} +REWARD_PROMPT_LENGTH=${REWARD_PROMPT_LENGTH:-8192} +REWARD_RESPONSE_LENGTH=${REWARD_RESPONSE_LENGTH:-4096} + +ADV_ESTIMATOR=${ADV_ESTIMATOR:-grpo} +PROJECT_NAME=${PROJECT_NAME:-verl_full_hh_rlhf_examples} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-grpo_mistral13B-skyworkLlama8b-hhrlhf} +SAVE_FREQ=${SAVE_FREQ:-10} +TEST_FREQ=${TEST_FREQ:--1} +TOTAL_EPOCHS=${TOTAL_EPOCHS:-5} +# ---- end user-adjustable ---- +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=${ADV_ESTIMATOR} + data.train_files="${TRAIN_FILE}" + data.val_files="${TEST_FILE}" + data.train_batch_size=${TRAIN_BATCH_SIZE} + data.prompt_key="prompt" + data.return_raw_chat=True + data.max_prompt_length=${MAX_PROMPT_LENGTH} + data.max_response_length=${MAX_RESPONSE_LENGTH} + data.filter_overlong_prompts=True + data.truncation='error' + algorithm.use_kl_in_reward=False +) + +MODEL=( + actor_rollout_ref.model.path=${MODEL_PATH} + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${ACTOR_LR} + actor_rollout_ref.actor.ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE} + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${PPO_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.actor.use_kl_loss=False +) + +ROLLOUT=( + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.n=${ROLLOUT_N} + actor_rollout_ref.rollout.gpu_memory_utilization=${ROLLOUT_GPU_MEM_UTIL} + actor_rollout_ref.rollout.tensor_model_parallel_size=${ROLLOUT_TP} +) + +REWARD=( + reward.num_workers=${REWARD_NUM_WORKERS} + reward.reward_model.enable=True + reward.reward_model.model_path=${REWARD_MODEL_PATH} + reward.reward_model.rollout.name=vllm + reward.reward_model.rollout.gpu_memory_utilization=${REWARD_ROLLOUT_GPU_MEM_UTIL} + reward.reward_model.rollout.tensor_model_parallel_size=${REWARD_ROLLOUT_TP} + reward.reward_model.rollout.prompt_length=${REWARD_PROMPT_LENGTH} + reward.reward_model.rollout.response_length=${REWARD_RESPONSE_LENGTH} +) + +TRAINER=( + trainer.logger='["console","wandb"]' + trainer.val_before_train=False + trainer.project_name=${PROJECT_NAME} + trainer.experiment_name=${EXPERIMENT_NAME} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${SAVE_FREQ} + trainer.test_freq=${TEST_FREQ} + trainer.total_epochs=${TOTAL_EPOCHS} +) + +EXTRA=( +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REWARD[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" \ No newline at end of file diff --git a/verl/examples/grpo_trainer/run_moonlight_16b_a3b_megatron.sh b/verl/examples/grpo_trainer/run_moonlight_16b_a3b_megatron.sh new file mode 100644 index 0000000000000000000000000000000000000000..e6a7013658537aa479475c32fcf56dbf517f328e --- /dev/null +++ b/verl/examples/grpo_trainer/run_moonlight_16b_a3b_megatron.sh @@ -0,0 +1,117 @@ +set -x + +export CUDA_DEVICE_MAX_CONNECTIONS=1 # For megatron communication/computation overlapping + +# ---- user-adjustable ---- +HF_MODEL_PATH=${HF_MODEL_PATH:-moonshotai/Moonlight-16B-A3B} +DIST_CKPT_PATH=${DIST_CKPT_PATH} +TRAIN_FILE=${TRAIN_FILE:-$HOME/data/gsm8k/train.parquet} +TEST_FILE=${TEST_FILE:-$HOME/data/gsm8k/test.parquet} +NNODES=${NNODES:-3} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-192} +PPO_MINI_BATCH_SIZE=${PPO_MINI_BATCH_SIZE:-64} +PPO_MICRO_BATCH_SIZE_PER_GPU=${PPO_MICRO_BATCH_SIZE_PER_GPU:-16} +LOG_PROB_MICRO_BATCH_SIZE_PER_GPU=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU:-16} +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-1024} +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-2048} + +ACTOR_LR=${ACTOR_LR:-1e-6} +KL_LOSS_COEF=${KL_LOSS_COEF:-0.001} +ENTROPY_COEFF=${ENTROPY_COEFF:-0} + +ACTOR_TP=${ACTOR_TP:-4} +ACTOR_PP=${ACTOR_PP:-3} +ACTOR_EP=${ACTOR_EP:-4} +ACTOR_ETP=${ACTOR_ETP:-1} +ROLLOUT_TP=${ROLLOUT_TP:-4} +ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.6} +ROLLOUT_N=${ROLLOUT_N:-5} + +PROJECT_NAME=${PROJECT_NAME:-verl_grpo_example_gsm8k_math} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-moonlight_megatron_ep} +SAVE_FREQ=${SAVE_FREQ:-20} +TEST_FREQ=${TEST_FREQ:-5} +TOTAL_EPOCHS=${TOTAL_EPOCHS:-15} +# ---- end user-adjustable ---- +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + data.train_files="$TRAIN_FILE" + data.val_files="$TEST_FILE" + data.train_batch_size=${TRAIN_BATCH_SIZE} + data.max_prompt_length=${MAX_PROMPT_LENGTH} + data.max_response_length=${MAX_RESPONSE_LENGTH} + data.filter_overlong_prompts=True + data.truncation='error' + data.trust_remote_code=True + algorithm.use_kl_in_reward=False +) + +MODEL=( + actor_rollout_ref.model.path=$HF_MODEL_PATH + actor_rollout_ref.model.trust_remote_code=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${ACTOR_LR} + actor_rollout_ref.actor.ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE} + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${PPO_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${ACTOR_PP} + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${ACTOR_TP} + actor_rollout_ref.actor.megatron.expert_model_parallel_size=${ACTOR_EP} + actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${ACTOR_ETP} + actor_rollout_ref.actor.megatron.use_dist_checkpointing=True + actor_rollout_ref.actor.megatron.dist_checkpointing_path=$DIST_CKPT_PATH + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${KL_LOSS_COEF} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${ENTROPY_COEFF} +) + +ROLLOUT=( + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.rollout.tensor_model_parallel_size=${ROLLOUT_TP} + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.gpu_memory_utilization=${ROLLOUT_GPU_MEM_UTIL} + actor_rollout_ref.rollout.n=${ROLLOUT_N} +) + +REF=( + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${ACTOR_PP} + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${ACTOR_TP} + actor_rollout_ref.ref.megatron.expert_model_parallel_size=${ACTOR_EP} + actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=${ACTOR_ETP} + actor_rollout_ref.ref.megatron.use_dist_checkpointing=True + actor_rollout_ref.ref.megatron.dist_checkpointing_path=$DIST_CKPT_PATH +) + +TRAINER=( + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${PROJECT_NAME} + trainer.experiment_name=${EXPERIMENT_NAME} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${SAVE_FREQ} + trainer.test_freq=${TEST_FREQ} + trainer.total_epochs=${TOTAL_EPOCHS} +) + +EXTRA=( + model_engine=megatron +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/grpo_trainer/run_nemotron_nano_v3_30b_a3b_megatron.sh b/verl/examples/grpo_trainer/run_nemotron_nano_v3_30b_a3b_megatron.sh new file mode 100644 index 0000000000000000000000000000000000000000..23d38500f492bc093ec4db013d9001fe0a49fc87 --- /dev/null +++ b/verl/examples/grpo_trainer/run_nemotron_nano_v3_30b_a3b_megatron.sh @@ -0,0 +1,225 @@ +#!/usr/bin/env bash +set -xeuo pipefail + +################################################### environment ################################################### +### # 1. use docker image `verlai/verl:vllm015.dev`` and install correct dependencies: +# pip install nvidia-modelopt +# MAX_JOBS=32 pip install git+https://github.com/Dao-AILab/causal-conv1d.git --no-build-isolation --no-cache-dir +# MAX_JOBS=32 pip install git+https://github.com/state-spaces/mamba.git --no-build-isolation --no-cache-dir +# pip install --no-deps git+https://github.com/NVIDIA-NeMo/Megatron-Bridge +# pip install --no-deps git+https://github.com/NVIDIA/Megatron-LM.git@core_dev_r0.16.0 +# unset ROCR_VISIBLE_DEVICES +# unset PYTORCH_CUDA_ALLOC_CONF + +################################################### quick config ################################################### + +# ---- user-adjustable ---- +rollout_mode="async" +return_raw_chat="False" +rollout_name="vllm" # sglang or vllm +if [ "$rollout_mode" = "async" ]; then + export VLLM_USE_V1=1 + return_raw_chat="True" +fi +dtype="bfloat16" + +project_name='DAPO' +exp_name='nano_30b' + +adv_estimator=grpo + +use_kl_in_reward=False +kl_coef=0.0 +use_kl_loss=False +kl_loss_coef=0.0 + +clip_ratio_low=0.2 +clip_ratio_high=0.28 + +max_prompt_length=$((1024 * 2)) +max_response_length=$((1024 * 2)) +enable_overlong_buffer=True +overlong_buffer_len=$((1024 * 1)) +overlong_penalty_factor=1.0 + +loss_agg_mode="token-mean" + +train_prompt_bsz=32 +n_resp_per_prompt=16 +train_prompt_mini_bsz=32 + +# Ray +RAY_ADDRESS=${RAY_ADDRESS:-"http://localhost:8265"} +WORKING_DIR=${WORKING_DIR:-"${PWD}"} +RUNTIME_ENV=${RUNTIME_ENV:-"${WORKING_DIR}/verl/verl/trainer/runtime_env.yaml"} +NNODES=${NNODES:-1} +# Paths +MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"} +CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"} +TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/dapo-math-17k.parquet"} +TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/aime-2024.parquet"} + +# Algorithm +temperature=1.0 +top_p=1.0 +top_k=-1 # 0 for HF rollout, -1 for vLLM rollout +val_top_p=0.7 + +# Performance Related Parameter +use_dynamic_bsz=True +actor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 10 / 10)) +infer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 1)) +offload=True +gen_tp=8 +train_tp=8 +train_pp=1 +EP=8 +ETP=1 +# ---- end user-adjustable ---- + +# ---- no user adjustment needed below ---- +################################################### start of config ################################################### + +FP8=( + # train + # +actor_rollout_ref.actor.megatron.override_transformer_config.fp8="e4m3" # e4m3 or hybrid + # +actor_rollout_ref.actor.megatron.override_transformer_config.fp8_recipe="blockwise" + # +actor_rollout_ref.actor.optim.override_optimizer_config.fp8_recipe="blockwise" + # rollout + actor_rollout_ref.actor.megatron.dtype=${dtype} + actor_rollout_ref.rollout.dtype=${dtype} + # +actor_rollout_ref.rollout.quantization="fp8" +) + +DATA=( + data.train_files="${TRAIN_FILE}" + data.val_files="${TEST_FILE}" + data.prompt_key=prompt + data.return_raw_chat=$return_raw_chat + data.truncation='left' + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.train_batch_size=${train_prompt_bsz} +) + +REWARD_MODEL=( + +reward_model.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} + +reward_model.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} + +reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} + +reward_model.reward_kwargs.overlong_buffer_cfg.log=False + +reward_model.reward_kwargs.max_resp_len=${max_response_length} + reward_model.reward_manager=dapo +) + +PERF_OPT=( + +actor_rollout_ref.actor.megatron.override_transformer_config.apply_rope_fusion=True + actor_rollout_ref.model.use_fused_kernels=False + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32 + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 + +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=1 \ + +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True \ + +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=True \ + +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True \ +) + +ACTOR=( + actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} + actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} + actor_rollout_ref.actor.clip_ratio_c=10.0 + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 + actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} + actor_rollout_ref.actor.optim.lr=1e-6 + actor_rollout_ref.actor.optim.lr_warmup_steps=10 + actor_rollout_ref.actor.optim.weight_decay=0.1 + actor_rollout_ref.actor.optim.clip_grad=1.0 + actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} + actor_rollout_ref.actor.megatron.param_offload=${offload} + actor_rollout_ref.actor.megatron.optimizer_offload=${offload} + actor_rollout_ref.actor.megatron.grad_offload=${offload} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${train_pp} + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${train_tp} + actor_rollout_ref.actor.megatron.expert_model_parallel_size=$EP + actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=$ETP + actor_rollout_ref.actor.entropy_coeff=0 + actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} + actor_rollout_ref.actor.megatron.use_mbridge=True + actor_rollout_ref.actor.megatron.vanilla_mbridge=False +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=${rollout_name} + actor_rollout_ref.rollout.mode=${rollout_mode} + actor_rollout_ref.rollout.gpu_memory_utilization=0.70 + actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} + actor_rollout_ref.rollout.enable_chunked_prefill=True + actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) + actor_rollout_ref.rollout.temperature=${temperature} + actor_rollout_ref.rollout.top_p=${top_p} + actor_rollout_ref.rollout.top_k=${top_k} + actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} + actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} + actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} + actor_rollout_ref.rollout.val_kwargs.do_sample=True + actor_rollout_ref.rollout.val_kwargs.n=1 + actor_rollout_ref.rollout.calculate_log_probs=True + actor_rollout_ref.rollout.n=${n_resp_per_prompt} +) + +TRAINER=( + trainer.logger=['console','wandb'] + trainer.project_name="${project_name}" + trainer.experiment_name="${exp_name}" + trainer.n_gpus_per_node=8 + trainer.nnodes="${NNODES}" + trainer.val_before_train=False + trainer.test_freq=10 + trainer.save_freq=-1 + trainer.total_epochs=10 + trainer.default_local_dir="${CKPTS_DIR}" + trainer.resume_mode=auto + trainer.log_val_generations=10 +) + +FORWARD_ONLY_SETS=( + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} +) + +MODEL=( + actor_rollout_ref.model.path="${MODEL_PATH}" + actor_rollout_ref.model.trust_remote_code=True +) + +ALGORITHM=( + algorithm.adv_estimator=${adv_estimator} + algorithm.use_kl_in_reward=${use_kl_in_reward} + algorithm.kl_ctrl.kl_coef=${kl_coef} +) +EXTRA=( + model_engine=megatron +) + +################################################### start script ################################################### + +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${ALGORITHM[@]}" \ + "${MODEL[@]}" \ + "${ROLLOUT[@]}" \ + "${ACTOR[@]}" \ + "${REWARD_MODEL[@]}" \ + "${FP8[@]}" \ + "${PERF_OPT[@]}" \ + "${TRAINER[@]}" \ + "${FORWARD_ONLY_SETS[@]}" \ + "${EXTRA[@]}" \ + "$@" \ No newline at end of file diff --git a/verl/examples/grpo_trainer/run_qwen2-7b_math_megatron_fsdp.sh b/verl/examples/grpo_trainer/run_qwen2-7b_math_megatron_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..664cad814d44ec94afa7473c91fdfa832eb9f321 --- /dev/null +++ b/verl/examples/grpo_trainer/run_qwen2-7b_math_megatron_fsdp.sh @@ -0,0 +1,110 @@ +#!/usr/bin/env bash +set -xeuo pipefail + +export CUDA_DEVICE_MAX_CONNECTIONS=1 # For megatron communication/computation overlapping +unset ROCR_VISIBLE_DEVICES +export VLLM_USE_V1=1 +export VLLM_ALLREDUCE_USE_SYMM_MEM=0 + +########################### Quick Config ########################### + +TP=${TP:-4} +PP=${PP:-1} +GEN_TP=${GEN_TP:-4} + +rollout_mode=${rollout_mode:-async} +return_raw_chat=${return_raw_chat:-True} +USE_FUSED_KERNELS=${USE_FUSED_KERNELS:-False} + +HF_MODEL_PATH=${HF_MODEL_PATH:-Qwen/Qwen2.5-Math-7B} +gsm8k_train_path=${gsm8k_train_path:-$HOME/data/gsm8k/train.parquet} +gsm8k_test_path=${gsm8k_test_path:-$HOME/data/gsm8k/test.parquet} +math_train_path=${math_train_path:-$HOME/data/math/train.parquet} +math_test_path=${math_test_path:-$HOME/data/math/test.parquet} + +train_files=${train_files:-"['$gsm8k_train_path', '$math_train_path']"} +test_files=${test_files:-"['$gsm8k_test_path', '$math_test_path']"} + +########################### Parameter Arrays ########################### + +DATA=( + "data.train_files=${train_files}" + "data.val_files=${test_files}" + "data.return_raw_chat=${return_raw_chat}" + data.train_batch_size=32 + data.max_prompt_length=512 + data.max_response_length=512 + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + "actor_rollout_ref.model.path=${HF_MODEL_PATH}" + "actor_rollout_ref.model.use_fused_kernels=${USE_FUSED_KERNELS}" +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=1e-6 + actor_rollout_ref.actor.ppo_mini_batch_size=16 + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=0.001 + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=0 + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${PP} + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${TP} + actor_rollout_ref.actor.megatron.use_mbridge=True + actor_rollout_ref.actor.megatron.vanilla_mbridge=False + actor_rollout_ref.actor.megatron.use_megatron_fsdp=True + ++actor_rollout_ref.actor.megatron.override_transformer_config.gradient_accumulation_fusion=False +) + +ROLLOUT=( + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=2 + actor_rollout_ref.rollout.tensor_model_parallel_size=${GEN_TP} + actor_rollout_ref.rollout.name=vllm + "actor_rollout_ref.rollout.mode=${rollout_mode}" + actor_rollout_ref.rollout.gpu_memory_utilization=0.4 + actor_rollout_ref.rollout.n=2 +) + +REF=( + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=2 + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${PP} + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${TP} + actor_rollout_ref.ref.megatron.use_mbridge=True + actor_rollout_ref.ref.megatron.vanilla_mbridge=False + actor_rollout_ref.ref.megatron.use_megatron_fsdp=True + ++actor_rollout_ref.ref.megatron.override_transformer_config.gradient_accumulation_fusion=False +) + +ALGORITHM=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False +) + +TRAINER=( + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name='verl_grpo_example_gsm8k_math' + trainer.experiment_name='qwen2_7b_megatron_fsdp' + trainer.n_gpus_per_node=8 + trainer.nnodes=1 + trainer.save_freq=20 + trainer.test_freq=5 + trainer.total_epochs=15 +) + +########################### Launch ########################### + +python3 -m verl.trainer.main_ppo \ + --config-path=config \ + --config-name='ppo_megatron_trainer.yaml' \ + "${DATA[@]}" \ + "${ALGORITHM[@]}" \ + "${MODEL[@]}" \ + "${ROLLOUT[@]}" \ + "${ACTOR[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "$@" diff --git a/verl/examples/grpo_trainer/run_qwen2_5_32b_fsdp.sh b/verl/examples/grpo_trainer/run_qwen2_5_32b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..4e3899c6671a28bc12fc4d3358dea7fb2a030f19 --- /dev/null +++ b/verl/examples/grpo_trainer/run_qwen2_5_32b_fsdp.sh @@ -0,0 +1,123 @@ +#!/usr/bin/env bash +# GRPO | Qwen2.5-32B | FSDP training | NVIDIA GPUs or Ascend NPUs +# +# INFER_BACKEND controls rollout backend: vllm + +set -xeuo pipefail + +# ---- user-adjustable ---- +INFER_BACKEND=${INFER_BACKEND:-vllm} +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen2.5-32B} +TRAIN_FILE=${TRAIN_FILE:-$HOME/data/gsm8k/train.parquet} +TEST_FILE=${TEST_FILE:-$HOME/data/gsm8k/test.parquet} +NNODES=${NNODES:-2} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-1024} +PPO_MINI_BATCH_SIZE=${PPO_MINI_BATCH_SIZE:-128} +PPO_MICRO_BATCH_SIZE_PER_GPU=${PPO_MICRO_BATCH_SIZE_PER_GPU:-2} +LOG_PROB_MICRO_BATCH_SIZE_PER_GPU=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU:-2} +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-1024} +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-1024} + +ACTOR_LR=${ACTOR_LR:-1e-6} +KL_LOSS_COEF=${KL_LOSS_COEF:-0.001} +ENTROPY_COEFF=${ENTROPY_COEFF:-0} + +ROLLOUT_TP=${ROLLOUT_TP:-8} +ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.4} +ROLLOUT_N=${ROLLOUT_N:-5} + +PROJECT_NAME=${PROJECT_NAME:-verl_grpo_example_geo3k} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-glm41v_9b_function_rm} +SAVE_FREQ=${SAVE_FREQ:-20} +TEST_FREQ=${TEST_FREQ:-5} +TOTAL_EPOCHS=${TOTAL_EPOCHS:-15} +# ---- end user-adjustable ---- + +case "${DEVICE}" in + gpu) + ;; + npu) + NGPUS_PER_NODE=16 + ;; + *) + echo "Unsupported DEVICE=${DEVICE}. Expected 'gpu' or 'npu'." >&2 + exit 1 + ;; +esac + +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + data.train_files=${TRAIN_FILE} + data.val_files=${TEST_FILE} + data.train_batch_size=${TRAIN_BATCH_SIZE} + data.max_prompt_length=${MAX_PROMPT_LENGTH} + data.max_response_length=${MAX_RESPONSE_LENGTH} + data.filter_overlong_prompts=True + data.truncation='error' + algorithm.use_kl_in_reward=False +) + +MODEL=( + actor_rollout_ref.model.path=${MODEL_PATH} + actor_rollout_ref.model.use_remove_padding=False + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${ACTOR_LR} + actor_rollout_ref.actor.ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE} + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${PPO_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${KL_LOSS_COEF} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${ENTROPY_COEFF} + actor_rollout_ref.actor.fsdp_config.param_offload=False + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False +) + +ROLLOUT=( + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.rollout.tensor_model_parallel_size=${ROLLOUT_TP} + actor_rollout_ref.rollout.name=${INFER_BACKEND} + actor_rollout_ref.rollout.gpu_memory_utilization=${ROLLOUT_GPU_MEM_UTIL} + actor_rollout_ref.rollout.enable_chunked_prefill=False + actor_rollout_ref.rollout.enforce_eager=False + actor_rollout_ref.rollout.free_cache_engine=True + actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=4096 + actor_rollout_ref.rollout.n=${ROLLOUT_N} +) + +REF=( + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.ref.fsdp_config.param_offload=True +) + +TRAINER=( + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${PROJECT_NAME} + trainer.experiment_name=${EXPERIMENT_NAME} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${SAVE_FREQ} + trainer.test_freq=${TEST_FREQ} + trainer.total_epochs=${TOTAL_EPOCHS} +) + +EXTRA=( +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/grpo_trainer/run_qwen2_5_vl_7b_fsdp.sh b/verl/examples/grpo_trainer/run_qwen2_5_vl_7b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..38ab8558231c85502e0048e74f4728d4853ec8c1 --- /dev/null +++ b/verl/examples/grpo_trainer/run_qwen2_5_vl_7b_fsdp.sh @@ -0,0 +1,151 @@ +#!/usr/bin/env bash +# GRPO | Qwen2.5-VL-7B | FSDP training | NVIDIA GPUs or Ascend NPUs +# +# INFER_BACKEND controls rollout backend: vllm | sglang | trtllm. + +set -xeuo pipefail + +########################### user-adjustable ########################### +INFER_BACKEND=${INFER_BACKEND:-vllm} + +# DEVICE is auto-detected by probing torch_npu; override only for special cases. +DEVICE=${DEVICE:-$(python3 -c 'import torch_npu' 2>/dev/null && echo npu || echo gpu)} + +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen2.5-VL-7B-Instruct} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +train_batch_size=${TRAIN_BATCH_SIZE:-512} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-128} +max_prompt_length=${MAX_PROMPT_LENGTH:-1024} +max_response_length=${MAX_RESPONSE_LENGTH:-2048} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-24576} + +actor_lr=${ACTOR_LR:-1e-6} +kl_loss_coef=${KL_LOSS_COEF:-0.01} +entropy_coeff=${ENTROPY_COEFF:-0} + +rollout_tp=${ROLLOUT_TP:-} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-} +rollout_n=${ROLLOUT_N:-5} + +total_epochs=${TOTAL_EPOCHS:-15} +save_freq=${SAVE_FREQ:-20} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_grpo_geo3k} +experiment_name=${EXPERIMENT_NAME:-qwen2_5_vl_7b_${INFER_BACKEND}_fsdp} +########################### end user-adjustable ########################### + +########################### derived defaults ########################### +rollout_tp=${rollout_tp:-2} +rollout_gpu_mem_util=${rollout_gpu_mem_util:-0.6} + +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files=$HOME/data/geo3k/train.parquet + data.val_files=$HOME/data/geo3k/test.parquet + data.image_key=images + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.fsdp_config.param_offload=False + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=${INFER_BACKEND} + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.fsdp_config.param_offload=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +# Conservative rollout extras shared by all inference backends. +EXTRA=( + actor_rollout_ref.actor.strategy=fsdp2 + actor_rollout_ref.model.use_fused_kernels=True + actor_rollout_ref.rollout.multi_stage_wake_up=True + actor_rollout_ref.rollout.enable_chunked_prefill=False + actor_rollout_ref.rollout.enforce_eager=False + actor_rollout_ref.rollout.free_cache_engine=True +) + + +case "${DEVICE}" in + gpu) + ;; + npu) + TRAINER+=(trainer.n_gpus_per_node=16 + ) + ROLLOUT+=( + actor_rollout_ref.rollout.gpu_memory_utilization=0.5 + +actor_rollout_ref.rollout.engine_kwargs.vllm.mm_processor_cache_gb=0 + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 + ) + REF+=( + actor_rollout_ref.ref.fsdp_config.param_offload=True + ) + EXTRA+=( + actor_rollout_ref.model.use_fused_kernels=False + actor_rollout_ref.rollout.multi_stage_wake_up=False + actor_rollout_ref.rollout.free_cache_engine=False + ) + ;; + *) + echo "Unsupported DEVICE=${DEVICE}. Expected 'gpu' or 'npu'." >&2 + exit 1 + ;; +esac + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/grpo_trainer/run_qwen2_5_vl_7b_megatron.sh b/verl/examples/grpo_trainer/run_qwen2_5_vl_7b_megatron.sh new file mode 100644 index 0000000000000000000000000000000000000000..a41275140de2dbbe4b0629e9ea90c8a56bc03610 --- /dev/null +++ b/verl/examples/grpo_trainer/run_qwen2_5_vl_7b_megatron.sh @@ -0,0 +1,114 @@ +#!/usr/bin/env bash +# GRPO | vision | vLLM rollout | Megatron training | NVIDIA GPUs + +set -xeuo pipefail + +export CUDA_DEVICE_MAX_CONNECTIONS=1 + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen2.5-VL-7B-Instruct} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +train_batch_size=${TRAIN_BATCH_SIZE:-512} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-128} +max_prompt_length=${MAX_PROMPT_LENGTH:-1024} +max_response_length=${MAX_RESPONSE_LENGTH:-2048} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-24576} + +actor_lr=${ACTOR_LR:-1e-6} +kl_loss_coef=${KL_LOSS_COEF:-0.01} +entropy_coeff=${ENTROPY_COEFF:-0} + +actor_tp=${ACTOR_TP:-2} +actor_pp=${ACTOR_PP:-2} + +rollout_tp=${ROLLOUT_TP:-2} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.6} +rollout_n=${ROLLOUT_N:-5} + +total_epochs=${TOTAL_EPOCHS:-15} +save_freq=${SAVE_FREQ:-20} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_grpo_geo3k} +experiment_name=${EXPERIMENT_NAME:-qwen2_5_vl_7b_vllm_megatron} +# ---- end user-adjustable ---- +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files=$HOME/data/geo3k/train.parquet + data.val_files=$HOME/data/geo3k/test.parquet + data.image_key=images + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.use_fused_kernels=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${actor_pp} + actor_rollout_ref.actor.megatron.use_mbridge=True +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${actor_pp} + actor_rollout_ref.ref.megatron.use_mbridge=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( + model_engine=megatron +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/grpo_trainer/run_qwen3_235b_a22b_megatron.sh b/verl/examples/grpo_trainer/run_qwen3_235b_a22b_megatron.sh new file mode 100644 index 0000000000000000000000000000000000000000..9f51813ed873961681ea1ec44ba595b9603c1705 --- /dev/null +++ b/verl/examples/grpo_trainer/run_qwen3_235b_a22b_megatron.sh @@ -0,0 +1,199 @@ +#!/usr/bin/env bash +# GRPO scale demo | Qwen3-235B-A22B | vLLM rollout | Megatron training | GPU/NPU +# Requires multi-node clusters. DEVICE is auto-detected from torch_npu; override +# with DEVICE=gpu|npu only when the auto-detection is wrong. + +set -xeuo pipefail + +########################### user-adjustable ########################### +# DEVICE is auto-detected by probing torch_npu; override only for special cases. +DEVICE=${DEVICE:-$(python3 -c 'import torch_npu' 2>/dev/null && echo npu || echo gpu)} +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-235B-A22B} +MCORE_MODEL_PATH=${MCORE_MODEL_PATH:-} +NNODES=${NNODES:-8} +NDEVICES_PER_NODE=${NDEVICES_PER_NODE:-} + +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-128} +PPO_MINI_BATCH_SIZE=${PPO_MINI_BATCH_SIZE:-128} +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-8192} +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-4096} +PPO_MAX_TOKEN_LEN_PER_GPU=${PPO_MAX_TOKEN_LEN_PER_GPU:-$((MAX_PROMPT_LENGTH + MAX_RESPONSE_LENGTH))} + +ACTOR_LR=${ACTOR_LR:-1e-6} +KL_LOSS_COEF=${KL_LOSS_COEF:-0.001} +ENTROPY_COEFF=${ENTROPY_COEFF:-0} +CLIP_RATIO_LOW=${CLIP_RATIO_LOW:-0.2} +CLIP_RATIO_HIGH=${CLIP_RATIO_HIGH:-0.28} + +ACTOR_TP=${ACTOR_TP:-4} +ACTOR_PP=${ACTOR_PP:-8} +ACTOR_EP=${ACTOR_EP:-4} +ALL_OFFLOAD=${ALL_OFFLOAD:-True} + +ROLLOUT_TP=${ROLLOUT_TP:-8} +ROLLOUT_DP=${ROLLOUT_DP:-} +ROLLOUT_EP=${ROLLOUT_EP:-64} +ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.75} +ROLLOUT_N=${ROLLOUT_N:-16} +ROLLOUT_MAX_NUM_BATCHED_TOKENS=${ROLLOUT_MAX_NUM_BATCHED_TOKENS:-1024} + +TOTAL_EPOCHS=${TOTAL_EPOCHS:-1} +SAVE_FREQ=${SAVE_FREQ:-100} +TEST_FREQ=${TEST_FREQ:--1} + +PROJECT_NAME=${PROJECT_NAME:-verl_grpo_scale_demo} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-qwen3_235b_a22b_grpo_vllm_megatron_$(date +%Y%m%d_%H%M)} +CKPTS_DIR=${CKPTS_DIR:-.ckpt} + +TRAIN_FILE=${TRAIN_FILE:-$HOME/data/gsm8k/train.parquet} +TEST_FILE=${TEST_FILE:-$HOME/data/gsm8k/test.parquet} +########################### end user-adjustable ########################### + +########################### derived defaults ########################### +n_devices_per_node=${NDEVICES_PER_NODE:-8} + +case "${DEVICE}" in + gpu) + export CUDA_DEVICE_MAX_CONNECTIONS=1 + ;; + npu) + export HCCL_CONNECT_TIMEOUT=1500 + export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050 + export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050 + export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 + ;; + *) + echo "Unsupported DEVICE=${DEVICE}. Expected 'gpu' or 'npu'." >&2 + exit 1 + ;; +esac + +########################### parameter arrays ########################### + +ALGORITHM=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False +) + +DATA=( + data.train_files="$TRAIN_FILE" + data.val_files="$TEST_FILE" + data.train_batch_size=${TRAIN_BATCH_SIZE} + data.max_prompt_length=${MAX_PROMPT_LENGTH} + data.max_response_length=${MAX_RESPONSE_LENGTH} + data.filter_overlong_prompts=False + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${ACTOR_LR} + actor_rollout_ref.actor.ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${KL_LOSS_COEF} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${ENTROPY_COEFF} + actor_rollout_ref.actor.clip_ratio_low=${CLIP_RATIO_LOW} + actor_rollout_ref.actor.clip_ratio_high=${CLIP_RATIO_HIGH} + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${ACTOR_TP} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${ACTOR_PP} + actor_rollout_ref.actor.megatron.expert_model_parallel_size=${ACTOR_EP} + actor_rollout_ref.actor.megatron.param_offload=${ALL_OFFLOAD} + actor_rollout_ref.actor.megatron.optimizer_offload=${ALL_OFFLOAD} + actor_rollout_ref.actor.megatron.grad_offload=${ALL_OFFLOAD} + +actor_rollout_ref.actor.megatron.override_transformer_config.num_layers_in_first_pipeline_stage=11 + +actor_rollout_ref.actor.megatron.override_transformer_config.num_layers_in_last_pipeline_stage=11 + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${ROLLOUT_TP} + actor_rollout_ref.rollout.expert_parallel_size=${ROLLOUT_EP} + actor_rollout_ref.rollout.gpu_memory_utilization=${ROLLOUT_GPU_MEM_UTIL} + actor_rollout_ref.rollout.n=${ROLLOUT_N} + actor_rollout_ref.rollout.max_num_batched_tokens=${ROLLOUT_MAX_NUM_BATCHED_TOKENS} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} + actor_rollout_ref.rollout.enable_chunked_prefill=True + actor_rollout_ref.rollout.enable_prefix_caching=True + actor_rollout_ref.rollout.free_cache_engine=True +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${ACTOR_TP} + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${ACTOR_PP} + actor_rollout_ref.ref.megatron.expert_model_parallel_size=${ACTOR_EP} + actor_rollout_ref.ref.megatron.param_offload=${ALL_OFFLOAD} +) + +TRAINER=( + actor_rollout_ref.nccl_timeout=7200 + trainer.balance_batch=True + trainer.logger='["console","wandb"]' + trainer.project_name=${PROJECT_NAME} + trainer.experiment_name=${EXPERIMENT_NAME} + trainer.n_gpus_per_node=${n_devices_per_node} + trainer.nnodes=${NNODES} + trainer.val_before_train=False + trainer.save_freq=${SAVE_FREQ} + trainer.test_freq=${TEST_FREQ} + trainer.total_epochs=${TOTAL_EPOCHS} + trainer.default_local_dir="${CKPTS_DIR}" +) + +# ---- conditional / per-device extras (rolled into a single trailing array) ---- +# Seed with the always-present rollout mode so the array is never empty (Bash 3.x + set -u safe). +EXTRA=( + model_engine=megatron +) + +if [ "${DEVICE}" = npu ]; then + ACTOR+=( + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 + ) + TRAINER+=( + trainer.n_gpus_per_node=16 + ) + + EXTRA+=( + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 + actor_rollout_ref.rollout.data_parallel_size=${ROLLOUT_DP:-8} + actor_rollout_ref.rollout.enforce_eager=False + +actor_rollout_ref.rollout.engine_kwargs.vllm.compilation_config.cudagraph_capture_sizes=[8,16,32,64,128] + +actor_rollout_ref.rollout.engine_kwargs.vllm.compilation_config.cudagraph_mode=FULL_DECODE_ONLY + ) +elif [ -n "${ROLLOUT_DP}" ]; then + EXTRA+=(actor_rollout_ref.rollout.data_parallel_size=${ROLLOUT_DP}) +fi + +if [ -n "$MCORE_MODEL_PATH" ]; then + EXTRA+=( + actor_rollout_ref.actor.megatron.dist_checkpointing_path=${MCORE_MODEL_PATH} + actor_rollout_ref.actor.megatron.use_dist_checkpointing=True + actor_rollout_ref.ref.megatron.dist_checkpointing_path=${MCORE_MODEL_PATH} + actor_rollout_ref.ref.megatron.use_dist_checkpointing=True + ) +fi + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${ALGORITHM[@]}" \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/grpo_trainer/run_qwen3_30b_a3b_fsdp.sh b/verl/examples/grpo_trainer/run_qwen3_30b_a3b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..c6608db2eb08e830534c208e4646f4eff1025c2a --- /dev/null +++ b/verl/examples/grpo_trainer/run_qwen3_30b_a3b_fsdp.sh @@ -0,0 +1,113 @@ +#!/usr/bin/env bash +# GRPO | MoE text | vLLM rollout | FSDP2 training | NVIDIA GPUs + +set -xeuo pipefail + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-30B-A3B} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +train_batch_size=${TRAIN_BATCH_SIZE:-512} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-128} +max_prompt_length=${MAX_PROMPT_LENGTH:-2048} +max_response_length=${MAX_RESPONSE_LENGTH:-4096} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-32768} + +actor_lr=${ACTOR_LR:-1e-6} +kl_loss_coef=${KL_LOSS_COEF:-0.001} +entropy_coeff=${ENTROPY_COEFF:-0} + +rollout_tp=${ROLLOUT_TP:-4} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.6} +rollout_n=${ROLLOUT_N:-8} + +total_epochs=${TOTAL_EPOCHS:-15} +save_freq=${SAVE_FREQ:-20} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_grpo_gsm8k_math} +experiment_name=${EXPERIMENT_NAME:-qwen3_30b_a3b_vllm_fsdp} +# ---- end user-adjustable ---- + +gsm8k_train=$HOME/data/gsm8k/train.parquet +gsm8k_test=$HOME/data/gsm8k/test.parquet +math_train=$HOME/data/math/train.parquet +math_test=$HOME/data/math/test.parquet + +train_files="['$gsm8k_train', '$math_train']" +val_files="['$gsm8k_test', '$math_test']" +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files="$train_files" + data.val_files="$val_files" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.strategy=fsdp2 + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.fsdp_config.param_offload=True + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.fsdp_config.param_offload=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/grpo_trainer/run_qwen3_30b_a3b_megatron.sh b/verl/examples/grpo_trainer/run_qwen3_30b_a3b_megatron.sh new file mode 100644 index 0000000000000000000000000000000000000000..88dc6d83c816b959b6b9e85fc605d706eea2cf64 --- /dev/null +++ b/verl/examples/grpo_trainer/run_qwen3_30b_a3b_megatron.sh @@ -0,0 +1,267 @@ +#!/usr/bin/env bash +# GRPO | Qwen3-30B-A3B (MoE) | Megatron training | NVIDIA GPUs +# DAPO-style recipe on DAPO-Math-17k / AIME-2024. +# +# Knobs: +# INFER_BACKEND rollout backend: vllm | sglang | trtllm (default: vllm) +# ROLLOUT_QUANTIZATION fp8 to enable TRT-LLM FP8 rollout (default: unset) +# +# Ascend NPU users: see run_qwen3_30b_a3b_mindspeed.sh. + +set -xeuo pipefail +export CUDA_DEVICE_MAX_CONNECTIONS=1 + +########################### user-adjustable ########################### +INFER_BACKEND=${INFER_BACKEND:-vllm} +ROLLOUT_QUANTIZATION=${ROLLOUT_QUANTIZATION:-} + +DATA_DIR=${DATA_DIR:-"$PWD"} +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-30B-A3B-Base} +MCORE_MODEL_PATH=${MCORE_MODEL_PATH:-} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-${GPUS_PER_NODE:-8}} + +train_files=${TRAIN_FILES:-${DAPO_MATH_TRAIN:-${DATA_DIR}/data/DAPO-Math-17k/data/dapo-math-17k.parquet}} +val_files=${VAL_FILES:-${AIME_VAL:-${DATA_DIR}/data/AIME-2024/data/aime-2024.parquet}} + +train_batch_size=${TRAIN_BATCH_SIZE:-512} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-16} +max_prompt_length=${MAX_PROMPT_LENGTH:-2048} +max_response_length=${MAX_RESPONSE_LENGTH:-8192} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-30720} + +actor_lr=${ACTOR_LR:-1e-5} +entropy_coeff=${ENTROPY_COEFF:-0} +actor_clip_ratio_low=${ACTOR_CLIP_RATIO_LOW:-0.2} +actor_clip_ratio_high=${ACTOR_CLIP_RATIO_HIGH:-0.28} +actor_clip_ratio_c=${ACTOR_CLIP_RATIO_C:-10.0} +actor_ppo_micro_batch_size_per_gpu=${ACTOR_PPO_MICRO_BATCH_SIZE_PER_GPU:-2} + +actor_tp=${ACTOR_TP:-4} +actor_pp=${ACTOR_PP:-2} +actor_vpp=${ACTOR_VPP:-2} +actor_ep=${ACTOR_EP:-2} +actor_cp=${ACTOR_CP:-1} +ref_tp=${REF_TP:-${actor_tp}} +ref_pp=${REF_PP:-${actor_pp}} +ref_vpp=${REF_VPP:-2} +ref_ep=${REF_EP:-${actor_ep}} +ref_cp=${REF_CP:-1} +all_offload=${ALL_OFFLOAD:-True} + +rollout_tp=${ROLLOUT_TP:-4} +infer_tp=${INFER_TP:-${rollout_tp}} +gen_moe_tp=${GEN_MOE_TP:-2} +gen_moe_ep=${GEN_MOE_EP:-2} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.6} +rollout_n=${ROLLOUT_N:-8} +rollout_max_num_batched_tokens=${ROLLOUT_MAX_NUM_BATCHED_TOKENS:-10240} +rollout_max_model_len=${ROLLOUT_MAX_MODEL_LEN:-10240} +rollout_temperature=${ROLLOUT_TEMPERATURE:-1.0} +rollout_top_p=${ROLLOUT_TOP_P:-1} +trtllm_moe_backend=${TRTLLM_MOE_BACKEND:-DEEPGEMM} + +ref_log_prob_max_token_len_per_gpu=${REF_LOG_PROB_MAX_TOKEN_LEN_PER_GPU:-40960} +ref_log_prob_micro_batch_size_per_gpu=${REF_LOG_PROB_MICRO_BATCH_SIZE_PER_GPU:-4} +rollout_log_prob_max_token_len_per_gpu=${ROLLOUT_LOG_PROB_MAX_TOKEN_LEN_PER_GPU:-40960} +rollout_log_prob_micro_batch_size_per_gpu=${ROLLOUT_LOG_PROB_MICRO_BATCH_SIZE_PER_GPU:-4} + +val_do_sample=${VAL_DO_SAMPLE:-True} +val_temperature=${VAL_TEMPERATURE:-1.0} +val_top_p=${VAL_TOP_P:-0.7} +val_n=${VAL_N:-1} +log_val_generations=${LOG_VAL_GENERATIONS:-10} + +total_epochs=${TOTAL_EPOCHS:-1000} +save_freq=${SAVE_FREQ:-50} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_grpo_dapo_math} +experiment_name=${EXPERIMENT_NAME:-qwen3_30b_a3b_${INFER_BACKEND}_megatron${ROLLOUT_QUANTIZATION:+_${ROLLOUT_QUANTIZATION}}} +########################### end user-adjustable ########################### + +########################### derived defaults ########################### +if [ -n "${ROLLOUT_QUANTIZATION}" ] && [ "${INFER_BACKEND}" != trtllm ]; then + echo "ROLLOUT_QUANTIZATION is only supported with INFER_BACKEND=trtllm, got: ${INFER_BACKEND}" >&2 + exit 1 +fi + +if [ "${infer_tp}" = 4 ] || [ "${infer_tp}" = 2 ]; then + rollout_max_num_seqs=${ROLLOUT_MAX_NUM_SEQS:-1024} +else + rollout_max_num_seqs=${ROLLOUT_MAX_NUM_SEQS:-384} +fi + +[ "${actor_pp}" -gt 1 ] && actor_vpp_override=${actor_vpp} || actor_vpp_override=null +[ "${ref_pp}" -gt 1 ] && ref_vpp_override=${ref_vpp} || ref_vpp_override=null + +########################### parameter arrays ########################### + +ALGORITHM=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + algorithm.kl_ctrl.kl_coef=0.0 +) + +REWARD=( + reward_model.reward_manager=dapo + +reward_model.reward_kwargs.overlong_buffer_cfg.enable=True + +reward_model.reward_kwargs.overlong_buffer_cfg.len=4096 + +reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=1.0 + +reward_model.reward_kwargs.overlong_buffer_cfg.log=False + +reward_model.reward_kwargs.max_resp_len=${max_response_length} +) + +DATA=( + data.train_files="$train_files" + data.val_files="$val_files" + data.train_batch_size=${train_batch_size} + data.prompt_key=prompt + data.return_raw_chat=True + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=False + data.truncation=left +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_fused_kernels=True + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${actor_ppo_micro_batch_size_per_gpu} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=False + actor_rollout_ref.actor.kl_loss_coef=0.0 + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.clip_ratio_low=${actor_clip_ratio_low} + actor_rollout_ref.actor.clip_ratio_high=${actor_clip_ratio_high} + actor_rollout_ref.actor.clip_ratio_c=${actor_clip_ratio_c} + actor_rollout_ref.actor.loss_agg_mode=token-mean + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${actor_pp} + actor_rollout_ref.actor.megatron.virtual_pipeline_model_parallel_size=${actor_vpp_override} + actor_rollout_ref.actor.megatron.expert_model_parallel_size=${actor_ep} + actor_rollout_ref.actor.megatron.context_parallel_size=${actor_cp} + actor_rollout_ref.actor.megatron.param_offload=${all_offload} + actor_rollout_ref.actor.megatron.optimizer_offload=${all_offload} + actor_rollout_ref.actor.megatron.grad_offload=${all_offload} + actor_rollout_ref.actor.megatron.use_mbridge=True + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 + +actor_rollout_ref.actor.megatron.override_transformer_config.apply_rope_fusion=True + +actor_rollout_ref.actor.megatron.override_transformer_config.masked_softmax_fusion=True + +actor_rollout_ref.actor.megatron.override_transformer_config.bias_activation_fusion=True + +actor_rollout_ref.actor.megatron.override_transformer_config.bias_dropout_fusion=True + +actor_rollout_ref.actor.megatron.override_transformer_config.gradient_accumulation_fusion=True + +actor_rollout_ref.actor.megatron.override_transformer_config.deallocate_pipeline_outputs=True + +actor_rollout_ref.actor.megatron.override_transformer_config.persist_layer_norm=True + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_grouped_gemm=True + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=True +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=${INFER_BACKEND} + actor_rollout_ref.rollout.tensor_model_parallel_size=${infer_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.calculate_log_probs=True + actor_rollout_ref.rollout.enable_chunked_prefill=True + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${rollout_log_prob_max_token_len_per_gpu} + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${rollout_log_prob_micro_batch_size_per_gpu} + actor_rollout_ref.rollout.max_num_seqs=${rollout_max_num_seqs} + actor_rollout_ref.rollout.max_num_batched_tokens=${rollout_max_num_batched_tokens} + actor_rollout_ref.rollout.max_model_len=${rollout_max_model_len} + actor_rollout_ref.rollout.prompt_length=${max_prompt_length} + actor_rollout_ref.rollout.response_length=${max_response_length} + actor_rollout_ref.rollout.temperature=${rollout_temperature} + actor_rollout_ref.rollout.top_p=${rollout_top_p} + actor_rollout_ref.rollout.val_kwargs.do_sample=${val_do_sample} + actor_rollout_ref.rollout.val_kwargs.temperature=${val_temperature} + actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} + actor_rollout_ref.rollout.val_kwargs.n=${val_n} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ref_log_prob_max_token_len_per_gpu} + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${ref_log_prob_micro_batch_size_per_gpu} + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${ref_tp} + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${ref_pp} + actor_rollout_ref.ref.megatron.virtual_pipeline_model_parallel_size=${ref_vpp_override} + actor_rollout_ref.ref.megatron.expert_model_parallel_size=${ref_ep} + actor_rollout_ref.ref.megatron.context_parallel_size=${ref_cp} + actor_rollout_ref.ref.megatron.param_offload=${all_offload} + actor_rollout_ref.ref.megatron.use_mbridge=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} + trainer.resume_mode=auto + trainer.val_before_train=False + trainer.log_val_generations=${log_val_generations} +) + +# ---- conditional extras (rolled into a single trailing array) ---- +EXTRA=( + model_engine=megatron + actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=${TRTLLM_UPDATE_WEIGHTS_BUCKET_MEGABYTES:-4096} + +actor_rollout_ref.rollout.moe_tensor_parallel_size=${gen_moe_tp} + actor_rollout_ref.rollout.expert_parallel_size=${gen_moe_ep} +) + +if [ "${INFER_BACKEND}" = trtllm ]; then + EXTRA+=( + +actor_rollout_ref.rollout.engine_kwargs.trtllm.batch_wait_timeout_iters=32 + +actor_rollout_ref.rollout.engine_kwargs.trtllm.batch_wait_max_tokens_ratio=0.5 + +actor_rollout_ref.rollout.engine_kwargs.trtllm.moe_config.backend=${trtllm_moe_backend} + ) +fi + +if [ "${ROLLOUT_QUANTIZATION}" = fp8 ]; then + # Avoid stray launcher env leaking into Ray workers under FP8 TRT-LLM. + for v in $(env | awk -F= '/^(PMI|PMIX|MPI|OMPI|SLURM)_/{print $1}'); do + unset "$v" + done + export RAY_DEDUP_LOGS=0 + EXTRA+=(+actor_rollout_ref.rollout.quantization=fp8) +fi + +if [ -n "$MCORE_MODEL_PATH" ]; then + EXTRA+=( + actor_rollout_ref.actor.megatron.dist_checkpointing_path=${MCORE_MODEL_PATH} + actor_rollout_ref.actor.megatron.use_dist_checkpointing=True + actor_rollout_ref.ref.megatron.dist_checkpointing_path=${MCORE_MODEL_PATH} + actor_rollout_ref.ref.megatron.use_dist_checkpointing=True + ) +fi + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${ALGORITHM[@]}" \ + "${REWARD[@]}" \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" \ No newline at end of file diff --git a/verl/examples/grpo_trainer/run_qwen3_30b_veomni.sh b/verl/examples/grpo_trainer/run_qwen3_30b_veomni.sh new file mode 100644 index 0000000000000000000000000000000000000000..28c66490b2e35c8e78eb45f4a6fa1503d2b7e516 --- /dev/null +++ b/verl/examples/grpo_trainer/run_qwen3_30b_veomni.sh @@ -0,0 +1,134 @@ +#!/usr/bin/env bash +# GRPO | Qwen3-30B-A3B (MoE) | VeOmni training | NVIDIA GPUs or Ascend NPU +# Knobs: +# INFER_BACKEND controls rollout backend: vllm + +set -x +ENGINE=${1:-vllm} +DEVICE=${DEVICE:-$(python3 -c 'import torch_npu' 2>/dev/null && echo npu || echo gpu)} + +TRAIN_FILE=dapo-math-17k.parquet +TEST_FILE=aime-2024.parquet +max_prompt_length=$((1024 * 2)) +max_response_length=$((1024 * 8)) +rollout_max_num_seqs=$((128)) +n_devices_per_node=$((8)) +actor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 1)) +infer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3)) + +case "${DEVICE}" in + gpu) + ;; + npu) + export TASK_QUEUE_ENABLE=1 + export HCCL_OP_EXPANSION_MODE="AIV" + export VLLM_USE_V1=1 + export VLLM_VERSION=0.13.0 + export VLLM_ASCEND_ENABLE_NZ=0 + export HCCL_BUFFSIZE=610 + export CKPT_DIR="./ckpt30b" + export PYTORCH_NPU_ALLOC_CONF=max_split_size_mb:1024 + export CUDA_DEVICE_MAX_CONNECTIONS=1 + n_devices_per_node=16 + ;; + *) + echo "Unsupported DEVICE=${DEVICE}. Expected 'gpu' or 'npu'." >&2 + exit 1 + ;; +esac + +DATA=( + algorithm.use_kl_in_reward=False + algorithm.adv_estimator=grpo + data.train_files=${TRAIN_FILE} + data.val_files=${TEST_FILE} + data.train_batch_size=16 + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=False + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path=/Qwen3-30B-MoE-merge + actor_rollout_ref.model.enable_gradient_checkpointing=True + actor_rollout_ref.model.use_remove_padding=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=1e-6 + actor_rollout_ref.actor.veomni.param_offload=True + actor_rollout_ref.actor.veomni.optimizer_offload=True + actor_rollout_ref.actor.ppo_mini_batch_size=8 + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=0.001 + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=0 + actor_rollout_ref.actor.use_torch_compile=False + actor_rollout_ref.actor.veomni.fsdp_size=-1 + actor_rollout_ref.actor.veomni.expert_parallel_size=1 + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} +) + +ROLLOUT=( + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 + actor_rollout_ref.rollout.data_parallel_size=8 + actor_rollout_ref.rollout.expert_parallel_size=8 + actor_rollout_ref.rollout.tensor_model_parallel_size=1 + actor_rollout_ref.rollout.name=$ENGINE + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} + actor_rollout_ref.rollout.gpu_memory_utilization=0.7 + actor_rollout_ref.rollout.enable_chunked_prefill=True + actor_rollout_ref.rollout.enable_prefix_caching=True + actor_rollout_ref.rollout.max_model_len=$((max_prompt_length + max_response_length)) + actor_rollout_ref.rollout.max_num_batched_tokens=$((1024)) + actor_rollout_ref.rollout.max_num_seqs=${rollout_max_num_seqs} + actor_rollout_ref.rollout.enforce_eager=False + actor_rollout_ref.rollout.free_cache_engine=True + actor_rollout_ref.rollout.n=4 + +actor_rollout_ref.rollout.engine_kwargs.vllm.compilation_config.cudagraph_capture_sizes="[1, 8, 16, 32, 40, 48, 64, 96, 128, 256]" + +actor_rollout_ref.rollout.engine_kwargs.vllm.mm_processor_cache_gb=0 +) + +REF=( + actor_rollout_ref.ref.veomni.param_offload=True + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 + actor_rollout_ref.ref.veomni.param_offload=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} +) + +TRAINER=( + +trainer.use_legacy_worker_impl=disable + trainer.critic_warmup=0 + trainer.logger=console + trainer.project_name='verl_qwen3_veomni' + trainer.experiment_name='qwen3_30b_veomni' + trainer.n_gpus_per_node=${n_devices_per_node} + trainer.nnodes=1 + trainer.save_freq=100 + trainer.default_local_dir=$CKPT_DIR + trainer.test_freq=-1 + trainer.total_training_steps=100 +) + +if [ "${DEVICE}" = "npu" ]; then + ROLLOUT+=( + +actor_rollout_ref.rollout.engine_kwargs.vllm.compilation_config.cudagraph_mode="FULL_DECODE_ONLY" + ) +fi + +EXTRA=( + model_engine=veomni +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/grpo_trainer/run_qwen3_4b_fsdp.sh b/verl/examples/grpo_trainer/run_qwen3_4b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..4eb2fa17e6f34811f9827a27e6c546acaf285283 --- /dev/null +++ b/verl/examples/grpo_trainer/run_qwen3_4b_fsdp.sh @@ -0,0 +1,135 @@ +#!/usr/bin/env bash +# GRPO | Qwen3-4B | FSDP training | NVIDIA GPUs or Ascend NPUs +# +# INFER_BACKEND controls rollout backend: vllm + +set -xeuo pipefail + +# ---- user-adjustable ---- +INFER_BACKEND=${INFER_BACKEND:-vllm} +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-4B} +TRAIN_FILE=${TRAIN_FILE:-$HOME/data/gsm8k/train.parquet} +TEST_FILE=${TEST_FILE:-$HOME/data/gsm8k/test.parquet} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-512} +PPO_MINI_BATCH_SIZE=${PPO_MINI_BATCH_SIZE:-256} +PPO_MICRO_BATCH_SIZE_PER_GPU=${PPO_MICRO_BATCH_SIZE_PER_GPU:-2} +LOG_PROB_MICRO_BATCH_SIZE_PER_GPU=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU:-2} +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-1024} +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-1024} + +ACTOR_LR=${ACTOR_LR:-1e-6} +KL_LOSS_COEF=${KL_LOSS_COEF:-0.001} +ENTROPY_COEFF=${ENTROPY_COEFF:-0} + +ROLLOUT_TP=${ROLLOUT_TP:-2} +ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.6} +ROLLOUT_N=${ROLLOUT_N:-5} + +PROJECT_NAME=${PROJECT_NAME:-verl_grpo_example_geo3k} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-glm41v_9b_function_rm} +SAVE_FREQ=${SAVE_FREQ:-20} +TEST_FREQ=${TEST_FREQ:-5} +TOTAL_EPOCHS=${TOTAL_EPOCHS:-15} + +# ---- end user-adjustable ---- + +case "${DEVICE}" in + gpu) + ;; + npu) + export VLLM_USE_V1=1 + export TASK_QUEUE_ENABLE=2 + export CPU_AFFINITY_CONF=1 + export LD_PRELOAD="/usr/lib/aarch64-linux-gnu/libjemalloc.so.2${LD_PRELOAD:+:$LD_PRELOAD}" + NGPUS_PER_NODE=16 + ROLLOUT_GPU_MEM_UTIL=0.9 + ;; + *) + echo "Unsupported DEVICE=${DEVICE}. Expected 'gpu' or 'npu'." >&2 + exit 1 + ;; +esac + +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + data.train_files=${TRAIN_FILE} + data.val_files=${TEST_FILE} + data.train_batch_size=${TRAIN_BATCH_SIZE} + data.max_prompt_length=${MAX_PROMPT_LENGTH} + data.max_response_length=${MAX_RESPONSE_LENGTH} + data.filter_overlong_prompts=True + data.truncation='error' + algorithm.use_kl_in_reward=False +) + +MODEL=( + actor_rollout_ref.model.path=${MODEL_PATH} + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${ACTOR_LR} + actor_rollout_ref.actor.ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE} + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${PPO_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${KL_LOSS_COEF} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${ENTROPY_COEFF} + actor_rollout_ref.actor.fsdp_config.param_offload=False + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=3000 + actor_rollout_ref.actor.use_dynamic_bsz=True +) + +ROLLOUT=( + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.rollout.tensor_model_parallel_size=${ROLLOUT_TP} + actor_rollout_ref.rollout.name=${INFER_BACKEND} + actor_rollout_ref.rollout.gpu_memory_utilization=${ROLLOUT_GPU_MEM_UTIL} + actor_rollout_ref.rollout.enable_chunked_prefill=False + actor_rollout_ref.rollout.enforce_eager=False + actor_rollout_ref.rollout.free_cache_engine=True + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=4096 + actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=4096 + actor_rollout_ref.rollout.n=${ROLLOUT_N} +) + +REF=( + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.ref.fsdp_config.param_offload=True + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=8192 +) + +TRAINER=( + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${PROJECT_NAME} + trainer.experiment_name=${EXPERIMENT_NAME} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${SAVE_FREQ} + trainer.test_freq=${TEST_FREQ} + trainer.total_epochs=${TOTAL_EPOCHS} +) + +EXTRA=( +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/grpo_trainer/run_qwen3_5-122b-a10b_veomni.sh b/verl/examples/grpo_trainer/run_qwen3_5-122b-a10b_veomni.sh new file mode 100644 index 0000000000000000000000000000000000000000..21b92b552416cf81e841dce07e1009cff8d7aa15 --- /dev/null +++ b/verl/examples/grpo_trainer/run_qwen3_5-122b-a10b_veomni.sh @@ -0,0 +1,173 @@ +#!/usr/bin/env bash +# This script is a demo for GRPO training of Qwen3.5-122B-A10B using VeOmniEngine. +# +# Environment: +# - transformers==5.3.0 +# - sglang==0.5.9 +# - flash-linear-attention==0.4.1 +# - veomni==0.1.9a5 +# Tested configuration: +# - Model: Qwen3.5-122B-A10B +# - Sequence Parallel (SP): Tested with sp=1 and sp=2 (ulysses_parallel_size). +# - Expert Parallel: Tested with ep=8 (expert_parallel_size). + +set -xeuo pipefail + +data_path=${data_path:-$HOME/data/geo3k} +model_path=${model_path:-$HOME/model/Qwen3.5-122B-A10B} +output_path=${output_path:-$HOME/output} + +usp_size=${usp_size:-2} +expert_size=${expert_size:-8} +nnodes=${nnodes:-8} + +backend=fsdp2 +model_engine=veomni +project_name='verl_grpo_qwen3_5_122b_a10b_geo3k' +exp_name='qwen3_5_122b_a10b_veomni_sp2_ep8' +default_local_dir=$output_path/$project_name/$exp_name + + +# ===================================== Algorithm ===================================== +adv_estimator=grpo +loss_mode=gspo + +# reference policy +use_kl_in_reward=False +kl_coef=0.001 +use_kl_loss=False +kl_loss_coef=0.001 + +clip_ratio_low=3e-4 +clip_ratio_high=4e-4 + +actor_lr=1e-6 +critic_lr=2e-6 +gae_gamma=1.0 +gae_lam=0.95 +critic_warmup=0 + +# ===================================== Data/Model ===================================== +train_files=$data_path/train.parquet +test_files=$data_path/test.parquet + +actor_model_path=$model_path + +max_prompt_length=$((1024 * 1)) +max_response_length=$((1024 * 2)) + +train_batch_size=128 +ppo_mini_batch_size=32 +n_resp_per_prompt=8 +n_resp_per_prompt_val=1 + +use_remove_padding=True +use_dynamic_bsz=False + +# ===================================== Training ===================================== +actor_max_token_len_per_gpu=$(((max_prompt_length + max_response_length) * 8)) +ppo_micro_batch_size_per_gpu=1 + +# VeOmni config +ACTOR_VEOMNI_CONFIG=" + actor_rollout_ref.actor.veomni.param_offload=True \ + actor_rollout_ref.actor.veomni.optimizer_offload=True \ + actor_rollout_ref.actor.veomni.enable_full_shard=True \ + actor_rollout_ref.actor.veomni.ulysses_parallel_size=$usp_size \ + actor_rollout_ref.actor.veomni.expert_parallel_size=$expert_size \ + actor_rollout_ref.actor.veomni.attn_implementation=flash_attention_2 \ + actor_rollout_ref.actor.veomni.moe_implementation=fused_triton \ + actor_rollout_ref.actor.veomni.cross_entropy_loss_implementation=liger_kernel" + +# Actor model config +ACTOR_CONFIG=" + actor_rollout_ref.actor.optim.lr=$actor_lr \ + actor_rollout_ref.model.path=$actor_model_path \ + actor_rollout_ref.model.use_remove_padding=$use_remove_padding \ + actor_rollout_ref.actor.use_kl_loss=$use_kl_loss \ + actor_rollout_ref.actor.kl_loss_coef=$kl_loss_coef \ + actor_rollout_ref.actor.clip_ratio_low=$clip_ratio_low \ + actor_rollout_ref.actor.clip_ratio_high=$clip_ratio_high \ + actor_rollout_ref.actor.clip_ratio_c=10.0 \ + actor_rollout_ref.actor.policy_loss.loss_mode=${loss_mode} \ + actor_rollout_ref.actor.use_dynamic_bsz=$use_dynamic_bsz \ + actor_rollout_ref.actor.ppo_mini_batch_size=$ppo_mini_batch_size \ + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_max_token_len_per_gpu} \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=$ppo_micro_batch_size_per_gpu" + +CONFIG_NAME=ppo_trainer +ACTOR_CONFIG="$ACTOR_CONFIG $ACTOR_VEOMNI_CONFIG" + +# ===================================== Inference ===================================== +rollout_name=sglang +infer_tp=8 +infer_dp=1 +infer_ep=8 +gpu_memory_utilization=0.6 + +ROLLOUT_CONFIG=" + actor_rollout_ref.rollout.name=$rollout_name \ + actor_rollout_ref.rollout.mode=async \ + actor_rollout_ref.rollout.tensor_model_parallel_size=$infer_tp \ + actor_rollout_ref.rollout.data_parallel_size=$infer_dp \ + actor_rollout_ref.rollout.expert_parallel_size=$infer_ep \ + actor_rollout_ref.rollout.gpu_memory_utilization=$gpu_memory_utilization \ + actor_rollout_ref.rollout.n=$n_resp_per_prompt \ + actor_rollout_ref.rollout.val_kwargs.top_p=0.7 \ + actor_rollout_ref.rollout.val_kwargs.temperature=1.0 \ + actor_rollout_ref.rollout.val_kwargs.n=$n_resp_per_prompt_val \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=False" +########################### parameter arrays ########################### + +CONFIG=( + --config-path=./config + --config-name=$CONFIG_NAME +) + +DATA=( + algorithm.adv_estimator=$adv_estimator + algorithm.use_kl_in_reward=$use_kl_in_reward + algorithm.kl_ctrl.kl_coef=$kl_coef + algorithm.gamma=$gae_gamma + algorithm.lam=$gae_lam + data.train_files="$train_files" + data.val_files="$test_files" + data.return_raw_chat=True + data.train_batch_size=$train_batch_size + data.max_prompt_length=$max_prompt_length + data.max_response_length=$max_response_length + data.filter_overlong_prompts=True + data.filter_overlong_prompts_workers=64 + data.truncation='error' +) + +TRAINER=( + trainer.critic_warmup=$critic_warmup + trainer.logger=['console','wandb'] + trainer.project_name=$project_name + trainer.experiment_name=$exp_name + trainer.n_gpus_per_node=8 + trainer.nnodes=$nnodes + trainer.val_before_train=False + trainer.log_val_generations=100 + trainer.save_freq=-1 + trainer.test_freq=10 + trainer.total_epochs=10 + trainer.total_training_steps=500 +) + +EXTRA=( + model_engine=$model_engine + $ACTOR_CONFIG + $ROLLOUT_CONFIG +) + +########################### launch ########################### +python -m verl.trainer.main_ppo \ + "${CONFIG[@]}" \ + "${DATA[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" + diff --git a/verl/examples/grpo_trainer/run_qwen3_5_122b_a10b_megatron.sh b/verl/examples/grpo_trainer/run_qwen3_5_122b_a10b_megatron.sh new file mode 100644 index 0000000000000000000000000000000000000000..47279ffbf8b890b9d254a859496995a1ff120126 --- /dev/null +++ b/verl/examples/grpo_trainer/run_qwen3_5_122b_a10b_megatron.sh @@ -0,0 +1,194 @@ +#!/usr/bin/env bash +# Qwen3.5-122B-A10B MoE GRPO RL with Megatron (four nodes, 8 GPUs, H20, 96G, geo3k dataset) +# Using verlai/verl:vllm017.latest docker image +# Requirements: +# - 32 GPUs (96GB each, e.g. 4x8 H20) +# - Additional packages on top of the base image: +# pip install --upgrade transformers +# pip install flash-linear-attention +# pip install -U git+https://github.com/ISEEKYAN/mbridge.git +# - Megatron-LM==0.16.0 +# +# Qwen3.5 architecture notes: +# Qwen3.5 uses Gated Delta Net (GDN) linear attention which currently does +# NOT support packed sequences (THD format) in Megatron-LM. Therefore: +# - model.use_remove_padding=False (deprecated option, will be removed in the future forces bshd compute format) +# - actor.megatron.use_remove_padding=False (forces bshd compute format) +# - actor.use_dynamic_bsz=False (required for bshd mode) +# +# Once Megatron-LM adds THD support for Qwen3.5 GDN, use_remove_padding +# can be set to True for better performance. +# +# Tested parallelism config (32 GPUs / 4 node): +# TP=2 PP=2 CP=1 EP=8 ETP=1 GEN_TP=8 +# + +export CUDA_DEVICE_MAX_CONNECTIONS=1 +export VLLM_USE_V1=1 +export VLLM_ALLREDUCE_USE_SYMM_MEM=0 +export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 + +set -xeuo pipefail +unset http_proxy +unset https_proxy +# download geo3k dataset +hf download tyzhu/geo3k --repo-type dataset --local-dir $HOME/data/geo3k + +# ---- user-adjustable ---- +test_files=${test_files:-$HOME/data/geo3k/test.parquet} +train_files=${train_files:-$HOME/data/geo3k/train.parquet} +HF_MODEL_PATH=${HF_MODEL_PATH:-"Qwen/Qwen3.5-122B-A10B"} + +save_contents="['model', 'extra', 'optimizer']" + +project_name=${project_name:-'verl_grpo_qwen3_5_122b_geo3k'} +exp_name=${exp_name:-'qwen3_5_122b_megatron'} + +rollout_backend="vllm" + +save_path=${save_path:-"Qwen/Qwen3.5-122B/verl_checkpoint"} +save_freq=50 + +train_batch_size=128 +max_prompt_length=3240 +max_response_length=4096 +adv_estimator=${adv_estimator:-grpo} + +TP=${TP:-2} +PP=${PP:-2} +CP=${CP:-1} +EP=${EP:-8} +ETP=${ETP:-1} +GEN_TP=${GEN_TP:-8} +ACTOR_VPP=${ACTOR_VPP:-null} +ALL_OFFLOAD=${ALL_OFFLOAD:-True} + +NODE_GPU_NUM=${NODE_GPU_NUM:-8} +NODES_NUM=${NODES_NUM:-4} +# ---- end user-adjustable ---- + +# ---- no user adjustment needed below ---- +########################### Parameter Arrays ########################### + +ACTOR=( + actor_rollout_ref.actor.optim.lr=1e-6 + actor_rollout_ref.actor.ppo_mini_batch_size=64 + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=8192 + actor_rollout_ref.actor.use_dynamic_bsz=False + actor_rollout_ref.actor.use_kl_loss=False + actor_rollout_ref.actor.kl_loss_coef=0.01 + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=0 + actor_rollout_ref.actor.megatron.vanilla_mbridge=True + actor_rollout_ref.actor.megatron.use_mbridge=True + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${TP} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${PP} + actor_rollout_ref.actor.megatron.context_parallel_size=${CP} + actor_rollout_ref.actor.megatron.expert_model_parallel_size=${EP} + actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${ETP} + actor_rollout_ref.actor.megatron.param_offload=${ALL_OFFLOAD} + actor_rollout_ref.actor.megatron.optimizer_offload=${ALL_OFFLOAD} + actor_rollout_ref.actor.megatron.grad_offload=${ALL_OFFLOAD} + actor_rollout_ref.actor.megatron.dtype=bfloat16 + actor_rollout_ref.actor.megatron.virtual_pipeline_model_parallel_size=$ACTOR_VPP + actor_rollout_ref.actor.megatron.use_remove_padding=False + actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full + actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform + actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_load_balancing_type=\"none\" + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=True + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_grouped_gemm=True + +actor_rollout_ref.actor.megatron.override_transformer_config.apply_rope_fusion=False + +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=1 + +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True + +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=True + +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True + actor_rollout_ref.actor.use_torch_compile=True + actor_rollout_ref.actor.checkpoint.save_contents="${save_contents}" + +) + + +ROLLOUT=( + actor_rollout_ref.rollout.name=${rollout_backend} + actor_rollout_ref.rollout.tensor_model_parallel_size=${GEN_TP} + actor_rollout_ref.rollout.gpu_memory_utilization=0.66 + actor_rollout_ref.rollout.n=6 + actor_rollout_ref.rollout.dtype=bfloat16 + actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=4096 + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=False + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=8192 + +actor_rollout_ref.rollout.engine_kwargs.vllm.max_model_len=15768 +) + +REF=( + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=False + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=8192 + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${TP} + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${PP} + actor_rollout_ref.ref.megatron.context_parallel_size=${CP} + actor_rollout_ref.ref.megatron.expert_model_parallel_size=${EP} + actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=${ETP} + actor_rollout_ref.ref.megatron.param_offload=${ALL_OFFLOAD} +) + +MODEL=( + actor_rollout_ref.model.path=$HF_MODEL_PATH + actor_rollout_ref.model.enable_gradient_checkpointing=True + actor_rollout_ref.model.use_remove_padding=False +) + +ACTOR_ROLLOUT_REF_COMMON=( + actor_rollout_ref.nccl_timeout=10800 +) + +ALGORITHM=( + algorithm.adv_estimator=${adv_estimator} + algorithm.use_kl_in_reward=False +) + + +DATA=( + data.train_files=$train_files + data.val_files=$test_files + data.train_batch_size=$train_batch_size + data.max_prompt_length=$max_prompt_length + data.max_response_length=$max_response_length + data.truncation='right' + data.filter_overlong_prompts=True + data.filter_overlong_prompts_workers=64 +) + +TRAINER=( + trainer.logger=['console','wandb'] + trainer.project_name=$project_name + trainer.experiment_name=$exp_name + trainer.n_gpus_per_node=$NODE_GPU_NUM + trainer.nnodes=$NODES_NUM + trainer.save_freq=$save_freq + trainer.default_local_dir=${save_path} + trainer.test_freq=10 + trainer.val_before_train=False + trainer.total_epochs=20 +) + +EXTRA=( + model_engine=megatron +) + +########################### Launch ########################### +export HYDRA_FULL_ERROR=1 +PYTHONUNBUFFERED=1 python3 -m verl.trainer.main_ppo \ + "${ALGORITHM[@]}" \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR_ROLLOUT_REF_COMMON[@]}" \ + "${TRAINER[@]}" \ + "${ROLLOUT[@]}" \ + "${ACTOR[@]}" \ + "${REF[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/grpo_trainer/run_qwen3_5_27b_fsdp.sh b/verl/examples/grpo_trainer/run_qwen3_5_27b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..f8b06940796010a26d67394d279e0f69616a6be7 --- /dev/null +++ b/verl/examples/grpo_trainer/run_qwen3_5_27b_fsdp.sh @@ -0,0 +1,146 @@ +#!/usr/bin/env bash +# dependency: GPU vllm==0.18.0, transformers@ +# dependency: NPU vllm==0.18.0, vllm-ascend@<54879467>, transformers@ + +set -xeuo pipefail + +########################### user-adjustable ########################### +# DEVICE is auto-detected by probing torch_npu; override only for special cases. +DEVICE=${DEVICE:-$(python3 -c 'import torch_npu' 2>/dev/null && echo npu || echo gpu)} +INFER_BACKEND=${INFER_BACKEND:-vllm} +PROJECT_NAME=${PROJECT_NAME:-GRPO-Qwen3_5} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-GRPO-Qwen3_5-27B} +NDEVICES_PER_NODE=${NDEVICES_PER_NODE:-} +NNODES=${NNODES:-1} + +GEN_TP=${GEN_TP:-4} +SP_SIZE=${SP_SIZE:-1} +FSDP_SIZE=${FSDP_SIZE:-} +ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.6} + +RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"} +MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen3.5-27B"} +CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${PROJECT_NAME}/${EXPERIMENT_NAME}"} +TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/geo3k/train.parquet"} +TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/geo3k/test.parquet"} +WORKING_DIR=${WORKING_DIR:-"${PWD}"} +RUNTIME_ENV=${RUNTIME_ENV:-"${WORKING_DIR}/verl/trainer/runtime_env.yaml"} +########################### end user-adjustable ########################### + +########################### derived defaults ########################### +n_devices_per_node=${NDEVICES_PER_NODE:-8} +fsdp_size=${FSDP_SIZE:-8} + +case "${DEVICE}" in + gpu) + ;; + npu) + export HCCL_CONNECT_TIMEOUT=1500 + export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050 + export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050 + export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 + n_devices_per_node=16 + fsdp_size=16 + ;; + *) + echo "Unsupported DEVICE=${DEVICE}. Expected 'gpu' or 'npu'." >&2 + exit 1 + ;; +esac + +start_time=$(date +%Y%m%d)_$(date +%H%M%S) +mkdir -p logs + +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files="${TRAIN_FILE}" + data.val_files="${TEST_FILE}" + data.train_batch_size=64 + data.max_prompt_length=1024 + data.max_response_length=2048 + data.filter_overlong_prompts=True + data.truncation='error' + data.image_key=images + data.shuffle=False +) + +MODEL=( + actor_rollout_ref.model.path=${MODEL_PATH} + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=1e-6 + actor_rollout_ref.actor.ppo_mini_batch_size=16 + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.entropy_coeff=0 + actor_rollout_ref.actor.kl_loss_coef=0.01 + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.use_torch_compile=False + actor_rollout_ref.actor.strategy=fsdp2 + actor_rollout_ref.actor.use_dynamic_bsz=False + actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} + actor_rollout_ref.actor.fsdp_config.reshard_after_forward=True + actor_rollout_ref.actor.fsdp_config.entropy_checkpointing=True + actor_rollout_ref.actor.entropy_from_logits_with_chunking=True + actor_rollout_ref.actor.fsdp_config.offload_policy=True + actor_rollout_ref.actor.fsdp_config.ulysses_sequence_parallel_size=${SP_SIZE} + actor_rollout_ref.actor.fsdp_config.param_offload=True + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True +) + +REF=( + actor_rollout_ref.ref.strategy=fsdp2 + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 + actor_rollout_ref.ref.fsdp_config.param_offload=True + actor_rollout_ref.ref.fsdp_config.reshard_after_forward=True + actor_rollout_ref.ref.entropy_from_logits_with_chunking=True + actor_rollout_ref.ref.fsdp_config.ulysses_sequence_parallel_size=${SP_SIZE} + actor_rollout_ref.ref.use_torch_compile=False + actor_rollout_ref.ref.fsdp_config.offload_policy=True +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=${INFER_BACKEND} + actor_rollout_ref.rollout.ignore_eos=False + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 + actor_rollout_ref.rollout.tensor_model_parallel_size=${GEN_TP} + actor_rollout_ref.rollout.gpu_memory_utilization=${ROLLOUT_GPU_MEM_UTIL} + actor_rollout_ref.rollout.n=5 + actor_rollout_ref.rollout.enable_chunked_prefill=True + actor_rollout_ref.rollout.max_num_batched_tokens=8192 + actor_rollout_ref.rollout.free_cache_engine=True + actor_rollout_ref.rollout.enforce_eager=False + actor_rollout_ref.rollout.enable_prefix_caching=False + actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=6144 +) + +TRAINER=( + trainer.critic_warmup=0 + trainer.logger=['console','wandb'] + trainer.project_name="${PROJECT_NAME}" + trainer.experiment_name="${EXPERIMENT_NAME}" + trainer.n_gpus_per_node=${n_devices_per_node} + trainer.nnodes=${NNODES} + trainer.balance_batch=False + trainer.resume_from_path=checkpoints/ + trainer.val_before_train=True + trainer.save_freq=5 + trainer.test_freq=5 + trainer.total_epochs=15 +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${REF[@]}" \ + "${ROLLOUT[@]}" \ + "${TRAINER[@]}" \ + "$@" 2>&1 | tee logs/qwen3_5-27b-${start_time}.log diff --git a/verl/examples/grpo_trainer/run_qwen3_5_35b_a3b_veomni.sh b/verl/examples/grpo_trainer/run_qwen3_5_35b_a3b_veomni.sh new file mode 100644 index 0000000000000000000000000000000000000000..c549ec383c091400fe170a9092435976c76157df --- /dev/null +++ b/verl/examples/grpo_trainer/run_qwen3_5_35b_a3b_veomni.sh @@ -0,0 +1,167 @@ +#!/usr/bin/env bash +# This script is a demo for GRPO training of Qwen3.5-35B-A3B using VeOmniEngine. +# +# Environment: +# - transformers==5.3.0 +# - sglang==0.5.9 +# - flash-linear-attention==0.4.1 +# - veomni==0.1.9a1 +# Tested configuration: +# - Model: Qwen3.5-35B-A3B +# - Sequence Parallel (SP): Tested with sp=1 and sp=2 (ulysses_parallel_size). + +set -xeuo pipefail + +# ---- user-adjustable ---- +data_path=${data_path:-$HOME/data/geo3k} +model_path=${model_path:-$HOME/model/Qwen3.5-35B-A3B} +usp_size=${usp_size:-2} +nnodes=${nnodes:-2} + +backend=fsdp2 +model_engine=veomni +project_name='verl_grpo_qwen3_5_35b_a3b_geo3k' +exp_name='qwen3_5_35b_a3b_veomni_sp2' + + +# ===================================== Algorithm ===================================== +adv_estimator=grpo +loss_mode=gspo + +# reference policy +use_kl_in_reward=False +kl_coef=0.001 +use_kl_loss=False +kl_loss_coef=0.001 + +clip_ratio_low=3e-4 +clip_ratio_high=4e-4 + +actor_lr=1e-6 +critic_lr=2e-6 +gae_gamma=1.0 +gae_lam=0.95 +critic_warmup=0 + +# ===================================== Data/Model ===================================== +train_files=$data_path/train.parquet +test_files=$data_path/test.parquet + +actor_model_path=$model_path + +max_prompt_length=$((1024 * 1)) +max_response_length=$((1024 * 2)) + +train_batch_size=128 +ppo_mini_batch_size=32 +n_resp_per_prompt=8 +n_resp_per_prompt_val=1 + +use_remove_padding=True +use_dynamic_bsz=False + +# ===================================== Training ===================================== +actor_max_token_len_per_gpu=$(((max_prompt_length + max_response_length) * 8)) +ppo_micro_batch_size_per_gpu=1 + +# VeOmni config +ACTOR_VEOMNI_CONFIG=" + actor_rollout_ref.actor.veomni.param_offload=True \ + actor_rollout_ref.actor.veomni.optimizer_offload=True \ + actor_rollout_ref.actor.veomni.enable_full_shard=True \ + actor_rollout_ref.actor.veomni.ulysses_parallel_size=$usp_size \ + actor_rollout_ref.actor.veomni.expert_parallel_size=1 \ + actor_rollout_ref.actor.veomni.attn_implementation=flash_attention_2" +# ---- end user-adjustable ---- + +# ---- no user adjustment needed below ---- +# Actor model config +ACTOR_CONFIG=" + actor_rollout_ref.actor.optim.lr=$actor_lr \ + actor_rollout_ref.model.path=$actor_model_path \ + actor_rollout_ref.model.use_remove_padding=$use_remove_padding \ + actor_rollout_ref.actor.use_kl_loss=$use_kl_loss \ + actor_rollout_ref.actor.kl_loss_coef=$kl_loss_coef \ + actor_rollout_ref.actor.clip_ratio_low=$clip_ratio_low \ + actor_rollout_ref.actor.clip_ratio_high=$clip_ratio_high \ + actor_rollout_ref.actor.clip_ratio_c=10.0 \ + actor_rollout_ref.actor.policy_loss.loss_mode=${loss_mode} \ + actor_rollout_ref.actor.use_dynamic_bsz=$use_dynamic_bsz \ + actor_rollout_ref.actor.ppo_mini_batch_size=$ppo_mini_batch_size \ + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_max_token_len_per_gpu} \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=$ppo_micro_batch_size_per_gpu" + +CONFIG_NAME=ppo_trainer +ACTOR_CONFIG="$ACTOR_CONFIG $ACTOR_VEOMNI_CONFIG" + +# ===================================== Inference ===================================== +rollout_name=sglang +infer_tp=4 +infer_dp=1 +infer_ep=1 +gpu_memory_utilization=0.6 + +ROLLOUT_CONFIG=" + actor_rollout_ref.rollout.name=$rollout_name \ + actor_rollout_ref.rollout.tensor_model_parallel_size=$infer_tp \ + actor_rollout_ref.rollout.data_parallel_size=$infer_dp \ + actor_rollout_ref.rollout.expert_parallel_size=$infer_ep \ + actor_rollout_ref.rollout.gpu_memory_utilization=$gpu_memory_utilization \ + actor_rollout_ref.rollout.n=$n_resp_per_prompt \ + actor_rollout_ref.rollout.val_kwargs.top_p=0.7 \ + actor_rollout_ref.rollout.val_kwargs.temperature=1.0 \ + actor_rollout_ref.rollout.val_kwargs.n=$n_resp_per_prompt_val \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=False" +########################### parameter arrays ########################### + +CONFIG=( + --config-path=./config + --config-name=$CONFIG_NAME +) + +DATA=( + algorithm.adv_estimator=$adv_estimator + algorithm.use_kl_in_reward=$use_kl_in_reward + algorithm.kl_ctrl.kl_coef=$kl_coef + algorithm.gamma=$gae_gamma + algorithm.lam=$gae_lam + data.train_files="$train_files" + data.val_files="$test_files" + data.return_raw_chat=True + data.train_batch_size=$train_batch_size + data.max_prompt_length=$max_prompt_length + data.max_response_length=$max_response_length + data.filter_overlong_prompts=True + data.filter_overlong_prompts_workers=64 + data.truncation='error' +) + +TRAINER=( + trainer.critic_warmup=$critic_warmup + trainer.logger=['console','wandb'] + trainer.project_name=$project_name + trainer.experiment_name=$exp_name + trainer.n_gpus_per_node=8 + trainer.nnodes=$nnodes + trainer.val_before_train=False + trainer.log_val_generations=100 + trainer.save_freq=-1 + trainer.test_freq=10 + trainer.total_epochs=10 + trainer.total_training_steps=500 +) + +EXTRA=( + model_engine=$model_engine + $ACTOR_CONFIG + $ROLLOUT_CONFIG +) + +########################### launch ########################### +python -m verl.trainer.main_ppo \ + "${CONFIG[@]}" \ + "${DATA[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/grpo_trainer/run_qwen3_5_35b_fsdp.sh b/verl/examples/grpo_trainer/run_qwen3_5_35b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..6c38dcb3d49421673e16d3115bb2c225af61889e --- /dev/null +++ b/verl/examples/grpo_trainer/run_qwen3_5_35b_fsdp.sh @@ -0,0 +1,146 @@ +#!/usr/bin/env bash +# dependency: GPU vllm==0.18.0, transformers@ +# dependency: NPU vllm==0.18.0, vllm-ascend@<54879467>, transformers@ + +set -xeuo pipefail + +########################### user-adjustable ########################### +# DEVICE is auto-detected by probing torch_npu; override only for special cases. +DEVICE=${DEVICE:-$(python3 -c 'import torch_npu' 2>/dev/null && echo npu || echo gpu)} +INFER_BACKEND=${INFER_BACKEND:-vllm} +PROJECT_NAME=${PROJECT_NAME:-GRPO-Qwen3_5} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-GRPO-Qwen3_5-35B} +NDEVICES_PER_NODE=${NDEVICES_PER_NODE:-} +NNODES=${NNODES:-1} + +GEN_TP=${GEN_TP:-4} +SP_SIZE=${SP_SIZE:-1} +FSDP_SIZE=${FSDP_SIZE:-} +ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.6} + +RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"} +MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen3.5-35B-A3B"} +CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${PROJECT_NAME}/${EXPERIMENT_NAME}"} +TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/geo3k/train.parquet"} +TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/geo3k/test.parquet"} +WORKING_DIR=${WORKING_DIR:-"${PWD}"} +RUNTIME_ENV=${RUNTIME_ENV:-"${WORKING_DIR}/verl/trainer/runtime_env.yaml"} +########################### end user-adjustable ########################### + +########################### derived defaults ########################### +n_devices_per_node=${NDEVICES_PER_NODE:-8} +fsdp_size=${FSDP_SIZE:-8} + +case "${DEVICE}" in + gpu) + ;; + npu) + export HCCL_CONNECT_TIMEOUT=1500 + export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050 + export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050 + export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 + n_devices_per_node=16 + fsdp_size=16 + ;; + *) + echo "Unsupported DEVICE=${DEVICE}. Expected 'gpu' or 'npu'." >&2 + exit 1 + ;; +esac + +start_time=$(date +%Y%m%d)_$(date +%H%M%S) +mkdir -p logs + +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files="${TRAIN_FILE}" + data.val_files="${TEST_FILE}" + data.train_batch_size=64 + data.max_prompt_length=1024 + data.max_response_length=2048 + data.filter_overlong_prompts=True + data.truncation='error' + data.image_key=images + data.shuffle=False +) + +MODEL=( + actor_rollout_ref.model.path=${MODEL_PATH} + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=1e-6 + actor_rollout_ref.actor.ppo_mini_batch_size=16 + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.entropy_coeff=0 + actor_rollout_ref.actor.kl_loss_coef=0.01 + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.use_torch_compile=False + actor_rollout_ref.actor.strategy=fsdp2 + actor_rollout_ref.actor.use_dynamic_bsz=False + actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} + actor_rollout_ref.actor.fsdp_config.reshard_after_forward=True + actor_rollout_ref.actor.fsdp_config.entropy_checkpointing=True + actor_rollout_ref.actor.entropy_from_logits_with_chunking=True + actor_rollout_ref.actor.fsdp_config.offload_policy=True + actor_rollout_ref.actor.fsdp_config.ulysses_sequence_parallel_size=${SP_SIZE} + actor_rollout_ref.actor.fsdp_config.param_offload=True + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True +) + +REF=( + actor_rollout_ref.ref.strategy=fsdp2 + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 + actor_rollout_ref.ref.fsdp_config.param_offload=True + actor_rollout_ref.ref.fsdp_config.reshard_after_forward=True + actor_rollout_ref.ref.entropy_from_logits_with_chunking=True + actor_rollout_ref.ref.fsdp_config.ulysses_sequence_parallel_size=${SP_SIZE} + actor_rollout_ref.ref.use_torch_compile=False + actor_rollout_ref.ref.fsdp_config.offload_policy=True +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=${INFER_BACKEND} + actor_rollout_ref.rollout.ignore_eos=False + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 + actor_rollout_ref.rollout.tensor_model_parallel_size=${GEN_TP} + actor_rollout_ref.rollout.gpu_memory_utilization=${ROLLOUT_GPU_MEM_UTIL} + actor_rollout_ref.rollout.n=5 + actor_rollout_ref.rollout.enable_chunked_prefill=True + actor_rollout_ref.rollout.max_num_batched_tokens=8192 + actor_rollout_ref.rollout.free_cache_engine=True + actor_rollout_ref.rollout.enforce_eager=False + actor_rollout_ref.rollout.enable_prefix_caching=False + actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=6144 +) + +TRAINER=( + trainer.critic_warmup=0 + trainer.logger=['console','wandb'] + trainer.project_name="${PROJECT_NAME}" + trainer.experiment_name="${EXPERIMENT_NAME}" + trainer.n_gpus_per_node=${n_devices_per_node} + trainer.nnodes=${NNODES} + trainer.balance_batch=False + trainer.resume_from_path=checkpoints/ + trainer.val_before_train=True + trainer.save_freq=5 + trainer.test_freq=5 + trainer.total_epochs=15 +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${REF[@]}" \ + "${ROLLOUT[@]}" \ + "${TRAINER[@]}" \ + "$@" 2>&1 | tee logs/qwen3_5-35b-${start_time}.log diff --git a/verl/examples/grpo_trainer/run_qwen3_5_35b_megatron.sh b/verl/examples/grpo_trainer/run_qwen3_5_35b_megatron.sh new file mode 100644 index 0000000000000000000000000000000000000000..8c8131c2cb2bfe13d69e9685ed45021593590569 --- /dev/null +++ b/verl/examples/grpo_trainer/run_qwen3_5_35b_megatron.sh @@ -0,0 +1,170 @@ +#!/usr/bin/env bash +# Qwen3.5-35B-A3B MoE GRPO RL with Megatron (single node, 8 GPUs, geo3k dataset) +# +# notes on vllm: +# by 20260225, the latest vllm nightly does not support qwen3.5 rollout, to use this script, you need to +# 1. wait until vllm supports qwen3.5 officially, and build a verl docker with that version of vllm +# 2. self build a verl docker image with vllm from source code with qwen3.5 support (main branch 20260225 is OK) +# I succeeded in running this script with the main branch of vllm on 20260225, yet there are still some minor issues +# the vllm qwen3.5 during initialization, need to be fixed. Also, the cuda_graph is somehow not working, need to be +# fixed, either by verl team with supoorts to vllm0.16, or by vllm team. +# Requirements: +# - 8 GPUs (80GB each, e.g. 1x8 H100/H200) +# - Additional packages on top of the base image: +# pip install --upgrade transformers +# pip install flash-linear-attention +# pip install -U git+https://github.com/ISEEKYAN/mbridge.git +# - Megatron-LM==0.16.0 +# +# Qwen3.5 architecture notes: +# Qwen3.5 uses Gated Delta Net (GDN) linear attention which currently does +# NOT support packed sequences (THD format) in Megatron-LM. Therefore: +# - model.use_remove_padding=False (deprecated option, will be removed in the future forces bshd compute format) +# - actor.megatron.use_remove_padding=False (forces bshd compute format) +# - actor.use_dynamic_bsz=False (required for bshd mode) +# +# Once Megatron-LM adds THD support for Qwen3.5 GDN, use_remove_padding +# can be set to True for better performance. +# +# Tested parallelism config (8 GPUs / 1 node): +# TP=2 PP=1 CP=1 EP=8 ETP=1 GEN_TP=8 +# + +export CUDA_DEVICE_MAX_CONNECTIONS=1 +export VLLM_USE_V1=1 +export VLLM_ALLREDUCE_USE_SYMM_MEM=0 + +set -xeuo pipefail + +########################### Quick Config ########################### + +# ---- user-adjustable ---- +TP=${TP:-2} +PP=${PP:-1} +CP=${CP:-1} +EP=${EP:-8} +ETP=${ETP:-1} +GEN_TP=${GEN_TP:-8} + +ALL_OFFLOAD=${ALL_OFFLOAD:-True} + +rollout_name="vllm" +project_name='verl_grpo_qwen3_5_35b_geo3k' +exp_name='qwen3_5_35b_megatron' +adv_estimator=grpo + +HF_MODEL_PATH=${HF_MODEL_PATH:-"Qwen3.5-35B-A3B"} +train_path=${train_path:-$HOME/data/geo3k/train.parquet} +test_path=${test_path:-$HOME/data/geo3k/test.parquet} +# ---- end user-adjustable ---- + +# ---- no user adjustment needed below ---- +########################### Parameter Arrays ########################### + +DATA=( + data.train_files=${train_path} + data.val_files=${test_path} + data.train_batch_size=32 + data.max_prompt_length=1024 + data.max_response_length=2048 + data.truncation='error' + data.filter_overlong_prompts=True +) + +MODEL=( + actor_rollout_ref.model.path=${HF_MODEL_PATH} + actor_rollout_ref.model.trust_remote_code=True + actor_rollout_ref.model.use_remove_padding=False +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=1e-6 + actor_rollout_ref.actor.ppo_mini_batch_size=32 + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=4096 + actor_rollout_ref.actor.use_dynamic_bsz=False + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=0.01 + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=0 + actor_rollout_ref.actor.megatron.use_mbridge=True + actor_rollout_ref.actor.megatron.vanilla_mbridge=True + actor_rollout_ref.actor.megatron.use_remove_padding=False + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${TP} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${PP} + actor_rollout_ref.actor.megatron.context_parallel_size=${CP} + actor_rollout_ref.actor.megatron.expert_model_parallel_size=${EP} + actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${ETP} + actor_rollout_ref.actor.megatron.param_offload=${ALL_OFFLOAD} + actor_rollout_ref.actor.megatron.optimizer_offload=${ALL_OFFLOAD} + actor_rollout_ref.actor.megatron.grad_offload=${ALL_OFFLOAD} + actor_rollout_ref.actor.megatron.dtype=bfloat16 + ++actor_rollout_ref.actor.megatron.override_transformer_config.attention_backend=auto + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_aux_loss_coeff=0.01 + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_z_loss_coeff=0.001 + +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=1 + +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True + +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=True + +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=${rollout_name} + actor_rollout_ref.rollout.tensor_model_parallel_size=${GEN_TP} + actor_rollout_ref.rollout.gpu_memory_utilization=0.6 + actor_rollout_ref.rollout.n=5 + actor_rollout_ref.rollout.dtype=bfloat16 + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=False + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=4096 +) + +REF=( + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=False + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=4096 + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${TP} + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${PP} + actor_rollout_ref.ref.megatron.context_parallel_size=${CP} + actor_rollout_ref.ref.megatron.expert_model_parallel_size=${EP} + actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=${ETP} + actor_rollout_ref.ref.megatron.param_offload=${ALL_OFFLOAD} +) + +ALGORITHM=( + algorithm.adv_estimator=${adv_estimator} + algorithm.use_kl_in_reward=False +) + +TRAINER=( + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${exp_name} + trainer.n_gpus_per_node=8 + trainer.nnodes=1 + trainer.save_freq=20 + trainer.val_before_train=False + trainer.test_freq=5 + trainer.total_epochs=15 +) + +EXTRA=( + model_engine=megatron +) + +########################### Launch ########################### + +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${ALGORITHM[@]}" \ + "${MODEL[@]}" \ + "${ROLLOUT[@]}" \ + "${ACTOR[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/grpo_trainer/run_qwen3_8b_fsdp.sh b/verl/examples/grpo_trainer/run_qwen3_8b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..fc0683c9f786e3bc272dd8ae867c49527bb5331b --- /dev/null +++ b/verl/examples/grpo_trainer/run_qwen3_8b_fsdp.sh @@ -0,0 +1,197 @@ +#!/usr/bin/env bash +# GRPO | Qwen3-8B | FSDP training | NVIDIA GPUs or Ascend NPUs +# +# Knobs: +# INFER_BACKEND rollout backend: vllm | sglang | trtllm (default: vllm) +# MACHINE free-form tag for hardware tweaks (e.g. gb200) (default: unset) +# (DEVICE is auto-detected from torch_npu; export DEVICE=gpu|npu only to override.) +# +# TensorRT-LLM is GPU-only. +# `MACHINE=gb200` (Blackwell SM100) bundles: enforce_eager=True, FSDP +# model_dtype=bfloat16, SGLang attention_backend=flashinfer (FA3 unsupported on +# SM>90), and ray_init.num_gpus pinned (Docker --privileged bypasses GPU +# autodetect). Unknown MACHINE values are accepted and only affect experiment_name. + +set -xeuo pipefail + +########################### user-adjustable ########################### +# DEVICE is auto-detected by probing torch_npu; override only for special cases. +DEVICE=${DEVICE:-$(python3 -c 'import torch_npu' 2>/dev/null && echo npu || echo gpu)} +INFER_BACKEND=${INFER_BACKEND:-vllm} +MACHINE=${MACHINE:-} + +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-8B} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-} +NPUS_PER_NODE=${NPUS_PER_NODE:-} + +train_batch_size=${TRAIN_BATCH_SIZE:-1024} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-256} +max_prompt_length=${MAX_PROMPT_LENGTH:-1024} +max_response_length=${MAX_RESPONSE_LENGTH:-2048} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-24576} + +actor_lr=${ACTOR_LR:-1e-6} +kl_loss_coef=${KL_LOSS_COEF:-0.001} +entropy_coeff=${ENTROPY_COEFF:-0} + +rollout_tp=${ROLLOUT_TP:-} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-} +rollout_n=${ROLLOUT_N:-5} +sp_size=${SP_SIZE:-1} + +total_epochs=${TOTAL_EPOCHS:-15} +save_freq=${SAVE_FREQ:-20} +test_freq=${TEST_FREQ:-5} + +PROJECT_NAME=${PROJECT_NAME:-verl_grpo_gsm8k_math} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-qwen3_8b_grpo_${INFER_BACKEND}_fsdp_$(date +%Y%m%d_%H%M)} +########################### end user-adjustable ########################### + +########################### derived defaults ########################### +case "${DEVICE}" in + gpu | npu) ;; + *) + echo "DEVICE must be gpu or npu, got: ${DEVICE}" >&2 + exit 1 + ;; +esac + +if [ "${DEVICE}" = npu ] && [ "${INFER_BACKEND}" = trtllm ]; then + echo "INFER_BACKEND=trtllm is only supported with DEVICE=gpu" >&2 + exit 1 +fi + +# Defaults and extras grouped by device. Backend / machine refinements stay +# nested in the device branch they apply to. +EXTRA=() +case "${DEVICE}" in + gpu) + actor_param_offload=False + actor_optimizer_offload=False + rollout_tp=${rollout_tp:-2} + rollout_gpu_mem_util=${rollout_gpu_mem_util:-0.6} + + case "${MACHINE}" in + gb200) + NGPUS_PER_NODE=${NGPUS_PER_NODE:-4} + # Blackwell SM100: see header comment for rationale of each override. + EXTRA+=( + actor_rollout_ref.rollout.enforce_eager=True + actor_rollout_ref.rollout.free_cache_engine=True + actor_rollout_ref.actor.fsdp_config.model_dtype=bfloat16 + "+ray_kwargs.ray_init.num_gpus=${NGPUS_PER_NODE}" + ) + if [ "${INFER_BACKEND}" = sglang ]; then + EXTRA+=(+actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend=flashinfer) + fi + ;; + *) + NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + ;; + esac + n_trainer_devices=${NGPUS_PER_NODE} + ;; + npu) + export HCCL_CONNECT_TIMEOUT=1500 + export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050 + export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050 + export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 + + NPUS_PER_NODE=8 + n_trainer_devices=${NPUS_PER_NODE} + actor_param_offload=True + actor_optimizer_offload=True + rollout_tp=${rollout_tp:-4} + sp_size=4 + train_batch_size=16 + max_prompt_length=$((1024 * 2)) + max_response_length=$((1024 * 32)) + ppo_mini_batch_size=16 + rollout_gpu_mem_util=0.3 + EXTRA+=( + actor_rollout_ref.actor.use_torch_compile=False + actor_rollout_ref.actor.fsdp_config.ulysses_sequence_parallel_size=${sp_size} + actor_rollout_ref.ref.fsdp_config.ulysses_sequence_parallel_size=${sp_size} + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 + actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=4096 + ) + if [ "${INFER_BACKEND}" = sglang ]; then + EXTRA+=( + +actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend=ascend + ) + fi + ;; +esac + +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files="['$HOME/data/gsm8k/train.parquet', '$HOME/data/math/train.parquet']" + data.val_files="['$HOME/data/gsm8k/test.parquet', '$HOME/data/math/test.parquet']" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.fsdp_config.param_offload=${actor_param_offload} + actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_optimizer_offload} +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=${INFER_BACKEND} + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.fsdp_config.param_offload=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.logger='["console","wandb"]' + trainer.project_name=${PROJECT_NAME} + trainer.experiment_name=${EXPERIMENT_NAME} + trainer.n_gpus_per_node=${n_trainer_devices} + trainer.nnodes=${NNODES} + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/grpo_trainer/run_qwen3_8b_megatron.sh b/verl/examples/grpo_trainer/run_qwen3_8b_megatron.sh new file mode 100644 index 0000000000000000000000000000000000000000..bfe67f180fb15652de9cbc9ee7052025167604a7 --- /dev/null +++ b/verl/examples/grpo_trainer/run_qwen3_8b_megatron.sh @@ -0,0 +1,114 @@ +#!/usr/bin/env bash +# GRPO | Qwen3-8B | Megatron training | NVIDIA GPUs +# +# INFER_BACKEND controls rollout backend: vllm | sglang | trtllm. + +set -xeuo pipefail +export CUDA_DEVICE_MAX_CONNECTIONS=1 + +########################### user-adjustable ########################### +INFER_BACKEND=${INFER_BACKEND:-vllm} + +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-8B} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +train_batch_size=${TRAIN_BATCH_SIZE:-1024} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-256} +max_prompt_length=${MAX_PROMPT_LENGTH:-1024} +max_response_length=${MAX_RESPONSE_LENGTH:-2048} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-24576} + +actor_lr=${ACTOR_LR:-1e-6} +kl_loss_coef=${KL_LOSS_COEF:-0.001} +entropy_coeff=${ENTROPY_COEFF:-0} + +actor_tp=${ACTOR_TP:-2} +actor_pp=${ACTOR_PP:-2} + +rollout_tp=${ROLLOUT_TP:-2} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.6} +rollout_n=${ROLLOUT_N:-5} + +total_epochs=${TOTAL_EPOCHS:-15} +save_freq=${SAVE_FREQ:-20} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_grpo_gsm8k_math} +experiment_name=${EXPERIMENT_NAME:-qwen3_8b_${INFER_BACKEND}_megatron} +########################### end user-adjustable ########################### + +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files="['$HOME/data/gsm8k/train.parquet', '$HOME/data/math/train.parquet']" + data.val_files="['$HOME/data/gsm8k/test.parquet', '$HOME/data/math/test.parquet']" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${actor_pp} +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=${INFER_BACKEND} + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${actor_pp} +) + +TRAINER=( + trainer.balance_batch=True + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( + model_engine=megatron +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/grpo_trainer/run_qwen3_8b_mindspeed.sh b/verl/examples/grpo_trainer/run_qwen3_8b_mindspeed.sh new file mode 100644 index 0000000000000000000000000000000000000000..3e5fb27ffa08f5c15e3b11d73327ae41a53495c8 --- /dev/null +++ b/verl/examples/grpo_trainer/run_qwen3_8b_mindspeed.sh @@ -0,0 +1,170 @@ +#!/usr/bin/env bash +# GRPO | text | MindSpeed-LLM training | Ascend NPU +# +# Set INFER_BACKEND=sglang (default). + +set -xeuo pipefail + +# ---- user-adjustable ---- +INFER_BACKEND=${INFER_BACKEND:-sglang} + +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-8B} +NNODES=${NNODES:-1} +NPUS_PER_NODE=${NPUS_PER_NODE:-8} + +train_batch_size=${TRAIN_BATCH_SIZE:-16} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-16} +max_prompt_length=${MAX_PROMPT_LENGTH:-2048} +max_response_length=${MAX_RESPONSE_LENGTH:-2048} +micro_bsz=${MICRO_BSZ:-1} + +actor_lr=${ACTOR_LR:-1e-6} +kl_loss_coef=${KL_LOSS_COEF:-0.001} +entropy_coeff=${ENTROPY_COEFF:-0} + +actor_tp=${ACTOR_TP:-4} +actor_pp=${ACTOR_PP:-4} +actor_cp=${ACTOR_CP:-1} +all_offload=${ALL_OFFLOAD:-True} + +rollout_tp=${ROLLOUT_TP:-4} +rollout_dp=${ROLLOUT_DP:-1} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.5} +rollout_n=${ROLLOUT_N:-8} + +total_epochs=${TOTAL_EPOCHS:-15} +save_freq=${SAVE_FREQ:--1} +test_freq=${TEST_FREQ:--1} + +project_name=${PROJECT_NAME:-verl_grpo_gsm8k_math} +experiment_name=${EXPERIMENT_NAME:-qwen3_8b_mindspeed} +CKPTS_DIR=${CKPTS_DIR:-"${HOME}/verl/ckpts/${project_name}/${experiment_name}"} +# ---- end user-adjustable ---- + +# ---- system defaults (normally leave as-is) ---- +export HCCL_CONNECT_TIMEOUT=1500 +export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050 +export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050 +export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 +export DISABLE_L2_CACHE=1 +export TASK_QUEUE_ENABLE=1 +# For CANN 8.5.0+, when using mbridge: +export HCCL_OP_EXPANSION_MODE=AIV +# ---- end system defaults ---- + +gsm8k_train=$HOME/data/gsm8k/train.parquet +gsm8k_test=$HOME/data/gsm8k/test.parquet +math_train=$HOME/data/math/train.parquet +math_test=$HOME/data/math/test.parquet + +train_files="['$gsm8k_train', '$math_train']" +val_files="['$gsm8k_test', '$math_test']" + +max_model_len=$((max_prompt_length + max_response_length)) +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + algorithm.kl_ctrl.kl_coef=0.0 + data.train_files="$train_files" + data.val_files="$val_files" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=False + data.truncation=left +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True +) + +ACTOR=( + actor_rollout_ref.actor.use_torch_compile=False + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${max_model_len} + actor_rollout_ref.actor.ppo_epochs=1 + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.mindspeed.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.actor.mindspeed.pipeline_model_parallel_size=${actor_pp} + actor_rollout_ref.actor.mindspeed.context_parallel_size=${actor_cp} + actor_rollout_ref.actor.mindspeed.param_offload=${all_offload} + actor_rollout_ref.actor.mindspeed.optimizer_offload=${all_offload} + actor_rollout_ref.actor.mindspeed.grad_offload=${all_offload} + actor_rollout_ref.actor.mindspeed.use_mbridge=True + actor_rollout_ref.actor.mindspeed.vanilla_mbridge=True + actor_rollout_ref.actor.mindspeed.llm_kwargs.spec='[mindspeed_llm.tasks.models.spec.qwen3_spec, layer_spec]' + actor_rollout_ref.actor.mindspeed.llm_kwargs.seq_length=${max_model_len} + actor_rollout_ref.actor.mindspeed.llm_kwargs.micro_batch_size=${micro_bsz} + +actor_rollout_ref.actor.mindspeed.llm_kwargs.num_query_groups=8 + +actor_rollout_ref.actor.mindspeed.llm_kwargs.recompute_method=uniform + +actor_rollout_ref.actor.mindspeed.llm_kwargs.recompute_granularity=full + +actor_rollout_ref.actor.mindspeed.llm_kwargs.recompute_num_layers=1 + +actor_rollout_ref.actor.mindspeed.llm_kwargs.overlap_grad_reduce=True + +actor_rollout_ref.actor.mindspeed.llm_kwargs.overlap_param_gather=True +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=${INFER_BACKEND} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${max_model_len} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.data_parallel_size=${rollout_dp} + actor_rollout_ref.rollout.enforce_eager=False +) + +REF=( + actor_rollout_ref.ref.use_torch_compile=False + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${max_model_len} + actor_rollout_ref.ref.mindspeed.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.ref.mindspeed.pipeline_model_parallel_size=${actor_pp} + actor_rollout_ref.ref.mindspeed.context_parallel_size=${actor_cp} + actor_rollout_ref.ref.mindspeed.param_offload=${all_offload} + actor_rollout_ref.ref.mindspeed.use_mbridge=True + actor_rollout_ref.ref.mindspeed.vanilla_mbridge=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.nnodes=${NNODES} + trainer.n_gpus_per_node=${NPUS_PER_NODE} + trainer.val_before_train=False + trainer.test_freq=${test_freq} + trainer.save_freq=${save_freq} + trainer.total_epochs=${total_epochs} + trainer.default_local_dir="${CKPTS_DIR}" +) + +EXTRA=( + model_engine=mindspeed +) + +if [ "${INFER_BACKEND}" = sglang ]; then + EXTRA+=( + +actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend=ascend + +actor_rollout_ref.rollout.engine_kwargs.sglang.chunked_prefill_size=-1 + ) +fi + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/grpo_trainer/run_qwen3_next_80b_a3b_fsdp.sh b/verl/examples/grpo_trainer/run_qwen3_next_80b_a3b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..01465296f04cfd0a2501e2e50a47cc803532948a --- /dev/null +++ b/verl/examples/grpo_trainer/run_qwen3_next_80b_a3b_fsdp.sh @@ -0,0 +1,190 @@ +#!/usr/bin/env bash +set -xeuo pipefail + +# ---- user-adjustable ---- +# DEVICE is auto-detected by probing torch_npu; override only for special cases. +DEVICE=${DEVICE:-$(python3 -c 'import torch_npu' 2>/dev/null && echo npu || echo gpu)} +PROJECT_NAME=${PROJECT_NAME:-verl_grpo_qwen3-next-80b} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-qwen3_next_80b_a3b_grpo_vllm_fsdp_$(date +%Y%m%d_%H%M)} + +# Paths +WORK_DIR=${WORK_DIR:-"${HOME}/verl"} +MODEL_PATH=${MODEL_PATH:-"${WORK_DIR}/Qwen3-Next-80B-A3B-Instruct"} +TRAIN_FILE=${TRAIN_FILE:-"${WORK_DIR}/datasets/dapo-math-17k/dapo-math-17k.parquet"} +TEST_FILE=${TEST_FILE:-"${WORK_DIR}/datasets/aime/aime-2024.parquet"} +NNODES=${NNODES:-4} +NDEVICES_PER_NODE=${NDEVICES_PER_NODE:-} + +# algorithm +adv_estimator=${ADV_ESTIMATOR:-grpo} +use_kl_in_reward=${USE_KL_IN_REWARD:-False} +kl_coef=${KL_COEF:-0.0} +use_kl_loss=${USE_KL_LOSS:-True} +kl_loss_coef=${KL_LOSS_COEF:-0.001} +clip_ratio_low=${CLIP_RATIO_LOW:-0.2} +clip_ratio_high=${CLIP_RATIO_HIGH:-0.28} + +temperature=${TEMPERATURE:-1.0} +top_p=${TOP_P:-1.0} +top_k=${TOP_K:--1} # 0 for HF rollout, -1 for vLLM rollout +val_top_p=${VAL_TOP_P:-0.7} + +# batch +train_batch_size=${TRAIN_BATCH_SIZE:-16} +rollout_n=${ROLLOUT_N:-16} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-8} + +# length +max_prompt_length=${MAX_PROMPT_LENGTH:-$((1024 * 2))} +max_response_length=${MAX_RESPONSE_LENGTH:-$((1024 * 20))} + +# optimizer +learning_rate=${ACTOR_LR:-1e-6} +warmup_steps=${WARMUP_STEPS:-0} + +# performance +sp_size=${SP_SIZE:-8} +gen_tp=${ROLLOUT_TP:-4} +use_dynamic_bsz=${USE_DYNAMIC_BSZ:-True} +offload=${OFFLOAD:-True} +# ---- end user-adjustable ---- + +# ---- no user adjustment needed below ---- +n_devices_per_node=${NDEVICES_PER_NODE:-8} + +case "${DEVICE}" in + gpu) + ;; + npu) + ;; + *) + echo "Unsupported DEVICE=${DEVICE}. Expected 'gpu' or 'npu'." >&2 + exit 1 + ;; +esac + +actor_ppo_max_token_len=$(((max_prompt_length + max_response_length) / sp_size)) +infer_ppo_max_token_len=$(((max_prompt_length + max_response_length) / sp_size)) + +DATA=( + data.train_files="${TRAIN_FILE}" + data.val_files="${TEST_FILE}" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.truncation='error' +) + +ACTOR=( + actor_rollout_ref.actor.strategy=fsdp2 + actor_rollout_ref.nccl_timeout=14400 + + # fsdp + actor_rollout_ref.actor.fsdp_config.use_orig_params=True + actor_rollout_ref.actor.fsdp_config.model_dtype=bfloat16 + actor_rollout_ref.actor.fsdp_config.param_offload=${offload} + actor_rollout_ref.actor.fsdp_config.optimizer_offload=${offload} + actor_rollout_ref.actor.fsdp_config.forward_prefetch=False + actor_rollout_ref.actor.fsdp_config.fsdp_size=-1 + +actor_rollout_ref.actor.fsdp_config.mixed_precision.reduce_dtype=bf16 + + # optimizer + actor_rollout_ref.actor.optim.lr=${learning_rate} + actor_rollout_ref.actor.optim.lr_warmup_steps=${warmup_steps} + + # ppo config + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} + actor_rollout_ref.actor.fsdp_config.ulysses_sequence_parallel_size=${sp_size} + + # entropy + actor_rollout_ref.actor.entropy_checkpointing=True + actor_rollout_ref.actor.entropy_from_logits_with_chunking=True + + actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} + actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} + actor_rollout_ref.actor.clip_ratio_c=10.0 + + actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} + actor_rollout_ref.actor.use_torch_compile=False +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=0.8 + actor_rollout_ref.rollout.load_format=auto + actor_rollout_ref.rollout.enforce_eager=True + actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) + actor_rollout_ref.rollout.calculate_log_probs=True + + actor_rollout_ref.rollout.temperature=${temperature} + actor_rollout_ref.rollout.top_p=${top_p} + actor_rollout_ref.rollout.top_k=${top_k} + actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} + actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} + actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} + actor_rollout_ref.rollout.val_kwargs.do_sample=True + actor_rollout_ref.rollout.val_kwargs.n=1 + + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} +) + +REF=( + actor_rollout_ref.ref.fsdp_config.ulysses_sequence_parallel_size=${sp_size} + actor_rollout_ref.ref.use_torch_compile=False + actor_rollout_ref.ref.fsdp_config.param_offload=${offload} + actor_rollout_ref.ref.fsdp_config.optimizer_offload=${offload} + actor_rollout_ref.ref.fsdp_config.forward_prefetch=False + + actor_rollout_ref.ref.entropy_checkpointing=True + actor_rollout_ref.ref.entropy_from_logits_with_chunking=True + + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} +) + +TRAINER=( + trainer.logger='["console","wandb"]' + trainer.project_name="${PROJECT_NAME}" + trainer.experiment_name="${EXPERIMENT_NAME}" + trainer.n_gpus_per_node=${n_devices_per_node} + trainer.nnodes=${NNODES} + trainer.val_before_train=False + trainer.save_freq=5 + trainer.test_freq=-1 + trainer.total_epochs=1 +) + +MODEL=( + actor_rollout_ref.model.path=${MODEL_PATH} + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_activation_offload=${offload} +) + +ALGORITHM=( + algorithm.adv_estimator=${adv_estimator} + algorithm.use_kl_in_reward=${use_kl_in_reward} + algorithm.kl_ctrl.kl_coef=${kl_coef} +) + +echo "Starting Training with:" +echo "Project: ${PROJECT_NAME}, Exp: ${EXPERIMENT_NAME}" +echo "Rollout N: ${rollout_n}, Batch Size: ${train_batch_size}, LR: ${learning_rate}" + +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${ALGORITHM[@]}" \ + "${MODEL[@]}" diff --git a/verl/examples/grpo_trainer/run_qwen3_vl_235b_a22b_megatron.sh b/verl/examples/grpo_trainer/run_qwen3_vl_235b_a22b_megatron.sh new file mode 100644 index 0000000000000000000000000000000000000000..ce7102b11debae2601d0bd86b015b93e9a9f3c95 --- /dev/null +++ b/verl/examples/grpo_trainer/run_qwen3_vl_235b_a22b_megatron.sh @@ -0,0 +1,137 @@ +set -x +export CUDA_DEVICE_MAX_CONNECTIONS=1 # For megatron communication/computation overlapping + +export VLLM_ALLREDUCE_USE_SYMM_MEM=0 # for vllm0.11.0 with TP + +# ---- user-adjustable ---- +INFER_BACKEND=${INFER_BACKEND:-vllm} +HF_MODEL_PATH=${HF_MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen3-VL-235B-A22B-Instruct"} +TRAIN_FILE=${TRAIN_FILE:-$HOME/data/geo3k/train.parquet} +TEST_FILE=${TEST_FILE:-$HOME/data/geo3k/test.parquet} +NNODES=${NNODES:-8} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +GEN_TP=${GEN_TP:-16} +CP=${CP:-2} +TP=${TP:-4} +PP=${PP:-8} +EP=${EP:-8} +ETP=${ETP:-1} + +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-512} +PPO_MINI_BATCH_SIZE=${PPO_MINI_BATCH_SIZE:-128} +PPO_MICRO_BATCH_SIZE_PER_GPU=${PPO_MICRO_BATCH_SIZE_PER_GPU:-1} +LOG_PROB_MICRO_BATCH_SIZE_PER_GPU=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU:-1} +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-1024} +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-2048} +PPO_MAX_TOKEN_LEN_PER_GPU=${PPO_MAX_TOKEN_LEN_PER_GPU:-4096} + +ACTOR_LR=${ACTOR_LR:-1e-6} +KL_LOSS_COEF=${KL_LOSS_COEF:-0.01} +ENTROPY_COEFF=${ENTROPY_COEFF:-0} + +ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.7} +ROLLOUT_N=${ROLLOUT_N:-5} + +PROJECT_NAME=${PROJECT_NAME:-verl_grpo_example_geo3k} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-qwen3_vl_235b_megatron} +SAVE_FREQ=${SAVE_FREQ:-20} +TEST_FREQ=${TEST_FREQ:-5} +TOTAL_EPOCHS=${TOTAL_EPOCHS:-15} +# ---- end user-adjustable ---- +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + data.train_files="$TRAIN_FILE" + data.val_files="$TEST_FILE" + data.train_batch_size=${TRAIN_BATCH_SIZE} + data.max_prompt_length=${MAX_PROMPT_LENGTH} + data.max_response_length=${MAX_RESPONSE_LENGTH} + data.filter_overlong_prompts=True + data.truncation='error' + algorithm.use_kl_in_reward=False +) + +MODEL=( + actor_rollout_ref.model.path=$HF_MODEL_PATH +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${ACTOR_LR} + actor_rollout_ref.actor.ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE} + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${PPO_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=$PP + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=$TP + actor_rollout_ref.actor.megatron.context_parallel_size=$CP + actor_rollout_ref.actor.megatron.expert_model_parallel_size=$EP + actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=$ETP + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${KL_LOSS_COEF} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${ENTROPY_COEFF} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} + actor_rollout_ref.actor.megatron.use_mbridge=True + actor_rollout_ref.actor.megatron.param_offload=True + actor_rollout_ref.actor.megatron.optimizer_offload=True + actor_rollout_ref.actor.megatron.grad_offload=True + +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=1 + +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True + +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=True + +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32 + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_enable_deepep=True + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_token_dispatcher_type=flex + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 + +actor_rollout_ref.actor.megatron.override_transformer_config.gradient_accumulation_fusion=True + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=True + +actor_rollout_ref.actor.megatron.override_transformer_config.account_for_loss_in_pipeline_split=True + +actor_rollout_ref.actor.megatron.override_transformer_config.account_for_embedding_in_pipeline_split=True +) + +ROLLOUT=( + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.rollout.tensor_model_parallel_size=$GEN_TP + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} + actor_rollout_ref.rollout.name=${INFER_BACKEND} + actor_rollout_ref.rollout.gpu_memory_utilization=${ROLLOUT_GPU_MEM_UTIL} + actor_rollout_ref.rollout.n=${ROLLOUT_N} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.ref.megatron.param_offload=True +) + +TRAINER=( + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${PROJECT_NAME} + trainer.experiment_name=${EXPERIMENT_NAME} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${SAVE_FREQ} + trainer.test_freq=${TEST_FREQ} + trainer.total_epochs=${TOTAL_EPOCHS} +) + +EXTRA=( + model_engine=megatron +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" \ No newline at end of file diff --git a/verl/examples/grpo_trainer/run_qwen3_vl_30b_a3b_fsdp.sh b/verl/examples/grpo_trainer/run_qwen3_vl_30b_a3b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..cbd5d3fa6a70940407db3f7162397683647c5007 --- /dev/null +++ b/verl/examples/grpo_trainer/run_qwen3_vl_30b_a3b_fsdp.sh @@ -0,0 +1,173 @@ +#!/usr/bin/env bash +# GRPO | Qwen3-VL-30B-A3B (MoE) | FSDP training | GPU/NPU +set -xeuo pipefail + +########################### user-adjustable ########################### +# DEVICE is auto-detected by probing torch_npu; override only for special cases. +DEVICE=${DEVICE:-$(python3 -c 'import torch_npu' 2>/dev/null && echo npu || echo gpu)} +PROJECT_NAME=${PROJECT_NAME:-verl_grpo_geo3k} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-qwen3_vl_30b_a3b_grpo_${INFER_BACKEND}_fsdp_$(date +%Y%m%d_%H%M)} +INFER_BACKEND=${INFER_BACKEND:-vllm} +RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"} +MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen3-VL-30B-A3B-Instruct"} +CKPTS_DIR=${CKPTS_DIR:-} +TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/geo3k/train.parquet"} +TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/geo3k/test.parquet"} +WORKING_DIR=${WORKING_DIR:-"${PWD}"} +RUNTIME_ENV=${RUNTIME_ENV:-"${WORKING_DIR}/verl/trainer/runtime_env.yaml"} +NNODES=${NNODES:-} +NDEVICES_PER_NODE=${NDEVICES_PER_NODE:-} + +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-512} +PPO_MINI_BATCH_SIZE=${PPO_MINI_BATCH_SIZE:-128} +PPO_MICRO_BATCH_SIZE_PER_GPU=${PPO_MICRO_BATCH_SIZE_PER_GPU:-10} +LOG_PROB_MICRO_BATCH_SIZE_PER_GPU=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU:-2} +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-1024} +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-2048} + +ACTOR_LR=${ACTOR_LR:-1e-6} +KL_LOSS_COEF=${KL_LOSS_COEF:-0.01} +ENTROPY_COEFF=${ENTROPY_COEFF:-0} +FSDP_SIZE=${FSDP_SIZE:-} +SP_SIZE=${SP_SIZE:-2} + +ROLLOUT_TP=${ROLLOUT_TP:-} +ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-} +ROLLOUT_N=${ROLLOUT_N:-5} +ROLLOUT_MAX_NUM_BATCHED_TOKENS=${ROLLOUT_MAX_NUM_BATCHED_TOKENS:-20000} + +ROLLOUT_IS=${ROLLOUT_IS:-sequence} +ROLLOUT_IS_THRESHOLD=${ROLLOUT_IS_THRESHOLD:-2.0} +ROLLOUT_IS_BATCH_NORMALIZE=${ROLLOUT_IS_BATCH_NORMALIZE:-true} +ROLLOUT_RS=${ROLLOUT_RS:-token_k1} +ROLLOUT_RS_THRESHOLD=${ROLLOUT_RS_THRESHOLD:-0.6_1.6} + +SAVE_FREQ=${SAVE_FREQ:-5} +TEST_FREQ=${TEST_FREQ:-5} +TOTAL_EPOCHS=${TOTAL_EPOCHS:-15} +########################### end user-adjustable ########################### + +########################### derived defaults ########################### +n_devices_per_node=${NDEVICES_PER_NODE:-8} + +case "${DEVICE}" in + gpu) + nnodes=${NNODES:-1} + fsdp_size=${FSDP_SIZE:-8} + rollout_tp=${ROLLOUT_TP:-4} + rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.6} + ;; + npu) + export HCCL_CONNECT_TIMEOUT=1500 + export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050 + export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050 + export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 + + nnodes=${NNODES:-2} + fsdp_size=${FSDP_SIZE:-16} + rollout_tp=${ROLLOUT_TP:-8} + rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.8} + ;; + *) + echo "Unsupported DEVICE=${DEVICE}. Expected 'gpu' or 'npu'." >&2 + exit 1 + ;; +esac + +ckpts_dir=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${PROJECT_NAME}/${EXPERIMENT_NAME}"} + +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + algorithm.rollout_correction.rollout_is=${ROLLOUT_IS} + algorithm.rollout_correction.rollout_is_threshold=${ROLLOUT_IS_THRESHOLD} + algorithm.rollout_correction.rollout_is_batch_normalize=${ROLLOUT_IS_BATCH_NORMALIZE} + algorithm.rollout_correction.rollout_rs=${ROLLOUT_RS} + algorithm.rollout_correction.rollout_rs_threshold=${ROLLOUT_RS_THRESHOLD} + data.train_files="${TRAIN_FILE}" + data.val_files="${TEST_FILE}" + data.train_batch_size=${TRAIN_BATCH_SIZE} + data.max_prompt_length=${MAX_PROMPT_LENGTH} + data.max_response_length=${MAX_RESPONSE_LENGTH} + data.filter_overlong_prompts=True + data.truncation='error' + data.image_key=images +) + +MODEL=( + actor_rollout_ref.model.path=${MODEL_PATH} + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.use_fused_kernels=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.strategy=fsdp2 + actor_rollout_ref.actor.optim.lr=${ACTOR_LR} + actor_rollout_ref.actor.ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE} + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${PPO_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${KL_LOSS_COEF} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${ENTROPY_COEFF} + actor_rollout_ref.actor.entropy_from_logits_with_chunking=True + actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} + actor_rollout_ref.actor.fsdp_config.reshard_after_forward=True + actor_rollout_ref.actor.fsdp_config.entropy_checkpointing=True + actor_rollout_ref.actor.fsdp_config.ulysses_sequence_parallel_size=${SP_SIZE} + actor_rollout_ref.actor.fsdp_config.param_offload=False + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False + actor_rollout_ref.actor.fsdp_config.forward_prefetch=True +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=${INFER_BACKEND} + actor_rollout_ref.rollout.max_num_batched_tokens=${ROLLOUT_MAX_NUM_BATCHED_TOKENS} + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.enable_chunked_prefill=True + actor_rollout_ref.rollout.enforce_eager=False + actor_rollout_ref.rollout.free_cache_engine=True + actor_rollout_ref.rollout.n=${ROLLOUT_N} + actor_rollout_ref.rollout.calculate_log_probs=True +) + +REF=( + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.ref.fsdp_config.param_offload=True + actor_rollout_ref.ref.fsdp_config.reshard_after_forward=True + actor_rollout_ref.ref.fsdp_config.forward_prefetch=True + actor_rollout_ref.ref.fsdp_config.ulysses_sequence_parallel_size=${SP_SIZE} + actor_rollout_ref.ref.entropy_from_logits_with_chunking=True +) + +TRAINER=( + trainer.critic_warmup=0 + trainer.logger='["console", "wandb"]' + trainer.project_name="${PROJECT_NAME}" + trainer.experiment_name="${EXPERIMENT_NAME}" + trainer.n_gpus_per_node=${n_devices_per_node} + trainer.nnodes=${nnodes} + trainer.default_local_dir=${ckpts_dir} + trainer.resume_mode=auto + trainer.val_before_train=True + trainer.save_freq=${SAVE_FREQ} + trainer.test_freq=${TEST_FREQ} + trainer.total_epochs=${TOTAL_EPOCHS} +) + +EXTRA=() + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/grpo_trainer/run_qwen3_vl_30b_a3b_megatron.sh b/verl/examples/grpo_trainer/run_qwen3_vl_30b_a3b_megatron.sh new file mode 100644 index 0000000000000000000000000000000000000000..bc9dc34c89185d7e99812e60794eb0e081f0d2d6 --- /dev/null +++ b/verl/examples/grpo_trainer/run_qwen3_vl_30b_a3b_megatron.sh @@ -0,0 +1,138 @@ +set -x +export CUDA_DEVICE_MAX_CONNECTIONS=1 # For megatron communication/computation overlapping + +export VLLM_ALLREDUCE_USE_SYMM_MEM=0 # for vllm0.11.0 with TP + +# ---- user-adjustable ---- +INFER_BACKEND=${INFER_BACKEND:-vllm} +HF_MODEL_PATH=${HF_MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen3-VL-30B-A3B-Instruct"} +TRAIN_FILE=${TRAIN_FILE:-$HOME/data/geo3k/train.parquet} +TEST_FILE=${TEST_FILE:-$HOME/data/geo3k/test.parquet} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +GEN_TP=${GEN_TP:-4} +CP=${CP:-2} +TP=${TP:-2} +PP=${PP:-1} +EP=${EP:-8} +ETP=${ETP:-1} + +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-512} +PPO_MINI_BATCH_SIZE=${PPO_MINI_BATCH_SIZE:-128} +PPO_MICRO_BATCH_SIZE_PER_GPU=${PPO_MICRO_BATCH_SIZE_PER_GPU:-1} +LOG_PROB_MICRO_BATCH_SIZE_PER_GPU=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU:-1} +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-1024} +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-2048} +PPO_MAX_TOKEN_LEN_PER_GPU=${PPO_MAX_TOKEN_LEN_PER_GPU:-4096} + +ACTOR_LR=${ACTOR_LR:-1e-6} +KL_LOSS_COEF=${KL_LOSS_COEF:-0.01} +ENTROPY_COEFF=${ENTROPY_COEFF:-0} + +ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.7} +ROLLOUT_N=${ROLLOUT_N:-5} + +PROJECT_NAME=${PROJECT_NAME:-verl_grpo_example_geo3k} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-qwen3_vl_30b_megatron} +SAVE_FREQ=${SAVE_FREQ:-20} +TEST_FREQ=${TEST_FREQ:-5} +TOTAL_EPOCHS=${TOTAL_EPOCHS:-15} +# ---- end user-adjustable ---- +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files="$TRAIN_FILE" + data.val_files="$TEST_FILE" + data.train_batch_size=${TRAIN_BATCH_SIZE} + data.max_prompt_length=${MAX_PROMPT_LENGTH} + data.max_response_length=${MAX_RESPONSE_LENGTH} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path=$HF_MODEL_PATH +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${ACTOR_LR} + actor_rollout_ref.actor.ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE} + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${PPO_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=$PP + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=$TP + actor_rollout_ref.actor.megatron.context_parallel_size=$CP + actor_rollout_ref.actor.megatron.expert_model_parallel_size=$EP + actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=$ETP + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${KL_LOSS_COEF} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${ENTROPY_COEFF} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} + actor_rollout_ref.actor.megatron.use_mbridge=True + actor_rollout_ref.actor.megatron.param_offload=True + actor_rollout_ref.actor.megatron.optimizer_offload=True + actor_rollout_ref.actor.megatron.grad_offload=True + +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=1 + +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True + +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=True + +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32 + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_enable_deepep=True + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_token_dispatcher_type=flex + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 + +actor_rollout_ref.actor.megatron.override_transformer_config.gradient_accumulation_fusion=True + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=True + # Use aux_loss and z_loss to mitigate expert load imbalance when training MoE models + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_aux_loss_coeff=0.01 + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_z_loss_coeff=0.001 +) + +ROLLOUT=( + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.rollout.tensor_model_parallel_size=$GEN_TP + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} + actor_rollout_ref.rollout.name=${INFER_BACKEND} + actor_rollout_ref.rollout.gpu_memory_utilization=${ROLLOUT_GPU_MEM_UTIL} + actor_rollout_ref.rollout.n=${ROLLOUT_N} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.ref.megatron.param_offload=True +) + +TRAINER=( + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${PROJECT_NAME} + trainer.experiment_name=${EXPERIMENT_NAME} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${SAVE_FREQ} + trainer.test_freq=${TEST_FREQ} + trainer.total_epochs=${TOTAL_EPOCHS} +) + +EXTRA=( + model_engine=megatron +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/grpo_trainer/run_qwen3_vl_30b_moe_veomni.sh b/verl/examples/grpo_trainer/run_qwen3_vl_30b_moe_veomni.sh new file mode 100644 index 0000000000000000000000000000000000000000..37781c46919326f323899e2389a6e9cbe4aa3caa --- /dev/null +++ b/verl/examples/grpo_trainer/run_qwen3_vl_30b_moe_veomni.sh @@ -0,0 +1,177 @@ +#!/usr/bin/env bash +# +# GRPO | Qwen3-VL-30B-A3B | VeOmni training | NVIDIA GPUs or Ascend NPUs +# +# INFER_BACKEND controls rollout backend: vllm + +set -xeuo pipefail + +# ---- user-adjustable ---- +data_path=${data_path:-$HOME/geo3k} +model_path=${model_path:-$HOME/Qwen3-VL-30B-A3B-Instruct} + +nnodes=${nnodes:-2} +num_gpus_per_node=8 + +# Parallelism settings +dp_size=${dp_size:-16} +usp_size=${usp_size:-1} +ep_size=${ep_size:-1} + +# Model and project settings +model_id=Qwen3_VL-30B-MOE +project_name=${model_id}-veomni +exp_name=grpo-${num_gpus_per_node}gpu + +backend=fsdp2 +model_engine=veomni + +# ===================================== Algorithm ===================================== +adv_estimator=grpo +use_kl_in_reward=False +kl_coef=0.001 + +# Actor settings +use_kl_loss=True +kl_loss_coef=0.01 +kl_loss_type=low_var_kl +entropy_coeff=0 +actor_lr=1e-6 +lr_scheduler_type=constant + +ppo_mini_batch_size=32 +critic_warmup=0 + +# ===================================== Data/Model ===================================== +train_files=$data_path/train.parquet +test_files=$data_path/test.parquet + +actor_model_path=$model_path + +max_prompt_length=1024 +max_response_length=2048 +train_batch_size=64 + +use_remove_padding=True +enable_gradient_checkpointing=True +max_position_embeddings=32768 + +# ===================================== Training ===================================== +ppo_micro_batch_size_per_gpu=1 + +# VeOmni config +ACTOR_VEOMNI_CONFIG=" + actor_rollout_ref.actor.veomni.param_offload=True \ + actor_rollout_ref.actor.veomni.optimizer_offload=True \ + actor_rollout_ref.actor.veomni.enable_full_shard=True \ + actor_rollout_ref.actor.veomni.fsdp_size=$dp_size \ + actor_rollout_ref.actor.veomni.ulysses_parallel_size=$usp_size \ + actor_rollout_ref.actor.veomni.expert_parallel_size=$ep_size \ + actor_rollout_ref.actor.veomni.attn_implementation=veomni_flash_attention_2_with_sp \ + actor_rollout_ref.actor.veomni.moe_implementation=fused" + +# Ref model config +REF_VEOMNI_CONFIG=" + actor_rollout_ref.ref.strategy=veomni \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.ref.veomni.param_offload=False \ + actor_rollout_ref.ref.veomni.expert_parallel_size=1" +# ---- end user-adjustable ---- + +# ---- no user adjustment needed below ---- +# Actor model config +ACTOR_CONFIG=" + actor_rollout_ref.actor.optim.lr=$actor_lr \ + actor_rollout_ref.actor.optim.lr_scheduler_type=$lr_scheduler_type \ + actor_rollout_ref.model.path=$actor_model_path \ + actor_rollout_ref.model.use_remove_padding=$use_remove_padding \ + actor_rollout_ref.model.enable_gradient_checkpointing=$enable_gradient_checkpointing \ + +actor_rollout_ref.model.override_config.max_position_embeddings=$max_position_embeddings \ + actor_rollout_ref.actor.use_kl_loss=$use_kl_loss \ + actor_rollout_ref.actor.kl_loss_coef=$kl_loss_coef \ + actor_rollout_ref.actor.kl_loss_type=$kl_loss_type \ + actor_rollout_ref.actor.entropy_coeff=$entropy_coeff \ + actor_rollout_ref.actor.use_torch_compile=False \ + actor_rollout_ref.actor.ppo_mini_batch_size=$ppo_mini_batch_size \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=$ppo_micro_batch_size_per_gpu" + +CONFIG_NAME=ppo_trainer +ACTOR_CONFIG="$ACTOR_CONFIG $ACTOR_VEOMNI_CONFIG $REF_VEOMNI_CONFIG" + +# ===================================== Inference ===================================== +rollout_name=vllm +infer_tp=4 +infer_dp=1 +infer_ep=1 +gpu_memory_utilization=0.6 +n_resp_per_prompt=8 +max_model_len=4096 +max_num_batched_tokens=5120 + +ROLLOUT_CONFIG=" + actor_rollout_ref.rollout.name=$rollout_name \ + actor_rollout_ref.rollout.tensor_model_parallel_size=$infer_tp \ + actor_rollout_ref.rollout.data_parallel_size=$infer_dp \ + actor_rollout_ref.rollout.expert_parallel_size=$infer_ep \ + actor_rollout_ref.rollout.gpu_memory_utilization=$gpu_memory_utilization \ + actor_rollout_ref.rollout.enable_chunked_prefill=False \ + actor_rollout_ref.rollout.enable_prefix_caching=False \ + actor_rollout_ref.rollout.enforce_eager=True \ + actor_rollout_ref.rollout.free_cache_engine=True \ + actor_rollout_ref.rollout.n=$n_resp_per_prompt \ + actor_rollout_ref.rollout.max_model_len=$max_model_len \ + actor_rollout_ref.rollout.max_num_batched_tokens=$max_num_batched_tokens \ + actor_rollout_ref.rollout.disable_log_stats=False \ + actor_rollout_ref.rollout.calculate_log_probs=True \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \ + +actor_rollout_ref.rollout.engine_kwargs.vllm.max_model_len=$max_model_len" + +########################### parameter arrays ########################### + +CONFIG=( + --config-path=./config + --config-name=$CONFIG_NAME +) + +DATA=( + algorithm.adv_estimator=$adv_estimator + algorithm.use_kl_in_reward=$use_kl_in_reward + algorithm.kl_ctrl.kl_coef=$kl_coef + data.train_files="$train_files" + data.val_files="$test_files" + data.train_batch_size=$train_batch_size + data.max_prompt_length=$max_prompt_length + data.max_response_length=$max_response_length + data.filter_overlong_prompts=True + data.shuffle=False + data.truncation='error' + data.image_key=images +) + +TRAINER=( + trainer.resume_mode=auto + trainer.critic_warmup=$critic_warmup + trainer.logger=['console','wandb'] + trainer.project_name=$project_name + trainer.experiment_name=$exp_name + trainer.n_gpus_per_node=$num_gpus_per_node + trainer.nnodes=$nnodes + trainer.save_freq=40 + trainer.test_freq=5 + trainer.total_epochs=20 + trainer.total_training_steps=200 +) + +EXTRA=( + model_engine=$model_engine + $ACTOR_CONFIG + $ROLLOUT_CONFIG +) + +########################### launch ########################### +python -m verl.trainer.main_ppo \ + "${CONFIG[@]}" \ + "${DATA[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" \ No newline at end of file diff --git a/verl/examples/grpo_trainer/run_qwen3_vl_8b_fsdp.sh b/verl/examples/grpo_trainer/run_qwen3_vl_8b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..9d6db0e03fa2d36de944ba070c10f74aefa62dec --- /dev/null +++ b/verl/examples/grpo_trainer/run_qwen3_vl_8b_fsdp.sh @@ -0,0 +1,152 @@ +#!/usr/bin/env bash +# GRPO | vision | vLLM rollout | FSDP training | GPU/NPU +# Canonical Qwen3-VL baseline on Geo3K. + +set -xeuo pipefail + +########################### user-adjustable ########################### +# DEVICE is auto-detected by probing torch_npu; override only for special cases. +DEVICE=${DEVICE:-$(python3 -c 'import torch_npu' 2>/dev/null && echo npu || echo gpu)} +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-VL-8B-Instruct} +NNODES=${NNODES:-1} +NDEVICES_PER_NODE=${NDEVICES_PER_NODE:-} + +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-512} +PPO_MINI_BATCH_SIZE=${PPO_MINI_BATCH_SIZE:-128} +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-1024} +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-2048} +PPO_MAX_TOKEN_LEN_PER_GPU=${PPO_MAX_TOKEN_LEN_PER_GPU:-24576} + +ACTOR_LR=${ACTOR_LR:-1e-6} +KL_LOSS_COEF=${KL_LOSS_COEF:-0.01} +ENTROPY_COEFF=${ENTROPY_COEFF:-0} + +ROLLOUT_TP=${ROLLOUT_TP:-2} +ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-} +ROLLOUT_N=${ROLLOUT_N:-5} +SP_SIZE=${SP_SIZE:-1} + +TOTAL_EPOCHS=${TOTAL_EPOCHS:-15} +SAVE_FREQ=${SAVE_FREQ:-20} +TEST_FREQ=${TEST_FREQ:-5} + +PROJECT_NAME=${PROJECT_NAME:-verl_grpo_geo3k} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-qwen3_vl_8b_grpo_vllm_fsdp_$(date +%Y%m%d_%H%M)} + +TRAIN_FILE=${TRAIN_FILE:-$HOME/data/geo3k/train.parquet} +TEST_FILE=${TEST_FILE:-$HOME/data/geo3k/test.parquet} +########################### end user-adjustable ########################### + +########################### derived defaults ########################### +n_devices_per_node=${NDEVICES_PER_NODE:-8} + +case "${DEVICE}" in + gpu) + rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.6} + ;; + npu) + export HCCL_CONNECT_TIMEOUT=1500 + export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050 + export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050 + export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 + + rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.5} + ;; + *) + echo "Unsupported DEVICE=${DEVICE}. Expected 'gpu' or 'npu'." >&2 + exit 1 + ;; +esac + +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files=${TRAIN_FILE} + data.val_files=${TEST_FILE} + data.image_key=images + data.train_batch_size=${TRAIN_BATCH_SIZE} + data.max_prompt_length=${MAX_PROMPT_LENGTH} + data.max_response_length=${MAX_RESPONSE_LENGTH} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.strategy=fsdp2 + actor_rollout_ref.actor.optim.lr=${ACTOR_LR} + actor_rollout_ref.actor.ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${KL_LOSS_COEF} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${ENTROPY_COEFF} +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${ROLLOUT_TP} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.enable_chunked_prefill=False + actor_rollout_ref.rollout.n=${ROLLOUT_N} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} + actor_rollout_ref.ref.fsdp_config.param_offload=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.logger='["console","wandb"]' + trainer.project_name=${PROJECT_NAME} + trainer.experiment_name=${EXPERIMENT_NAME} + trainer.n_gpus_per_node=${n_devices_per_node} + trainer.nnodes=${NNODES} + trainer.save_freq=${SAVE_FREQ} + trainer.test_freq=${TEST_FREQ} + trainer.total_epochs=${TOTAL_EPOCHS} +) + +# Per-device extras (single trailing array, never empty). +if [ "${DEVICE}" = npu ]; then + EXTRA=( + actor_rollout_ref.actor.use_torch_compile=False + actor_rollout_ref.actor.fsdp_config.param_offload=True + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True + actor_rollout_ref.actor.fsdp_config.ulysses_sequence_parallel_size=${SP_SIZE} + actor_rollout_ref.ref.fsdp_config.ulysses_sequence_parallel_size=${SP_SIZE} + +actor_rollout_ref.rollout.engine_kwargs.vllm.mm_processor_cache_gb=0 + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=20 + ) +else + EXTRA=( + actor_rollout_ref.model.use_fused_kernels=True + actor_rollout_ref.actor.fsdp_config.param_offload=False + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False + actor_rollout_ref.rollout.enforce_eager=False + actor_rollout_ref.rollout.free_cache_engine=True + ) +fi + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/grpo_trainer/run_qwen3_vl_8b_megatron.sh b/verl/examples/grpo_trainer/run_qwen3_vl_8b_megatron.sh new file mode 100644 index 0000000000000000000000000000000000000000..b8216b23c2e93d64435bae6ed742f644809efbdc --- /dev/null +++ b/verl/examples/grpo_trainer/run_qwen3_vl_8b_megatron.sh @@ -0,0 +1,114 @@ +#!/usr/bin/env bash +# GRPO | vision | vLLM rollout | Megatron training | NVIDIA GPUs + +set -xeuo pipefail + +export CUDA_DEVICE_MAX_CONNECTIONS=1 + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-VL-8B-Instruct} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +train_batch_size=${TRAIN_BATCH_SIZE:-512} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-128} +max_prompt_length=${MAX_PROMPT_LENGTH:-1024} +max_response_length=${MAX_RESPONSE_LENGTH:-2048} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-24576} + +actor_lr=${ACTOR_LR:-1e-6} +kl_loss_coef=${KL_LOSS_COEF:-0.01} +entropy_coeff=${ENTROPY_COEFF:-0} + +actor_tp=${ACTOR_TP:-2} +actor_pp=${ACTOR_PP:-2} + +rollout_tp=${ROLLOUT_TP:-2} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.6} +rollout_n=${ROLLOUT_N:-5} + +total_epochs=${TOTAL_EPOCHS:-15} +save_freq=${SAVE_FREQ:-20} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_grpo_geo3k} +experiment_name=${EXPERIMENT_NAME:-qwen3_vl_8b_vllm_megatron} +# ---- end user-adjustable ---- +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files=$HOME/data/geo3k/train.parquet + data.val_files=$HOME/data/geo3k/test.parquet + data.image_key=images + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.use_fused_kernels=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${actor_pp} + actor_rollout_ref.actor.megatron.use_mbridge=True +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${actor_pp} + actor_rollout_ref.ref.megatron.use_mbridge=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( + model_engine=megatron +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/grpo_trainer/run_seed_oss_36b_fsdp.sh b/verl/examples/grpo_trainer/run_seed_oss_36b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..7e92a158caa4d34ebbe67f2f6452151c6d7dbfa9 --- /dev/null +++ b/verl/examples/grpo_trainer/run_seed_oss_36b_fsdp.sh @@ -0,0 +1,108 @@ +set -x + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-ByteDance-Seed/Seed-OSS-36B-Base} +TRAIN_FILE=${TRAIN_FILE:-$HOME/data/gsm8k/train.parquet} +TEST_FILE=${TEST_FILE:-$HOME/data/gsm8k/test.parquet} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-64} +PPO_MINI_BATCH_SIZE=${PPO_MINI_BATCH_SIZE:-8} +PPO_MICRO_BATCH_SIZE_PER_GPU=${PPO_MICRO_BATCH_SIZE_PER_GPU:-2} +LOG_PROB_MICRO_BATCH_SIZE_PER_GPU=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU:-2} +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-512} +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-1024} + +ACTOR_LR=${ACTOR_LR:-1e-6} +KL_LOSS_COEF=${KL_LOSS_COEF:-0.001} +ENTROPY_COEFF=${ENTROPY_COEFF:-0} + +ROLLOUT_TP=${ROLLOUT_TP:-2} +ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.6} +ROLLOUT_N=${ROLLOUT_N:-2} + +PROJECT_NAME=${PROJECT_NAME:-verl_grpo_seed_oss_36b} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-seed_oss_36b} +SAVE_FREQ=${SAVE_FREQ:-20} +TEST_FREQ=${TEST_FREQ:-5} +TOTAL_EPOCHS=${TOTAL_EPOCHS:-15} +# ---- end user-adjustable ---- +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + data.train_files=${TRAIN_FILE} + data.val_files=${TEST_FILE} + data.train_batch_size=${TRAIN_BATCH_SIZE} + data.max_prompt_length=${MAX_PROMPT_LENGTH} + data.max_response_length=${MAX_RESPONSE_LENGTH} + data.filter_overlong_prompts=True + data.truncation='error' + algorithm.use_kl_in_reward=False +) + +MODEL=( + actor_rollout_ref.model.path=${MODEL_PATH} + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True + actor_rollout_ref.model.use_fused_kernels=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${ACTOR_LR} + actor_rollout_ref.actor.ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE} + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${PPO_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${KL_LOSS_COEF} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${ENTROPY_COEFF} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.strategy=fsdp2 + actor_rollout_ref.actor.fsdp_config.param_offload=True + actor_rollout_ref.actor.fsdp_config.param_offload=True +) + +ROLLOUT=( + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.rollout.tensor_model_parallel_size=${ROLLOUT_TP} + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.gpu_memory_utilization=${ROLLOUT_GPU_MEM_UTIL} + actor_rollout_ref.rollout.n=${ROLLOUT_N} + actor_rollout_ref.rollout.free_cache_engine=True +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.ref.fsdp_config.param_offload=True + actor_rollout_ref.ref.strategy=fsdp2 +) + +TRAINER=( + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${PROJECT_NAME} + trainer.experiment_name=${EXPERIMENT_NAME} + trainer.val_before_train=False + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${SAVE_FREQ} + trainer.test_freq=${TEST_FREQ} + trainer.total_epochs=${TOTAL_EPOCHS} +) + +EXTRA=( +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" \ No newline at end of file diff --git a/verl/examples/gspo_trainer/README.md b/verl/examples/gspo_trainer/README.md new file mode 100644 index 0000000000000000000000000000000000000000..6d594a25e111f143e6f7eae0a02883d777211434 --- /dev/null +++ b/verl/examples/gspo_trainer/README.md @@ -0,0 +1,21 @@ +# GSPO (Group Sequence Policy Optimization) + +GSPO is a GRPO-family policy-loss variant that aggregates the IS-ratio at the sequence level (`seq-mean-token-mean`) and uses a very tight clip window. It is especially useful for large MoE models. + +Reference: [Group Sequence Policy Optimization](https://arxiv.org/abs/2507.18071). + +## Canonical Scripts + +| Script | Infer | Train | Platform | +|---------------------------------------------|-------|----------|----------| +| `run_qwen3_8b_fsdp.sh` | vLLM | FSDP | NVIDIA | +| `run_qwen3_8b_fsdp.sh` | vLLM | FSDP | Ascend | +| `run_qwen3_30b_a3b_megatron.sh` | vLLM | Megatron | NVIDIA | + +## Key Flags + +- `actor_rollout_ref.actor.policy_loss.loss_mode=gspo` +- `actor_rollout_ref.actor.loss_agg_mode=seq-mean-token-mean` +- `actor_rollout_ref.actor.clip_ratio_low=3e-4` +- `actor_rollout_ref.actor.clip_ratio_high=4e-4` +- `actor_rollout_ref.actor.clip_ratio_c=10.0` (dual-clip guard) diff --git a/verl/examples/gspo_trainer/run_qwen3_30b_a3b_megatron.sh b/verl/examples/gspo_trainer/run_qwen3_30b_a3b_megatron.sh new file mode 100644 index 0000000000000000000000000000000000000000..bc2a1fd7ab1db7dfbe6f67c32a5b4492d565a1cf --- /dev/null +++ b/verl/examples/gspo_trainer/run_qwen3_30b_a3b_megatron.sh @@ -0,0 +1,143 @@ +#!/usr/bin/env bash +# GSPO | MoE | vLLM rollout | Megatron training | NVIDIA GPUs +# GSPO with Qwen3-30B-A3B (MoE). + +set -xeuo pipefail + +export CUDA_DEVICE_MAX_CONNECTIONS=1 +export VLLM_USE_V1=1 + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-30B-A3B-Base} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +train_batch_size=${TRAIN_BATCH_SIZE:-256} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-32} +max_prompt_length=${MAX_PROMPT_LENGTH:-2048} +max_response_length=${MAX_RESPONSE_LENGTH:-8192} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-30720} + +actor_lr=${ACTOR_LR:-1e-6} +entropy_coeff=${ENTROPY_COEFF:-0} + +clip_ratio_low=${CLIP_RATIO_LOW:-3e-4} +clip_ratio_high=${CLIP_RATIO_HIGH:-4e-4} + +actor_tp=${ACTOR_TP:-2} +actor_pp=${ACTOR_PP:-1} +actor_ep=${ACTOR_EP:-8} +actor_etp=${ACTOR_ETP:-1} + +rollout_tp=${ROLLOUT_TP:-4} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.8} +rollout_n=${ROLLOUT_N:-16} + +total_epochs=${TOTAL_EPOCHS:-10} +save_freq=${SAVE_FREQ:-50} +test_freq=${TEST_FREQ:-10} + +project_name=${PROJECT_NAME:-verl_gspo_qwen3_moe} +experiment_name=${EXPERIMENT_NAME:-qwen3_30b_a3b_vllm_megatron} + +train_file=${TRAIN_FILE:-$HOME/data/dapo-math-17k/train.parquet} +val_file=${VAL_FILE:-$HOME/data/aime-2024/test.parquet} +# ---- end user-adjustable ---- +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files="['$train_file']" + data.val_files="['$val_file']" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True +) + +ACTOR=( + actor_rollout_ref.actor.policy_loss.loss_mode=gspo + actor_rollout_ref.actor.loss_agg_mode=seq-mean-token-mean + actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} + actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} + actor_rollout_ref.actor.clip_ratio_c=10.0 + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=False + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${actor_pp} + actor_rollout_ref.actor.megatron.expert_model_parallel_size=${actor_ep} + actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${actor_etp} + actor_rollout_ref.actor.megatron.param_offload=True + actor_rollout_ref.actor.megatron.grad_offload=True + actor_rollout_ref.actor.megatron.optimizer_offload=True + actor_rollout_ref.actor.megatron.use_mbridge=True + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32 + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=True + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 + +actor_rollout_ref.actor.megatron.override_transformer_config.apply_rope_fusion=True +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.rollout.val_kwargs.n=1 + actor_rollout_ref.rollout.val_kwargs.temperature=1.0 + actor_rollout_ref.rollout.val_kwargs.top_p=0.7 +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${actor_pp} + actor_rollout_ref.ref.megatron.expert_model_parallel_size=${actor_ep} + actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=${actor_etp} + actor_rollout_ref.ref.megatron.param_offload=True + actor_rollout_ref.ref.megatron.use_mbridge=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.val_before_train=False + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( + model_engine=megatron +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/gspo_trainer/run_qwen3_8b_fsdp.sh b/verl/examples/gspo_trainer/run_qwen3_8b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..177d8144b38340cc0a79057cf6f0e89cd410942d --- /dev/null +++ b/verl/examples/gspo_trainer/run_qwen3_8b_fsdp.sh @@ -0,0 +1,173 @@ +#!/usr/bin/env bash +# GSPO | text | vLLM rollout | FSDP training | GPU/NPU +# GSPO is a sequence-mean policy-loss variant on top of GRPO (paper: arXiv:2507.18071). + +set -xeuo pipefail + +########################### user-adjustable ########################### +# DEVICE is auto-detected by probing torch_npu; override only for special cases. +DEVICE=${DEVICE:-$(python3 -c 'import torch_npu' 2>/dev/null && echo npu || echo gpu)} +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-8B} +NNODES=${NNODES:-1} +NDEVICES_PER_NODE=${NDEVICES_PER_NODE:-} + +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-} +PPO_MINI_BATCH_SIZE=${PPO_MINI_BATCH_SIZE:-} +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-2048} +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-8192} +PPO_MAX_TOKEN_LEN_PER_GPU=${PPO_MAX_TOKEN_LEN_PER_GPU:-20480} + +ACTOR_LR=${ACTOR_LR:-1e-6} +ENTROPY_COEFF=${ENTROPY_COEFF:-0} + +CLIP_RATIO_LOW=${CLIP_RATIO_LOW:-3e-4} +CLIP_RATIO_HIGH=${CLIP_RATIO_HIGH:-4e-4} + +SP_SIZE=${SP_SIZE:-} +ROLLOUT_TP=${ROLLOUT_TP:-} +ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-} +ROLLOUT_N=${ROLLOUT_N:-16} + +TOTAL_EPOCHS=${TOTAL_EPOCHS:-10} +SAVE_FREQ=${SAVE_FREQ:-20} +TEST_FREQ=${TEST_FREQ:-10} + +PROJECT_NAME=${PROJECT_NAME:-verl_gspo_gsm8k_math} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-qwen3_8b_vllm_fsdp} + +GSM8K_TRAIN_FILE=${GSM8K_TRAIN_FILE:-$HOME/data/gsm8k/train.parquet} +GSM8K_TEST_FILE=${GSM8K_TEST_FILE:-$HOME/data/gsm8k/test.parquet} +MATH_TRAIN_FILE=${MATH_TRAIN_FILE:-$HOME/data/math/train.parquet} +MATH_TEST_FILE=${MATH_TEST_FILE:-$HOME/data/math/test.parquet} +########################### end user-adjustable ########################### + +########################### derived defaults ########################### +n_devices_per_node=${NDEVICES_PER_NODE:-8} +save_freq=${SAVE_FREQ} +test_freq=${TEST_FREQ} + +case "${DEVICE}" in + gpu) + train_batch_size=${TRAIN_BATCH_SIZE:-512} + ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-128} + sp_size=${SP_SIZE:-1} + rollout_tp=${ROLLOUT_TP:-2} + rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.6} + ;; + npu) + ulimit -n 32768 + export RAY_DEDUP_LOGS=0 + export HYDRA_FULL_ERROR=1 + export TASK_QUEUE_ENABLE=1 + export HCCL_EXEC_TIMEOUT=3600 + export HCCL_CONNECT_TIMEOUT=3600 + export HCCL_ASYNC_ERROR_HANDLING=0 + export CPU_AFFINITY_CONF=1 + export VLLM_USE_V1=1 + + train_batch_size=${TRAIN_BATCH_SIZE:-256} + ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-64} + sp_size=${SP_SIZE:-4} + rollout_tp=${ROLLOUT_TP:-4} + rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.7} + ;; + *) + echo "Unsupported DEVICE=${DEVICE}. Expected 'gpu' or 'npu'." >&2 + exit 1 + ;; +esac + +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files="['$GSM8K_TRAIN_FILE', '$MATH_TRAIN_FILE']" + data.val_files="['$GSM8K_TEST_FILE', '$MATH_TEST_FILE']" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${MAX_PROMPT_LENGTH} + data.max_response_length=${MAX_RESPONSE_LENGTH} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.policy_loss.loss_mode=gspo + actor_rollout_ref.actor.loss_agg_mode=seq-mean-token-mean + actor_rollout_ref.actor.clip_ratio_low=${CLIP_RATIO_LOW} + actor_rollout_ref.actor.clip_ratio_high=${CLIP_RATIO_HIGH} + actor_rollout_ref.actor.clip_ratio_c=10.0 + actor_rollout_ref.actor.optim.lr=${ACTOR_LR} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} + actor_rollout_ref.actor.use_kl_loss=False + actor_rollout_ref.actor.entropy_coeff=${ENTROPY_COEFF} + actor_rollout_ref.actor.fsdp_config.param_offload=True + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${ROLLOUT_N} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} + actor_rollout_ref.ref.fsdp_config.param_offload=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${PROJECT_NAME} + trainer.experiment_name=${EXPERIMENT_NAME} + trainer.n_gpus_per_node=${n_devices_per_node} + trainer.nnodes=${NNODES} + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${TOTAL_EPOCHS} +) + +# Per-device extras (single trailing array, never empty under set -u). +EXTRA=() +if [ "${DEVICE}" = npu ]; then + EXTRA+=( + actor_rollout_ref.actor.strategy=fsdp2 + actor_rollout_ref.actor.use_torch_compile=False + actor_rollout_ref.actor.fsdp_config.forward_prefetch=True + actor_rollout_ref.actor.fsdp_config.ulysses_sequence_parallel_size=${sp_size} + actor_rollout_ref.actor.entropy_checkpointing=True + actor_rollout_ref.actor.entropy_from_logits_with_chunking=True + actor_rollout_ref.rollout.val_kwargs.n=1 + actor_rollout_ref.rollout.val_kwargs.temperature=1.0 + actor_rollout_ref.rollout.val_kwargs.top_p=0.7 + actor_rollout_ref.ref.strategy=fsdp2 + actor_rollout_ref.ref.use_torch_compile=False + actor_rollout_ref.ref.fsdp_config.ulysses_sequence_parallel_size=${sp_size} + trainer.val_before_train=False + ) +fi + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/mtp_trainer/README.md b/verl/examples/mtp_trainer/README.md new file mode 100644 index 0000000000000000000000000000000000000000..48a635a6f0946332a8a4f4675f165353c0fdb9f9 --- /dev/null +++ b/verl/examples/mtp_trainer/README.md @@ -0,0 +1,33 @@ +# Multi-Token-Prediction (MTP) Training + +MTP uses an auxiliary token-prediction head (speculative / draft head) during training. Currently supported on MiMo-7B-RL with Megatron backend. + +## Canonical Scripts + +| Script | Infer | Train | Mode | Platform | +|---------------------------------------------------------------------------|--------|----------|-----------------------------------|----------| +| `run_mimo_7b_mtp_megatron.sh` | SGLang | Megatron | Sync hybrid-engine | NVIDIA | +| `run_mimo_7b_mtp_fully_async_megatron_multinode.sh` | SGLang | Megatron | Fully-async split-placement (DAPO)| NVIDIA | + +IMPORTANT: after downloading MiMo-7B-RL, set `max_position_embeddings: 32768` in its `config.json`. + +## Key Flags + +- `actor_rollout_ref.model.mtp.enable=True` +- `actor_rollout_ref.model.mtp.enable_train=True` +- `actor_rollout_ref.model.mtp.mtp_loss_scaling_factor=0.1` +- `actor_rollout_ref.model.mtp.detach_encoder=True` + +## Multi-node fully-async layout + +The `*_multinode.sh` variant uses the fully-async one-step-off trainer +(`verl.experimental.fully_async_policy.fully_async_main`). Scale it via: + +```bash +TRAIN_NNODES=4 TRAIN_NGPUS_PER_NODE=8 \ +ROLLOUT_NNODES=4 ROLLOUT_NGPUS_PER_NODE=8 \ +bash examples/mtp_trainer/run_mimo_7b_mtp_fully_async_megatron_multinode.sh +``` + +Defaults to a single-node 4+4 split (trainer + rollout) for a smoke-test, +matching the historical `..._math_megatron_4_4.sh` layout. diff --git a/verl/examples/mtp_trainer/run_mimo_7b_mtp_fully_async_megatron_multinode.sh b/verl/examples/mtp_trainer/run_mimo_7b_mtp_fully_async_megatron_multinode.sh new file mode 100644 index 0000000000000000000000000000000000000000..2a2807873722d36918942029a89c1204ed671924 --- /dev/null +++ b/verl/examples/mtp_trainer/run_mimo_7b_mtp_fully_async_megatron_multinode.sh @@ -0,0 +1,152 @@ +#!/usr/bin/env bash +# MTP + DAPO | MiMo-7B (speculative MTP head) | SGLang rollout | Megatron training | NVIDIA GPUs +# Fully-async (one-step-off) split-placement training demo, multi-node capable. +# +# Layout defaults to the "4+4" split (4 trainer GPUs + 4 rollout GPUs on a single node). +# Scale via TRAIN_NNODES / ROLLOUT_NNODES for a true multi-node run (e.g. 4 trainer nodes + +# 4 rollout nodes, hence the historical name "dapo_mimo_7b_with_mtp_math_megatron_4_4"). + +set -xeuo pipefail + +export CUDA_DEVICE_MAX_CONNECTIONS=1 + +# ---- user-adjustable ---- +# NOTE: remember to set max_position_embeddings=32768 in the model's config.json after downloading. +MODEL_PATH=${MODEL_PATH:-XiaomiMiMo/MiMo-7B-RL} + +# Fully-async split-placement layout: trainer group + rollout group. +TRAIN_NNODES=${TRAIN_NNODES:-1} +TRAIN_NGPUS_PER_NODE=${TRAIN_NGPUS_PER_NODE:-4} +ROLLOUT_NNODES=${ROLLOUT_NNODES:-1} +ROLLOUT_NGPUS_PER_NODE=${ROLLOUT_NGPUS_PER_NODE:-4} + +max_prompt_length=${MAX_PROMPT_LENGTH:-2048} +max_response_length=${MAX_RESPONSE_LENGTH:-8192} +gen_batch_size=${GEN_BATCH_SIZE:-1} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-32} +n_resp_per_prompt=${N_RESP_PER_PROMPT:-16} + +clip_ratio_low=${CLIP_RATIO_LOW:-0.2} +clip_ratio_high=${CLIP_RATIO_HIGH:-0.28} + +mtp_loss_scaling_factor=${MTP_LOSS_SCALING_FACTOR:-0.1} + +actor_tp=${ACTOR_TP:-2} +actor_pp=${ACTOR_PP:-1} +actor_cp=${ACTOR_CP:-1} + +rollout_tp=${ROLLOUT_TP:-2} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.8} + +staleness_threshold=${STALENESS_THRESHOLD:-0.5} +trigger_parameter_sync_step=${TRIGGER_PARAMETER_SYNC_STEP:-4} +require_batches=${REQUIRE_BATCHES:-1} + +total_rollout_steps=${TOTAL_ROLLOUT_STEPS:-51200} +total_epochs=${TOTAL_EPOCHS:-10} +test_freq=${TEST_FREQ:-10} +save_freq=${SAVE_FREQ:--1} + +project_name=${PROJECT_NAME:-verl_mtp_fully_async} +experiment_name=${EXPERIMENT_NAME:-dapo_mimo_7b_mtp_fully_async} + +RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"} +CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${experiment_name}"} +TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/dapo-math-17k.parquet"} +TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/aime-2024.parquet"} +# ---- end user-adjustable ---- + +actor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2)) +infer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3)) + +python -m verl.experimental.fully_async_policy.fully_async_main \ + --config-path=config \ + --config-name='fully_async_ppo_megatron_trainer.yaml' \ + data.train_files="${TRAIN_FILE}" \ + data.val_files="${TEST_FILE}" \ + data.prompt_key=prompt \ + data.truncation='left' \ + data.max_prompt_length=${max_prompt_length} \ + data.max_response_length=${max_response_length} \ + data.train_batch_size=0 \ + data.gen_batch_size=${gen_batch_size} \ + data.trust_remote_code=True \ + algorithm.adv_estimator=grpo \ + algorithm.use_kl_in_reward=False \ + algorithm.kl_ctrl.kl_coef=0.0 \ + actor_rollout_ref.model.path="${MODEL_PATH}" \ + actor_rollout_ref.model.trust_remote_code=True \ + actor_rollout_ref.model.mtp.enable=True \ + actor_rollout_ref.model.mtp.enable_train=True \ + actor_rollout_ref.model.mtp.mtp_loss_scaling_factor=${mtp_loss_scaling_factor} \ + actor_rollout_ref.model.mtp.detach_encoder=True \ + actor_rollout_ref.model.mtp.enable_rollout=True \ + actor_rollout_ref.hybrid_engine=False \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.actor.optim.lr_warmup_steps=10 \ + actor_rollout_ref.actor.optim.lr_decay_steps=${total_rollout_steps} \ + actor_rollout_ref.actor.optim.weight_decay=0.1 \ + actor_rollout_ref.actor.optim.clip_grad=1.0 \ + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} \ + actor_rollout_ref.actor.use_dynamic_bsz=True \ + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \ + actor_rollout_ref.actor.use_kl_loss=False \ + actor_rollout_ref.actor.entropy_coeff=0 \ + actor_rollout_ref.actor.loss_agg_mode=token-mean \ + actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \ + actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \ + actor_rollout_ref.actor.clip_ratio_c=10.0 \ + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${actor_tp} \ + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${actor_pp} \ + actor_rollout_ref.actor.megatron.context_parallel_size=${actor_cp} \ + actor_rollout_ref.actor.megatron.param_offload=False \ + actor_rollout_ref.actor.megatron.grad_offload=False \ + actor_rollout_ref.actor.megatron.optimizer_offload=False \ + actor_rollout_ref.actor.megatron.use_mbridge=True \ + actor_rollout_ref.rollout.name=sglang \ + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} \ + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} \ + actor_rollout_ref.rollout.n=${n_resp_per_prompt} \ + actor_rollout_ref.rollout.calculate_log_probs=True \ + actor_rollout_ref.rollout.enable_chunked_prefill=True \ + actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \ + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \ + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \ + actor_rollout_ref.rollout.val_kwargs.temperature=1.0 \ + actor_rollout_ref.rollout.val_kwargs.top_p=0.7 \ + actor_rollout_ref.rollout.val_kwargs.top_k=-1 \ + actor_rollout_ref.rollout.val_kwargs.do_sample=True \ + actor_rollout_ref.rollout.val_kwargs.n=1 \ + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True \ + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \ + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${actor_tp} \ + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${actor_pp} \ + actor_rollout_ref.ref.megatron.context_parallel_size=${actor_cp} \ + actor_rollout_ref.ref.megatron.param_offload=False \ + reward_model.reward_manager=dapo \ + +reward_model.reward_kwargs.overlong_buffer_cfg.enable=True \ + +reward_model.reward_kwargs.overlong_buffer_cfg.len=$((1024 * 4)) \ + +reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=1.0 \ + +reward_model.reward_kwargs.overlong_buffer_cfg.log=False \ + +reward_model.reward_kwargs.max_resp_len=${max_response_length} \ + rollout.total_rollout_steps=${total_rollout_steps} \ + rollout.nnodes=${ROLLOUT_NNODES} \ + rollout.n_gpus_per_node=${ROLLOUT_NGPUS_PER_NODE} \ + async_training.staleness_threshold=${staleness_threshold} \ + async_training.trigger_parameter_sync_step=${trigger_parameter_sync_step} \ + async_training.require_batches=${require_batches} \ + async_training.partial_rollout=True \ + trainer.balance_batch=True \ + trainer.critic_warmup=0 \ + trainer.logger='["console","wandb"]' \ + trainer.project_name="${project_name}" \ + trainer.experiment_name="${experiment_name}" \ + trainer.nnodes=${TRAIN_NNODES} \ + trainer.n_gpus_per_node=${TRAIN_NGPUS_PER_NODE} \ + trainer.val_before_train=False \ + trainer.save_freq=${save_freq} \ + trainer.test_freq=${test_freq} \ + trainer.total_epochs=${total_epochs} \ + trainer.default_local_dir="${CKPTS_DIR}" \ + trainer.resume_mode=auto \ + trainer.log_val_generations=10 "$@" diff --git a/verl/examples/mtp_trainer/run_mimo_7b_mtp_megatron.sh b/verl/examples/mtp_trainer/run_mimo_7b_mtp_megatron.sh new file mode 100644 index 0000000000000000000000000000000000000000..5368b9812cff9b0ffc78261104adb267ebc57633 --- /dev/null +++ b/verl/examples/mtp_trainer/run_mimo_7b_mtp_megatron.sh @@ -0,0 +1,150 @@ +#!/usr/bin/env bash +# MTP | MiMo-7B (speculative MTP head) | SGLang rollout | Megatron training | NVIDIA GPUs +# Multi-token-prediction training flow for MiMo-style models. + +set -xeuo pipefail + +export CUDA_DEVICE_MAX_CONNECTIONS=1 + +# ---- user-adjustable ---- +# NOTE: remember to set max_position_embeddings=32768 in the model's config.json after downloading. +MODEL_PATH=${MODEL_PATH:-XiaomiMiMo/MiMo-7B-RL} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +train_batch_size=${TRAIN_BATCH_SIZE:-128} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-32} +max_prompt_length=${MAX_PROMPT_LENGTH:-2048} +max_response_length=${MAX_RESPONSE_LENGTH:-8192} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-20480} + +actor_lr=${ACTOR_LR:-1e-6} +entropy_coeff=${ENTROPY_COEFF:-0} + +clip_ratio_low=${CLIP_RATIO_LOW:-0.2} +clip_ratio_high=${CLIP_RATIO_HIGH:-0.28} + +mtp_loss_scaling_factor=${MTP_LOSS_SCALING_FACTOR:-0.1} + +actor_tp=${ACTOR_TP:-2} +actor_pp=${ACTOR_PP:-2} +actor_cp=${ACTOR_CP:-2} + +rollout_tp=${ROLLOUT_TP:-4} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.8} +rollout_n=${ROLLOUT_N:-16} + +total_epochs=${TOTAL_EPOCHS:-10} +total_training_steps=${TOTAL_TRAINING_STEPS:-400} +save_freq=${SAVE_FREQ:--1} +test_freq=${TEST_FREQ:-10} + +project_name=${PROJECT_NAME:-verl_mtp} +experiment_name=${EXPERIMENT_NAME:-mimo_7b_mtp_sglang_megatron} +# ---- end user-adjustable ---- + +train_file=${TRAIN_FILE:-$HOME/data/dapo-math-17k/train.parquet} +val_file=${VAL_FILE:-$HOME/data/aime-2024/test.parquet} +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files="['$train_file']" + data.val_files="['$val_file']" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.truncation='left' + data.trust_remote_code=True +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.trust_remote_code=True + actor_rollout_ref.model.mtp.enable=True + actor_rollout_ref.model.mtp.enable_train=True + actor_rollout_ref.model.mtp.mtp_loss_scaling_factor=${mtp_loss_scaling_factor} + actor_rollout_ref.model.mtp.detach_encoder=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.optim.lr_warmup_steps=10 + actor_rollout_ref.actor.optim.weight_decay=0.1 + actor_rollout_ref.actor.optim.clip_grad=1.0 + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=False + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.loss_agg_mode=token-mean + actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} + actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} + actor_rollout_ref.actor.clip_ratio_c=10.0 + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${actor_pp} + actor_rollout_ref.actor.megatron.context_parallel_size=${actor_cp} + actor_rollout_ref.actor.megatron.param_offload=True + actor_rollout_ref.actor.megatron.grad_offload=True + actor_rollout_ref.actor.megatron.optimizer_offload=True + actor_rollout_ref.actor.megatron.use_mbridge=True +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=sglang + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${actor_pp} + actor_rollout_ref.ref.megatron.context_parallel_size=${actor_cp} + actor_rollout_ref.ref.megatron.param_offload=True +) + +REWARD=( + reward.reward_manager.name=dapo + +reward.reward_kwargs.overlong_buffer_cfg.enable=True + +reward.reward_kwargs.overlong_buffer_cfg.len=4096 + +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=1.0 + +reward.reward_kwargs.max_resp_len=${max_response_length} +) + +TRAINER=( + trainer.balance_batch=True + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.val_before_train=False + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} + trainer.total_training_steps=${total_training_steps} +) + +EXTRA=( + model_engine=megatron +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${REWARD[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/on_policy_distillation_trainer/README.md b/verl/examples/on_policy_distillation_trainer/README.md new file mode 100644 index 0000000000000000000000000000000000000000..d3cbce25ebc967b9db7dae37d1415fd4c5499e2f --- /dev/null +++ b/verl/examples/on_policy_distillation_trainer/README.md @@ -0,0 +1,24 @@ +# On-Policy Distillation + +This trainer jointly trains a student model with policy-gradient on-policy rollouts and a distillation loss against a frozen teacher model served by a separate Ray cluster. Compared to pure SFT from teacher generations, on-policy distillation typically closes more of the teacher/student gap at the same compute budget. + +## Canonical Scripts + +| Script | Teachers | Modality | Infer | Train | Platform | +|---------------------------------|----------|------------|-------|----------|----------| +| `run_qwen3_8b_fsdp.sh` | single | text | vLLM | FSDP | NVIDIA | +| `run_qwen3_8b_megatron.sh` | single | text | vLLM | Megatron | NVIDIA | +| `run_qwen3_vl_8b_fsdp.sh` | single | VL | vLLM | FSDP | NVIDIA | +| `run_qwen3_8b_mopd_fsdp.sh` | multi | text + VL | vLLM | FSDP | NVIDIA | + +Override `STUDENT_MODEL` and `TEACHER_MODEL` via env vars to swap model pairs in +the single-teacher scripts. The MOPD script exposes per-teacher overrides. + +## Key Flags + +- `distillation.enabled=True` +- `distillation.teacher_models.teacher_model.model_path=` (single-teacher) +- `+distillation.teacher_models..{key,model_path,num_replicas,inference.*}` (multi-teacher) +- `distillation.distillation_loss.loss_mode={k1, k3, forward_kl_topk, ...}` +- `distillation.distillation_loss.use_policy_gradient=True|False` +- `distillation.distillation_loss.topk=64` diff --git a/verl/examples/on_policy_distillation_trainer/run_qwen3_8b_fsdp.sh b/verl/examples/on_policy_distillation_trainer/run_qwen3_8b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..befe0147fae3f7c0c4f5562063d760b7e9a5b575 --- /dev/null +++ b/verl/examples/on_policy_distillation_trainer/run_qwen3_8b_fsdp.sh @@ -0,0 +1,127 @@ +#!/usr/bin/env bash +# On-policy distillation | text | vLLM rollout | FSDP training | NVIDIA GPUs + +set -xeuo pipefail + +# ---- user-adjustable ---- +STUDENT_MODEL=${STUDENT_MODEL:-Qwen/Qwen3-8B} +TEACHER_MODEL=${TEACHER_MODEL:-Qwen/Qwen3-32B} + +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} +TEACHER_WORLD_SIZE=${TEACHER_WORLD_SIZE:-4} + +distillation_loss_mode=${DISTILLATION_LOSS_MODE:-k1} +use_policy_gradient=${USE_POLICY_GRADIENT:-True} +distillation_topk=${DISTILLATION_TOPK:-64} + +train_batch_size=${TRAIN_BATCH_SIZE:-128} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-128} +max_prompt_length=${MAX_PROMPT_LENGTH:-1024} +max_response_length=${MAX_RESPONSE_LENGTH:-2048} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-24576} + +actor_lr=${ACTOR_LR:-1e-6} + +rollout_tp=${ROLLOUT_TP:-2} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.4} +teacher_tp=${TEACHER_TP:-2} +teacher_gpu_mem_util=${TEACHER_GPU_MEM_UTIL:-0.4} + +total_epochs=${TOTAL_EPOCHS:-15} +save_freq=${SAVE_FREQ:-200} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_distill_gsm8k_math} +experiment_name=${EXPERIMENT_NAME:-qwen3_8b_from_qwen3_32b_vllm_fsdp} +# ---- end user-adjustable ---- + +gsm8k_train=$HOME/data/gsm8k/train.parquet +gsm8k_test=$HOME/data/gsm8k/test.parquet +math_train=$HOME/data/math/train.parquet +math_test=$HOME/data/math/test.parquet + +train_files="['$gsm8k_train', '$math_train']" +val_files="['$gsm8k_test', '$math_test']" + +max_num_tokens=$(( max_prompt_length + max_response_length + 1 )) +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files="$train_files" + data.val_files="$val_files" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' + data.shuffle=False +) + +MODEL=( + actor_rollout_ref.model.path="$STUDENT_MODEL" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.use_torch_compile=True + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.fsdp_config.param_offload=True + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=1 + actor_rollout_ref.rollout.max_model_len=${max_num_tokens} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +TRAINER=( + trainer.balance_batch=True + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.val_before_train=False + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( + distillation.enabled=True + distillation.n_gpus_per_node=${TEACHER_WORLD_SIZE} + distillation.nnodes=${NNODES} + distillation.teacher_models.teacher_model.model_path="$TEACHER_MODEL" + distillation.teacher_models.teacher_model.inference.tensor_model_parallel_size=${teacher_tp} + distillation.teacher_models.teacher_model.inference.name=vllm + distillation.teacher_models.teacher_model.inference.gpu_memory_utilization=${teacher_gpu_mem_util} + distillation.teacher_models.teacher_model.inference.max_model_len=${max_num_tokens} + distillation.distillation_loss.loss_mode=${distillation_loss_mode} + distillation.distillation_loss.topk=${distillation_topk} + distillation.distillation_loss.use_task_rewards=False + distillation.distillation_loss.use_policy_gradient=${use_policy_gradient} + distillation.distillation_loss.loss_max_clamp=10.0 + distillation.distillation_loss.log_prob_min_clamp=-10.0 +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/on_policy_distillation_trainer/run_qwen3_8b_megatron.sh b/verl/examples/on_policy_distillation_trainer/run_qwen3_8b_megatron.sh new file mode 100644 index 0000000000000000000000000000000000000000..0328f23f6298f1342b90622450be522d14cd9a8d --- /dev/null +++ b/verl/examples/on_policy_distillation_trainer/run_qwen3_8b_megatron.sh @@ -0,0 +1,132 @@ +#!/usr/bin/env bash +# On-policy distillation | text | vLLM rollout | Megatron training | NVIDIA GPUs + +set -xeuo pipefail + +export CUDA_DEVICE_MAX_CONNECTIONS=1 + +# ---- user-adjustable ---- +STUDENT_MODEL=${STUDENT_MODEL:-Qwen/Qwen3-8B} +TEACHER_MODEL=${TEACHER_MODEL:-Qwen/Qwen3-32B} + +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} +TEACHER_WORLD_SIZE=${TEACHER_WORLD_SIZE:-4} + +distillation_loss_mode=${DISTILLATION_LOSS_MODE:-forward_kl_topk} +use_policy_gradient=${USE_POLICY_GRADIENT:-False} +distillation_topk=${DISTILLATION_TOPK:-64} + +train_batch_size=${TRAIN_BATCH_SIZE:-128} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-128} +max_prompt_length=${MAX_PROMPT_LENGTH:-1024} +max_response_length=${MAX_RESPONSE_LENGTH:-2048} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-24576} + +actor_lr=${ACTOR_LR:-1e-6} + +actor_tp=${ACTOR_TP:-2} +actor_pp=${ACTOR_PP:-1} + +rollout_tp=${ROLLOUT_TP:-2} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.4} +teacher_tp=${TEACHER_TP:-2} +teacher_gpu_mem_util=${TEACHER_GPU_MEM_UTIL:-0.4} + +total_epochs=${TOTAL_EPOCHS:-15} +save_freq=${SAVE_FREQ:-200} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_distill_gsm8k_math} +experiment_name=${EXPERIMENT_NAME:-qwen3_8b_from_qwen3_32b_vllm_megatron} +# ---- end user-adjustable ---- + +gsm8k_train=$HOME/data/gsm8k/train.parquet +gsm8k_test=$HOME/data/gsm8k/test.parquet +math_train=$HOME/data/math/train.parquet +math_test=$HOME/data/math/test.parquet + +train_files="['$gsm8k_train', '$math_train']" +val_files="['$gsm8k_test', '$math_test']" + +max_num_tokens=$(( max_prompt_length + max_response_length + 1 )) +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files="$train_files" + data.val_files="$val_files" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$STUDENT_MODEL" + actor_rollout_ref.model.use_remove_padding=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${actor_pp} + actor_rollout_ref.actor.megatron.param_offload=True + actor_rollout_ref.actor.megatron.optimizer_offload=True +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=1 + actor_rollout_ref.rollout.max_model_len=${max_num_tokens} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +TRAINER=( + trainer.balance_batch=True + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.val_before_train=False + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( + model_engine=megatron + distillation.enabled=True + distillation.n_gpus_per_node=${TEACHER_WORLD_SIZE} + distillation.nnodes=${NNODES} + distillation.teacher_models.teacher_model.model_path="$TEACHER_MODEL" + distillation.teacher_models.teacher_model.inference.tensor_model_parallel_size=${teacher_tp} + distillation.teacher_models.teacher_model.inference.name=vllm + distillation.teacher_models.teacher_model.inference.gpu_memory_utilization=${teacher_gpu_mem_util} + distillation.teacher_models.teacher_model.inference.max_model_len=${max_num_tokens} + distillation.distillation_loss.loss_mode=${distillation_loss_mode} + distillation.distillation_loss.topk=${distillation_topk} + distillation.distillation_loss.use_task_rewards=False + distillation.distillation_loss.use_policy_gradient=${use_policy_gradient} + distillation.distillation_loss.loss_max_clamp=10.0 + distillation.distillation_loss.log_prob_min_clamp=-10.0 +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/on_policy_distillation_trainer/run_qwen3_8b_mopd_fsdp.sh b/verl/examples/on_policy_distillation_trainer/run_qwen3_8b_mopd_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..36bff7f48caaabb802e74b5dbf4807c06ed99413 --- /dev/null +++ b/verl/examples/on_policy_distillation_trainer/run_qwen3_8b_mopd_fsdp.sh @@ -0,0 +1,150 @@ +#!/usr/bin/env bash +# On-policy distillation | multi-teacher (gsm8k text + geo3k VL) | vLLM rollout | FSDP training | NVIDIA GPUs + +set -xeuo pipefail + +# ---- user-adjustable ---- +STUDENT_MODEL=${STUDENT_MODEL:-Qwen/Qwen3-VL-8B-Instruct} +GSM8K_TEACHER_MODEL=${GSM8K_TEACHER_MODEL:-Qwen/Qwen3-32B} +GEO3K_TEACHER_MODEL=${GEO3K_TEACHER_MODEL:-Qwen/Qwen3-VL-32B-Instruct} + +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +# Per-teacher replicas; total teacher GPUs = sum(num_replicas) * teacher_tp +TEACHER_NNODES=${TEACHER_NNODES:-1} +TEACHER_NUM_REPLICAS_GSM8K=${TEACHER_NUM_REPLICAS_GSM8K:-1} +TEACHER_NUM_REPLICAS_GEO3K=${TEACHER_NUM_REPLICAS_GEO3K:-1} +teacher_tp=${TEACHER_TP:-2} +TEACHER_WORLD_SIZE=$(( (TEACHER_NUM_REPLICAS_GSM8K + TEACHER_NUM_REPLICAS_GEO3K) * teacher_tp )) + +distillation_loss_mode=${DISTILLATION_LOSS_MODE:-k1} +use_policy_gradient=${USE_POLICY_GRADIENT:-True} +distillation_topk=${DISTILLATION_TOPK:-64} + +train_batch_size=${TRAIN_BATCH_SIZE:-128} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-128} +max_prompt_length=${MAX_PROMPT_LENGTH:-1024} +max_response_length=${MAX_RESPONSE_LENGTH:-2048} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-24576} + +actor_lr=${ACTOR_LR:-1e-6} + +rollout_tp=${ROLLOUT_TP:-2} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.4} +teacher_gpu_mem_util=${TEACHER_GPU_MEM_UTIL:-0.4} + +total_epochs=${TOTAL_EPOCHS:-15} +save_freq=${SAVE_FREQ:-200} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_distill_mopd_gsm8k_geo3k} +experiment_name=${EXPERIMENT_NAME:-qwen3_vl_8b_from_qwen3_32b_and_qwen3_vl_32b_mopd_vllm_fsdp} +# ---- end user-adjustable ---- + +gsm8k_train=$HOME/data/gsm8k/train.parquet +gsm8k_test=$HOME/data/gsm8k/test.parquet +geo3k_train=$HOME/data/geo3k/train.parquet +geo3k_test=$HOME/data/geo3k/test.parquet + +train_files="['$gsm8k_train', '$geo3k_train']" +val_files="['$gsm8k_test', '$geo3k_test']" + +max_num_tokens=$(( max_prompt_length + max_response_length + 1 )) +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files="$train_files" + data.val_files="$val_files" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' + data.shuffle=True + data.image_key=images +) + +MODEL=( + actor_rollout_ref.model.path="$STUDENT_MODEL" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.use_torch_compile=True + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.fsdp_config.param_offload=True + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=1 + actor_rollout_ref.rollout.max_model_len=${max_num_tokens} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +TRAINER=( + trainer.balance_batch=True + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.val_before_train=False + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +# Multi-teacher: one teacher per dataset, routed by the sample's `data_source` value. +# Use `+distillation.teacher_models..*` to add named teachers; the default `teacher_model` +# entry is silently popped when other teacher entries are added. +EXTRA=( + distillation.enabled=True + distillation.n_gpus_per_node=${TEACHER_WORLD_SIZE} + distillation.nnodes=${TEACHER_NNODES} + distillation.teacher_key=data_source + # --- gsm8k teacher (text) --- + +distillation.teacher_models.gsm8k.key="openai/gsm8k" + +distillation.teacher_models.gsm8k.model_path="$GSM8K_TEACHER_MODEL" + +distillation.teacher_models.gsm8k.num_replicas=${TEACHER_NUM_REPLICAS_GSM8K} + +distillation.teacher_models.gsm8k.inference.name=vllm + +distillation.teacher_models.gsm8k.inference.tensor_model_parallel_size=${teacher_tp} + +distillation.teacher_models.gsm8k.inference.gpu_memory_utilization=${teacher_gpu_mem_util} + +distillation.teacher_models.gsm8k.inference.max_model_len=${max_num_tokens} + # --- geo3k teacher (VL) --- + +distillation.teacher_models.geo3k.key="hiyouga/geometry3k" + +distillation.teacher_models.geo3k.model_path="$GEO3K_TEACHER_MODEL" + +distillation.teacher_models.geo3k.num_replicas=${TEACHER_NUM_REPLICAS_GEO3K} + +distillation.teacher_models.geo3k.inference.name=vllm + +distillation.teacher_models.geo3k.inference.tensor_model_parallel_size=${teacher_tp} + +distillation.teacher_models.geo3k.inference.gpu_memory_utilization=${teacher_gpu_mem_util} + +distillation.teacher_models.geo3k.inference.max_model_len=${max_num_tokens} + # --- loss --- + distillation.distillation_loss.loss_mode=${distillation_loss_mode} + distillation.distillation_loss.topk=${distillation_topk} + distillation.distillation_loss.use_task_rewards=False + distillation.distillation_loss.use_policy_gradient=${use_policy_gradient} + distillation.distillation_loss.loss_max_clamp=10.0 + distillation.distillation_loss.log_prob_min_clamp=-10.0 +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/on_policy_distillation_trainer/run_qwen3_vl_8b_fsdp.sh b/verl/examples/on_policy_distillation_trainer/run_qwen3_vl_8b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..7edbf258c348159a89dff12dac9a35e06019e5f8 --- /dev/null +++ b/verl/examples/on_policy_distillation_trainer/run_qwen3_vl_8b_fsdp.sh @@ -0,0 +1,121 @@ +#!/usr/bin/env bash +# On-policy distillation | vision (Qwen3-VL) | vLLM rollout | FSDP training | NVIDIA GPUs + +set -xeuo pipefail + +# ---- user-adjustable ---- +STUDENT_MODEL=${STUDENT_MODEL:-Qwen/Qwen3-VL-8B-Instruct} +TEACHER_MODEL=${TEACHER_MODEL:-Qwen/Qwen3-VL-32B-Instruct} + +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} +TEACHER_WORLD_SIZE=${TEACHER_WORLD_SIZE:-4} + +distillation_loss_mode=${DISTILLATION_LOSS_MODE:-k1} +use_policy_gradient=${USE_POLICY_GRADIENT:-True} +distillation_topk=${DISTILLATION_TOPK:-64} + +train_batch_size=${TRAIN_BATCH_SIZE:-128} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-128} +max_prompt_length=${MAX_PROMPT_LENGTH:-1024} +max_response_length=${MAX_RESPONSE_LENGTH:-2048} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-24576} + +actor_lr=${ACTOR_LR:-1e-6} + +rollout_tp=${ROLLOUT_TP:-2} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.4} +teacher_tp=${TEACHER_TP:-2} +teacher_gpu_mem_util=${TEACHER_GPU_MEM_UTIL:-0.4} + +total_epochs=${TOTAL_EPOCHS:-15} +save_freq=${SAVE_FREQ:-200} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_distill_geo3k} +experiment_name=${EXPERIMENT_NAME:-qwen3_vl_8b_from_qwen3_vl_32b_vllm_fsdp} +# ---- end user-adjustable ---- + +train_file=${TRAIN_FILE:-$HOME/data/geo3k/train.parquet} +val_file=${VAL_FILE:-$HOME/data/geo3k/test.parquet} + +max_num_tokens=$(( max_prompt_length + max_response_length + 1 )) +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files="['$train_file']" + data.val_files="['$val_file']" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' + data.image_key=images +) + +MODEL=( + actor_rollout_ref.model.path="$STUDENT_MODEL" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.fsdp_config.param_offload=True + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=1 + actor_rollout_ref.rollout.max_model_len=${max_num_tokens} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +TRAINER=( + trainer.balance_batch=True + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.val_before_train=False + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( + distillation.enabled=True + distillation.n_gpus_per_node=${TEACHER_WORLD_SIZE} + distillation.nnodes=${NNODES} + distillation.teacher_models.teacher_model.model_path="$TEACHER_MODEL" + distillation.teacher_models.teacher_model.inference.tensor_model_parallel_size=${teacher_tp} + distillation.teacher_models.teacher_model.inference.name=vllm + distillation.teacher_models.teacher_model.inference.gpu_memory_utilization=${teacher_gpu_mem_util} + distillation.teacher_models.teacher_model.inference.max_model_len=${max_num_tokens} + distillation.distillation_loss.loss_mode=${distillation_loss_mode} + distillation.distillation_loss.topk=${distillation_topk} + distillation.distillation_loss.use_task_rewards=False + distillation.distillation_loss.use_policy_gradient=${use_policy_gradient} + distillation.distillation_loss.loss_max_clamp=10.0 + distillation.distillation_loss.log_prob_min_clamp=-10.0 +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/otb_trainer/README.md b/verl/examples/otb_trainer/README.md new file mode 100644 index 0000000000000000000000000000000000000000..94cbea059faace34f76ecf518fa6094af0a0394b --- /dev/null +++ b/verl/examples/otb_trainer/README.md @@ -0,0 +1,16 @@ +# Optimal Token Baseline (OTB) + +OTB uses a token-wise optimal variance-reduction baseline derived from the per-prompt sample group. See `tir_optimal_token_baseline` variant for a TIR-specific form. + +## Canonical Scripts + +| Script | Infer | Train | Platform | +|--------------------------------------|-------|-------|----------| +| `run_qwen3_8b_fsdp.sh` | vLLM | FSDP | NVIDIA | + +Set `ADV_ESTIMATOR=tir_optimal_token_baseline` for the TIR variant. + +## Key Flags + +- `algorithm.adv_estimator=optimal_token_baseline` +- `actor_rollout_ref.actor.calculate_sum_pi_squared=True` diff --git a/verl/examples/otb_trainer/run_qwen3_8b_fsdp.sh b/verl/examples/otb_trainer/run_qwen3_8b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..fae4c2360fa4bf877995f9494d25070ba94c3947 --- /dev/null +++ b/verl/examples/otb_trainer/run_qwen3_8b_fsdp.sh @@ -0,0 +1,113 @@ +#!/usr/bin/env bash +# OTB (Optimal Token Baseline) | text | vLLM rollout | FSDP training | NVIDIA GPUs + +set -xeuo pipefail + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-8B} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +adv_estimator=${ADV_ESTIMATOR:-optimal_token_baseline} + +train_batch_size=${TRAIN_BATCH_SIZE:-128} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-128} +max_prompt_length=${MAX_PROMPT_LENGTH:-1024} +max_response_length=${MAX_RESPONSE_LENGTH:-2048} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-24576} + +actor_lr=${ACTOR_LR:-1e-6} +entropy_coeff=${ENTROPY_COEFF:-0} + +rollout_tp=${ROLLOUT_TP:-2} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.75} +rollout_n=${ROLLOUT_N:-8} + +total_epochs=${TOTAL_EPOCHS:-15} +save_freq=${SAVE_FREQ:-20} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_otb_gsm8k_math} +experiment_name=${EXPERIMENT_NAME:-qwen3_8b_vllm_fsdp} +# ---- end user-adjustable ---- + +gsm8k_train=$HOME/data/gsm8k/train.parquet +gsm8k_test=$HOME/data/gsm8k/test.parquet +math_train=$HOME/data/math/train.parquet +math_test=$HOME/data/math/test.parquet + +train_files="['$gsm8k_train', '$math_train']" +val_files="['$gsm8k_test', '$math_test']" +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=${adv_estimator} + algorithm.use_kl_in_reward=False + data.train_files="$train_files" + data.val_files="$val_files" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=False + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.calculate_sum_pi_squared=True + actor_rollout_ref.actor.fsdp_config.param_offload=False + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.fsdp_config.param_offload=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/ppo_trainer/README.md b/verl/examples/ppo_trainer/README.md new file mode 100644 index 0000000000000000000000000000000000000000..c5f16fffeccd730a493c9aace20388733e5a4b62 --- /dev/null +++ b/verl/examples/ppo_trainer/README.md @@ -0,0 +1,99 @@ +# Proximal Policy Optimization (PPO) + +Proximal Policy Optimization (PPO) is a family of policy gradient methods for reinforcement learning, proposed by OpenAI in 2017. PPO strikes a balance between simplicity, stability, and performance, making it one of the most widely used algorithms in modern RL applications, including large-scale language model fine-tuning. + +Traditional policy gradient methods like REINFORCE or Vanilla Policy Gradient suffer from: + +- High variance and sample inefficiency. +- Instability due to large policy updates. + +PPO addresses this problem using a clipped surrogate objective that avoids overly large updates without requiring second-order derivatives. + +For more technical details regarding PPO, we suggest reading the introduction in the [OpenAI spinning up tutorial](https://spinningup.openai.com/en/latest/algorithms/ppo.html), and the paper [Proximal Policy Optimization Algorithms](https://arxiv.org/abs/1707.06347). + +## Key Components + +- Actor-Critic Architecture: PPO requires both an actor model (policy) and a critic model (value function). This differs from other algorithms like GRPO and RLOO that don't require a critic model. + +- Generalized Advantage Estimation (GAE): PPO uses GAE for computing advantage values, which helps reduce variance in policy gradient estimates while maintaining low bias. + +- Clipped Surrogate Objective: The core of PPO is implemented through the clipped surrogate objective function that limits policy updates. + +## Configuration + +Note that all configs containing `micro_batch_size` are used to configure the maximum sample or token count per forward or backward pass to avoid GPU OOMs, whose value should not change algorithmic/convergence behavior. + +Most critic configs are similar to those of actors. Note that the critic model is omitted from the figure below. + +![image](https://github.com/user-attachments/assets/16aebad1-0da6-4eb3-806d-54a74e712c2d) + +- `data.train_batch_size`: The global batch size of prompts used to generate a set of sampled trajectories/rollouts. The number of responses/trajectories is `data.train_batch_size * actor_rollout.ref.rollout.n` + +- `actor_rollout_ref.actor.ppo_mini_batch_size`: The set of sampled trajectories is split into multiple mini-batches with batch_size=ppo_mini_batch_size for PPO actor updates. The ppo_mini_batch_size is a global size across all workers + +- `critic.ppo_mini_batch_size`: The set of sampled trajectories is split into multiple mini-batches with batch_size=ppo_mini_batch_size for PPO critic updates. The ppo_mini_batch_size is a global size across all workers + +- `actor_rollout_ref.actor.clip_ratio`: The PPO clip range. Default to 0.2 + +- `actor_rollout_ref.actor.ppo_epochs`: Number of epochs for PPO updates on one set of sampled trajectories for actor + +- `critic.ppo_epochs`: Number of epochs for PPO updates on one set of sampled trajectories for critic. Defaults to `actor_rollout_ref.actor.ppo_epochs` + +- `algorithm.gamma`: discount factor + +- `algorithm.lam`: The lambda term that trades off between bias and variance in the GAE estimator + +- `algorithm.adv_estimator`: Support gae, grpo, reinforce_plus_plus, reinforce_plus_plus_baseline, rloo, rloo_vectorized + +## Advanced Extensions + +### KL Divergence Control + +Options to prevent the policy from diverging too far from a reference policy. Two mechanisms are available: KL reward penalty and KL loss. For more technical details, see [Training language models to follow instructions with human feedback](https://arxiv.org/abs/2203.02155) + +Options to use KL loss for KL divergence control: + +- `actor_rollout_ref.actor.use_kl_loss`: to use kl loss in the actor. When used, we are not applying KL in the reward function. Default is False + +- `actor_rollout_ref.actor.kl_loss_coef`: The coefficient of kl loss. Default is 0.001. + +- `actor_rollout_ref.actor.kl_loss_type`: Support kl(k1), abs, mse(k2), low_var_kl(k3) and full. Appending "+" in the end (e.g., 'k1+' and 'k3+') would apply straight through to employ k2 for unbiased gradient estimation, regardless of the kl value estimation (see https://github.com/verl-project/verl/pull/2953#issuecomment-3162113848 for more details). How to calculate the kl divergence between actor and reference policy. See this blog post for detailed analysis: http://joschu.net/blog/kl-approx.html + +Options to use KL penalty in the reward: + +- `algorithm.use_kl_in_reward`: Whether to enable in-reward kl penalty. Default is False. + +- `algorithm.kl_penalty`: Support kl(k1), abs, mse(k2), low_var_kl(k3) and full. This defines the way to calculate the kl divergence between actor and reference policy. For specific options, refer to `kl_penalty` in core_algos.py. See this blog post for detailed analysis: http://joschu.net/blog/kl-approx.html + +- `algorithm.kl_ctrl.kl_coef`: The (initial) coefficient of in-reward kl_penalty. Default is 0.001. +- `algorithm.kl_ctrl.type`: 'fixed' for FixedKLController and 'adaptive' for AdaptiveKLController. +- `algorithm.kl_ctrl.horizon`: See source code of AdaptiveKLController for details. +- `algorithm.kl_ctrl.target_kl`: See source code of AdaptiveKLController for details. + +### Dual-clip PPO + +The Dual-Clip PPO introduces a approach by applying a lower bound to the policy ratio when the advantage is less than zero, when multiplied by a large raito, does not exceed a specified lower bound. + +![image](https://github.com/user-attachments/assets/fc232181-d8b0-4307-8dd2-4dc0a4c1c139) + +- `actor_rollout_ref.actor.clip_ratio_c`: lower bound of the value for Dual-clip PPO, defaults to 3.0 + +## Canonical Scripts + +All scripts follow the `run___[_].sh` naming convention, use `MODEL_PATH` as an env var (override like `MODEL_PATH=Qwen/Qwen3-14B bash run_qwen3_8b_fsdp.sh`), enable dynamic batch size and batch balancing by default, and only use current-API Hydra overrides. + +| Script | Infer | Train | Platform | +|-----------------------------------------|--------|----------|----------| +| `run_qwen3_8b_fsdp.sh` | vLLM | FSDP | NVIDIA | +| `run_qwen3_8b_fsdp.sh` | SGLang | FSDP | NVIDIA | +| `run_qwen3_8b_megatron.sh` | vLLM | Megatron | NVIDIA | +| `run_qwen3_8b_fsdp.sh` | vLLM | FSDP | Ascend | + +Default dataset: GSM8K + MATH. Override `TRAIN_BATCH_SIZE`, `PPO_MINI_BATCH_SIZE`, `ACTOR_LR`, `CRITIC_LR`, `ROLLOUT_TP`, `NNODES`, `NGPUS_PER_NODE`, etc. via env vars listed at the top of each script. + +Reference performance with verl v0.2: + +| Model | Method | Score | Link | +|-------------------------------|------------------|-------|------------------------------------------------------------------------------------------------| +| Qwen/Qwen2.5-0.5B-Instruct | pretrained model | 36.4 | [Qwen Blog](https://qwenlm.github.io/blog/qwen2.5-llm/) | +| Qwen/Qwen2.5-0.5B-Instruct | PPO | 56.7 | [PPO Command and Logs](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/Qwen2.5-0.5B-bsz256_2-prompt1024-resp512-0.567.log) | diff --git a/verl/examples/ppo_trainer/run_qwen3_8b_fsdp.sh b/verl/examples/ppo_trainer/run_qwen3_8b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..f94654fb7601c1c633577113eff478b6a54fb935 --- /dev/null +++ b/verl/examples/ppo_trainer/run_qwen3_8b_fsdp.sh @@ -0,0 +1,138 @@ +#!/usr/bin/env bash +# PPO | text | FSDP training | GPU/NPU +# Canonical PPO (actor + critic) baseline on GSM8K + MATH. + +set -xeuo pipefail + +########################### user-adjustable ########################### +# DEVICE is auto-detected by probing torch_npu; override only for special cases. +DEVICE=${DEVICE:-$(python3 -c 'import torch_npu' 2>/dev/null && echo npu || echo gpu)} +INFER_BACKEND=${INFER_BACKEND:-vllm} +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-8B} +CRITIC_MODEL_PATH=${CRITIC_MODEL_PATH:-$MODEL_PATH} +NNODES=${NNODES:-1} +NDEVICES_PER_NODE=${NDEVICES_PER_NODE:-} + +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-1024} +PPO_MINI_BATCH_SIZE=${PPO_MINI_BATCH_SIZE:-256} +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-1024} +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-2048} +PPO_MAX_TOKEN_LEN_PER_GPU=${PPO_MAX_TOKEN_LEN_PER_GPU:-24576} + +ACTOR_LR=${ACTOR_LR:-1e-6} +CRITIC_LR=${CRITIC_LR:-1e-5} +ENTROPY_COEFF=${ENTROPY_COEFF:-0} + +ROLLOUT_TP=${ROLLOUT_TP:-2} +ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.6} +ROLLOUT_N=${ROLLOUT_N:-1} + +TOTAL_EPOCHS=${TOTAL_EPOCHS:-15} +SAVE_FREQ=${SAVE_FREQ:-20} +TEST_FREQ=${TEST_FREQ:-5} + +PROJECT_NAME=${PROJECT_NAME:-verl_ppo_gsm8k_math} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-qwen3_8b_ppo_${INFER_BACKEND}_fsdp_$(date +%Y%m%d_%H%M)} + +GSM8K_TRAIN_FILE=${GSM8K_TRAIN_FILE:-$HOME/data/gsm8k/train.parquet} +GSM8K_TEST_FILE=${GSM8K_TEST_FILE:-$HOME/data/gsm8k/test.parquet} +MATH_TRAIN_FILE=${MATH_TRAIN_FILE:-$HOME/data/math/train.parquet} +MATH_TEST_FILE=${MATH_TEST_FILE:-$HOME/data/math/test.parquet} +########################### end user-adjustable ########################### + +########################### derived defaults ########################### +n_devices_per_node=${NDEVICES_PER_NODE:-8} + +case "${DEVICE}" in + gpu) + ;; + npu) + export HCCL_CONNECT_TIMEOUT=2400 + export HCCL_EXEC_TIMEOUT=2400 + export HCCL_OP_EXPANSION_MODE=AIV + export CLOSE_MATMUL_K_SHIFT=1 + ;; + *) + echo "Unsupported DEVICE=${DEVICE}. Expected 'gpu' or 'npu'." >&2 + exit 1 + ;; +esac + +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=gae + data.train_files="['$GSM8K_TRAIN_FILE', '$MATH_TRAIN_FILE']" + data.val_files="['$GSM8K_TEST_FILE', '$MATH_TEST_FILE']" + data.train_batch_size=${TRAIN_BATCH_SIZE} + data.max_prompt_length=${MAX_PROMPT_LENGTH} + data.max_response_length=${MAX_RESPONSE_LENGTH} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${ACTOR_LR} + actor_rollout_ref.actor.ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} + actor_rollout_ref.actor.entropy_coeff=${ENTROPY_COEFF} + actor_rollout_ref.actor.fsdp_config.param_offload=False + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=${INFER_BACKEND} + actor_rollout_ref.rollout.tensor_model_parallel_size=${ROLLOUT_TP} + actor_rollout_ref.rollout.gpu_memory_utilization=${ROLLOUT_GPU_MEM_UTIL} + actor_rollout_ref.rollout.n=${ROLLOUT_N} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} + actor_rollout_ref.ref.fsdp_config.param_offload=True +) + +CRITIC=( + critic.model.path="$CRITIC_MODEL_PATH" + critic.model.use_remove_padding=True + critic.model.enable_gradient_checkpointing=True + critic.optim.lr=${CRITIC_LR} + critic.use_dynamic_bsz=True + critic.ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} + critic.fsdp.param_offload=False + critic.fsdp.optimizer_offload=False +) + +TRAINER=( + trainer.balance_batch=True + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${PROJECT_NAME} + trainer.experiment_name=${EXPERIMENT_NAME} + trainer.n_gpus_per_node=${n_devices_per_node} + trainer.nnodes=${NNODES} + trainer.save_freq=${SAVE_FREQ} + trainer.test_freq=${TEST_FREQ} + trainer.total_epochs=${TOTAL_EPOCHS} +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${CRITIC[@]}" \ + "${TRAINER[@]}" \ + "$@" diff --git a/verl/examples/ppo_trainer/run_qwen3_8b_megatron.sh b/verl/examples/ppo_trainer/run_qwen3_8b_megatron.sh new file mode 100644 index 0000000000000000000000000000000000000000..02423f4c29420fb7401933ea8069df01f73a194d --- /dev/null +++ b/verl/examples/ppo_trainer/run_qwen3_8b_megatron.sh @@ -0,0 +1,130 @@ +#!/usr/bin/env bash +# PPO | text | vLLM rollout | Megatron training | NVIDIA GPUs + +set -xeuo pipefail + +export CUDA_DEVICE_MAX_CONNECTIONS=1 + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-8B} +CRITIC_MODEL_PATH=${CRITIC_MODEL_PATH:-$MODEL_PATH} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +train_batch_size=${TRAIN_BATCH_SIZE:-1024} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-256} +max_prompt_length=${MAX_PROMPT_LENGTH:-1024} +max_response_length=${MAX_RESPONSE_LENGTH:-2048} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-24576} + +actor_lr=${ACTOR_LR:-1e-6} +critic_lr=${CRITIC_LR:-1e-5} +entropy_coeff=${ENTROPY_COEFF:-0} + +actor_tp=${ACTOR_TP:-2} +actor_pp=${ACTOR_PP:-2} +critic_tp=${CRITIC_TP:-2} +critic_pp=${CRITIC_PP:-2} + +rollout_tp=${ROLLOUT_TP:-2} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.6} +rollout_n=${ROLLOUT_N:-1} + +total_epochs=${TOTAL_EPOCHS:-15} +save_freq=${SAVE_FREQ:-20} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_ppo_gsm8k_math} +experiment_name=${EXPERIMENT_NAME:-qwen3_8b_vllm_megatron} + +gsm8k_train=${GSM8K_TRAIN_FILE:-$HOME/data/gsm8k/train.parquet} +gsm8k_test=${GSM8K_TEST_FILE:-$HOME/data/gsm8k/test.parquet} +math_train=${MATH_TRAIN_FILE:-$HOME/data/math/train.parquet} +math_test=${MATH_TEST_FILE:-$HOME/data/math/test.parquet} +# ---- end user-adjustable ---- + +# ---- no user adjustment needed below ---- +train_files="['$gsm8k_train', '$math_train']" +val_files="['$gsm8k_test', '$math_test']" +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=gae + data.train_files="$train_files" + data.val_files="$val_files" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${actor_pp} +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${actor_pp} +) + +CRITIC=( + critic.model.path="$CRITIC_MODEL_PATH" + critic.model.use_remove_padding=True + critic.optim.lr=${critic_lr} + critic.use_dynamic_bsz=True + critic.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + critic.megatron.tensor_model_parallel_size=${critic_tp} + critic.megatron.pipeline_model_parallel_size=${critic_pp} +) + +TRAINER=( + trainer.balance_batch=True + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( + model_engine=megatron +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${CRITIC[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/prefix_grouper/README.md b/verl/examples/prefix_grouper/README.md new file mode 100644 index 0000000000000000000000000000000000000000..cd3f602c7ef83f820a7097b3e63674ac411246e7 --- /dev/null +++ b/verl/examples/prefix_grouper/README.md @@ -0,0 +1,85 @@ +# PrefixGrouper Examples + +This directory contains examples for using **PrefixGrouper**, an optimization technique that groups samples by shared prompts to reduce redundant computations in GRPO. + +## Introduction + +> Official Repository: [https://github.com/johncaged/PrefixGrouper](https://github.com/johncaged/PrefixGrouper) + +``PrefixGrouper`` is a plug-and-play efficient GRPO training tool that requires minimal modifications to existing codebases to achieve reduced computation, lower device memory consumption, and accelerated training. + +In current mainstream GRPO training pipelines, policy model training primarily involves copying prefixes (typically questions, multimodal inputs, etc.) `G` times. Consequently, when training data prefixes are sufficiently long (e.g., long-context reasoning, image/long-video inference), redundant computation during training becomes non-negligible. + +**PrefixGrouper** decomposes the original redundant self-attention operation into prefix self-attention + suffix concat-attention. + +

+ +

+ +## Installation + +```bash +pip install prefix_grouper +``` + +## Limitations + +- Currently only supports FSDP worker (Megatron worker is not supported yet). +- Incompatible with `use_dynamic_bsz=True`. +- Incompatible with `use_remove_padding=True` (Flash Attention V2 variable length). +- Incompatible with `use_fused_kernels=True`. +- Incompatible with Ulysses sequence parallelism (`use_ulysses_sp=True`) and ring-attention. + +Note: `balance_batch=True` is now supported with group-level balancing, which keeps samples with the same uid together on the same rank. However, this requires `batch_size % (world_size * rollout.n) == 0`. For example, with `world_size=8` and `rollout.n=4`, you need `batch_size` to be a multiple of 32. + +## How to Use + +### 1. Enable PrefixGrouper in Config + +Simply set `use_prefix_grouper=True` in your training config: + +```yaml +actor_rollout_ref: + actor: + use_prefix_grouper: True + model: + use_remove_padding: False +``` + +Optionally enable balance_batch for better load distribution: +```yaml +trainer: + balance_batch: True # Now supported with group-level balancing +``` + +### 2. Run Training + +Use the provided script `run_qwen3_8b_fsdp.sh` as an example: + +```bash +bash examples/prefix_grouper/run_qwen3_8b_fsdp.sh +``` + +## How It Works + +When `use_prefix_grouper=True`, verl automatically patches the attention functions in `transformers.modeling_utils.ALL_ATTENTION_FUNCTIONS` to support the `prefix_grouper` parameter. No model code modifications are needed. + +The patch wraps each attention function to: +1. Extract `prefix_grouper` from kwargs +2. If `prefix_grouper` is None, call original attention +3. If `prefix_grouper` is provided, use PrefixGrouper's optimized attention computation + +## Performance + +**Benchmark Results** (Qwen3-4B, 4×H800, `rollout.n=4`): + +| Context Length | Metric | PG | No PG | Speedup | +|----------------|--------|-----|-------|---------| +| **4K** | `old_log_prob` | 1.31s | 1.70s | **1.30x** | +| | `update_actor` | 4.80s | 6.07s | **1.26x** | +| | `step` | 17.08s | 19.40s | **1.14x** | +| **8K** | `old_log_prob` | 1.69s | 2.63s | **1.56x** | +| | `update_actor` | 5.98s | 10.18s | **1.70x** | +| | `step` | 19.48s | 24.71s | **1.27x** | + +As context length increases, the speedup becomes more pronounced. diff --git a/verl/examples/prefix_grouper/run_qwen3_8b_fsdp.sh b/verl/examples/prefix_grouper/run_qwen3_8b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..7651e4d8864817213235f759ae477c425f25bf21 --- /dev/null +++ b/verl/examples/prefix_grouper/run_qwen3_8b_fsdp.sh @@ -0,0 +1,107 @@ +#!/usr/bin/env bash +# GRPO + prefix-grouper | text | vLLM rollout | FSDP training | NVIDIA GPUs +# Demonstrates `actor.use_prefix_grouper=True` on Qwen3-8B / GSM8K. + +set -xeuo pipefail + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-8B} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +train_batch_size=${TRAIN_BATCH_SIZE:-1024} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-256} +max_prompt_length=${MAX_PROMPT_LENGTH:-512} +max_response_length=${MAX_RESPONSE_LENGTH:-1024} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-12288} + +actor_lr=${ACTOR_LR:-1e-6} +kl_loss_coef=${KL_LOSS_COEF:-0.001} +entropy_coeff=${ENTROPY_COEFF:-0} + +rollout_tp=${ROLLOUT_TP:-2} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.6} +rollout_n=${ROLLOUT_N:-5} + +total_epochs=${TOTAL_EPOCHS:-15} +save_freq=${SAVE_FREQ:-20} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_grpo_prefix_grouper_gsm8k} +experiment_name=${EXPERIMENT_NAME:-qwen3_8b_vllm_fsdp_prefix_grouper} +# ---- end user-adjustable ---- +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files=$HOME/data/gsm8k/train.parquet + data.val_files=$HOME/data/gsm8k/test.parquet + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=False + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_prefix_grouper=True + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.fsdp_config.param_offload=False + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.fsdp_config.param_offload=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/profile/README.md b/verl/examples/profile/README.md new file mode 100644 index 0000000000000000000000000000000000000000..6e45e8b2bacd9cd400c72a81971c0d0514fff772 --- /dev/null +++ b/verl/examples/profile/README.md @@ -0,0 +1,28 @@ +# Profiling Examples + +End-to-end GRPO runs that enable one of verl's profilers so you can capture a performance/memory trace without authoring a bespoke launcher. All scripts use the current `verl.trainer.main_ppo` entry point and the current Hydra API. + +## Canonical Scripts + +| Script | Profiler | Model | Infer | Train | Platform | +|-----------------------------------------|---------------|--------------------|--------|-------|----------| +| `run_qwen3_8b_npu_profile_e2e.sh` | NPU (E2E) | Qwen3-8B | vLLM | FSDP | NPU | +| `run_qwen3_8b_npu_profile_discrete.sh` | NPU (discrete)| Qwen3-8B | vLLM | FSDP | NPU | +| `run_qwen2_5_vl_7b_torch_memory.sh` | torch_memory | Qwen2.5-VL-7B | SGLang | FSDP | NVIDIA | + +### NPU profiling + +- `*_profile_e2e.sh` — one end-to-end timeline for all ranks. +- `*_profile_discrete.sh` — per-stage (rollout/ref/actor) discrete traces. + +Controlled via `global_profiler.tool=npu`, `global_profiler.steps=[...]`, `global_profiler.save_path=...`, plus per-role `actor_rollout_ref.*.profiler.*` overrides. Override any of `PROFILE_STEPS`, `PROFILE_SAVE_PATH`, `PROFILE_LEVEL`, `PROFILE_CONTENTS`, `PROFILE_DISCRETE`, `PROFILE_RANKS_ALL` to adjust behavior. + +### Torch memory profiling + +- `run_qwen2_5_vl_7b_torch_memory.sh` dumps `torch.cuda._record_memory_history` snapshots to `global_profiler.save_path` (default `./mem_snapshots`). Load the `.pickle` in PyTorch's memory viz UI. Override `TRACE_ALLOC_MAX_ENTRIES`, `STACK_DEPTH`, `PROFILE_SAVE_PATH` as needed. + +## Conventions + +- `VAR=${VAR:-default}` for `MODEL_PATH`, batch sizes, learning rate, rollout TP, profile options, etc. +- Dynamic batch size and `trainer.balance_batch=True` are enabled by default. +- No deprecated config (`ppo_megatron_trainer.yaml`, `ppo_micro_batch_size`, `data.val_batch_size`, top-level `reward_model.*`, `actor.ulysses_sequence_parallel_size`). diff --git a/verl/examples/profile/run_qwen2_5_vl_7b_torch_memory.sh b/verl/examples/profile/run_qwen2_5_vl_7b_torch_memory.sh new file mode 100644 index 0000000000000000000000000000000000000000..c2eca222c9ba1f0117e9a513196e79c0165dfea1 --- /dev/null +++ b/verl/examples/profile/run_qwen2_5_vl_7b_torch_memory.sh @@ -0,0 +1,120 @@ +#!/usr/bin/env bash +# GRPO profiling (torch memory) | vision | SGLang rollout | FSDP training | NVIDIA GPUs +# +# Captures Torch CUDA memory snapshots to disk for leak / fragmentation analysis. +# Snapshots are written under ./mem_snapshots; use torch.cuda._dump_snapshot to inspect. + +set -xeuo pipefail + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen2.5-VL-7B-Instruct} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +profile_save_path=${PROFILE_SAVE_PATH:-./mem_snapshots} +trace_alloc_max_entries=${TRACE_ALLOC_MAX_ENTRIES:-100000} +stack_depth=${STACK_DEPTH:-32} + +train_batch_size=${TRAIN_BATCH_SIZE:-512} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-128} +max_prompt_length=${MAX_PROMPT_LENGTH:-1024} +max_response_length=${MAX_RESPONSE_LENGTH:-2048} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-24576} + +actor_lr=${ACTOR_LR:-1e-6} +kl_loss_coef=${KL_LOSS_COEF:-0.01} +entropy_coeff=${ENTROPY_COEFF:-0} + +rollout_tp=${ROLLOUT_TP:-1} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.85} +rollout_n=${ROLLOUT_N:-5} + +total_epochs=${TOTAL_EPOCHS:-15} +save_freq=${SAVE_FREQ:-20} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_grpo_profile} +experiment_name=${EXPERIMENT_NAME:-qwen2_5_vl_7b_torch_memory} +# ---- end user-adjustable ---- +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files=$HOME/data/geo3k/train.parquet + data.val_files=$HOME/data/geo3k/test.parquet + data.image_key=images + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.fsdp_config.param_offload=False + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=sglang + actor_rollout_ref.rollout.multi_stage_wake_up=True + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.enable_chunked_prefill=False + actor_rollout_ref.rollout.enforce_eager=False + actor_rollout_ref.rollout.free_cache_engine=True + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.fsdp_config.param_offload=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( + global_profiler.tool=torch_memory + global_profiler.save_path=${profile_save_path} + global_profiler.global_tool_config.torch_memory.trace_alloc_max_entries=${trace_alloc_max_entries} + global_profiler.global_tool_config.torch_memory.stack_depth=${stack_depth} +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/profile/run_qwen3_8b_npu_profile_discrete.sh b/verl/examples/profile/run_qwen3_8b_npu_profile_discrete.sh new file mode 100644 index 0000000000000000000000000000000000000000..9cd2aaa961f482cbbd3096545f31ee5c859930a8 --- /dev/null +++ b/verl/examples/profile/run_qwen3_8b_npu_profile_discrete.sh @@ -0,0 +1,130 @@ +#!/usr/bin/env bash +# GRPO profiling (discrete per-rank) | text | vLLM rollout | FSDP training | Ascend NPU +# +# Captures NPU traces on a subset of ranks in "discrete" mode +# (one trace per worker module). Useful for targeted hot-path analysis. + +set -xeuo pipefail + +export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-8B} +NNODES=${NNODES:-1} +NPUS_PER_NODE=${NPUS_PER_NODE:-8} + +profile_steps=${PROFILE_STEPS:-"[2,4]"} +profile_ranks=${PROFILE_RANKS:-"[1,2]"} +profile_ranks_all=${PROFILE_RANKS_ALL:-False} +profile_discrete=${PROFILE_DISCRETE:-True} +profile_save_path=${PROFILE_SAVE_PATH:-$HOME/profile_data} +profile_level=${PROFILE_LEVEL:-level0} +profile_contents=${PROFILE_CONTENTS:-"['npu','cpu']"} +profile_analysis=${PROFILE_ANALYSIS:-True} + +train_batch_size=${TRAIN_BATCH_SIZE:-32} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-16} +max_prompt_length=${MAX_PROMPT_LENGTH:-1024} +max_response_length=${MAX_RESPONSE_LENGTH:-1024} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-8192} + +actor_lr=${ACTOR_LR:-5e-8} + +rollout_tp=${ROLLOUT_TP:-4} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.6} +rollout_n=${ROLLOUT_N:-4} + +project_name=${PROJECT_NAME:-verl_grpo_profile} +experiment_name=${EXPERIMENT_NAME:-qwen3_8b_npu_profile_discrete} +# ---- end user-adjustable ---- +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files=$HOME/data/gsm8k/train.parquet + data.val_files=$HOME/data/gsm8k/test.parquet + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=0.001 + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=0 + actor_rollout_ref.actor.fsdp_config.param_offload=False + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False + actor_rollout_ref.actor.profiler.enable=True + actor_rollout_ref.actor.profiler.ranks=${profile_ranks} + actor_rollout_ref.actor.profiler.all_ranks=${profile_ranks_all} + actor_rollout_ref.actor.profiler.tool_config.npu.discrete=${profile_discrete} + actor_rollout_ref.actor.profiler.tool_config.npu.contents=${profile_contents} + actor_rollout_ref.actor.profiler.tool_config.npu.level=${profile_level} + actor_rollout_ref.actor.profiler.tool_config.npu.analysis=${profile_analysis} +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.rollout.enable_chunked_prefill=False + actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=4096 +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.fsdp_config.param_offload=True + actor_rollout_ref.ref.profiler.enable=True + actor_rollout_ref.ref.profiler.ranks=${profile_ranks} + actor_rollout_ref.ref.profiler.all_ranks=${profile_ranks_all} + actor_rollout_ref.ref.profiler.tool_config.npu.discrete=${profile_discrete} + actor_rollout_ref.ref.profiler.tool_config.npu.contents=${profile_contents} + actor_rollout_ref.ref.profiler.tool_config.npu.level=${profile_level} + actor_rollout_ref.ref.profiler.tool_config.npu.analysis=${profile_analysis} +) + +TRAINER=( + trainer.balance_batch=True + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=-1 + trainer.test_freq=5 + trainer.total_epochs=5 +) + +EXTRA=( + global_profiler.tool=npu + global_profiler.steps=${profile_steps} + global_profiler.save_path=${profile_save_path} +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/profile/run_qwen3_8b_npu_profile_e2e.sh b/verl/examples/profile/run_qwen3_8b_npu_profile_e2e.sh new file mode 100644 index 0000000000000000000000000000000000000000..17ef9ff61a1bc1cf1a919a42882723e92aca0550 --- /dev/null +++ b/verl/examples/profile/run_qwen3_8b_npu_profile_e2e.sh @@ -0,0 +1,127 @@ +#!/usr/bin/env bash +# GRPO profiling (end-to-end, all ranks) | text | vLLM rollout | FSDP training | Ascend NPU +# +# Captures NPU traces on all ranks in non-discrete mode (one combined trace). +# Useful for full step-level timeline analysis. + +set -xeuo pipefail + +export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-8B} +NNODES=${NNODES:-1} +NPUS_PER_NODE=${NPUS_PER_NODE:-8} + +profile_steps=${PROFILE_STEPS:-"[2,4]"} +profile_ranks_all=${PROFILE_RANKS_ALL:-True} +profile_discrete=${PROFILE_DISCRETE:-False} +profile_save_path=${PROFILE_SAVE_PATH:-$HOME/profile_data} +profile_level=${PROFILE_LEVEL:-level0} +profile_contents=${PROFILE_CONTENTS:-"['npu','cpu']"} +profile_analysis=${PROFILE_ANALYSIS:-True} + +train_batch_size=${TRAIN_BATCH_SIZE:-32} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-16} +max_prompt_length=${MAX_PROMPT_LENGTH:-1024} +max_response_length=${MAX_RESPONSE_LENGTH:-1024} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-8192} + +actor_lr=${ACTOR_LR:-5e-8} + +rollout_tp=${ROLLOUT_TP:-4} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.6} +rollout_n=${ROLLOUT_N:-4} + +project_name=${PROJECT_NAME:-verl_grpo_profile} +experiment_name=${EXPERIMENT_NAME:-qwen3_8b_npu_profile_e2e} +# ---- end user-adjustable ---- +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files=$HOME/data/gsm8k/train.parquet + data.val_files=$HOME/data/gsm8k/test.parquet + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=0.001 + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=0 + actor_rollout_ref.actor.fsdp_config.param_offload=False + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False + actor_rollout_ref.actor.profiler.enable=True + actor_rollout_ref.actor.profiler.all_ranks=${profile_ranks_all} + actor_rollout_ref.actor.profiler.tool_config.npu.discrete=${profile_discrete} + actor_rollout_ref.actor.profiler.tool_config.npu.contents=${profile_contents} + actor_rollout_ref.actor.profiler.tool_config.npu.level=${profile_level} + actor_rollout_ref.actor.profiler.tool_config.npu.analysis=${profile_analysis} +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.rollout.enable_chunked_prefill=False + actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=4096 +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.fsdp_config.param_offload=True + actor_rollout_ref.ref.profiler.enable=True + actor_rollout_ref.ref.profiler.all_ranks=${profile_ranks_all} + actor_rollout_ref.ref.profiler.tool_config.npu.discrete=${profile_discrete} + actor_rollout_ref.ref.profiler.tool_config.npu.contents=${profile_contents} + actor_rollout_ref.ref.profiler.tool_config.npu.level=${profile_level} + actor_rollout_ref.ref.profiler.tool_config.npu.analysis=${profile_analysis} +) + +TRAINER=( + trainer.balance_batch=True + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=-1 + trainer.test_freq=5 + trainer.total_epochs=5 +) + +EXTRA=( + global_profiler.tool=npu + global_profiler.steps=${profile_steps} + global_profiler.save_path=${profile_save_path} +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/reinforce_plus_plus_trainer/README.md b/verl/examples/reinforce_plus_plus_trainer/README.md new file mode 100644 index 0000000000000000000000000000000000000000..ac54593fa0cf822bbb316da695a4be3093d4e653 --- /dev/null +++ b/verl/examples/reinforce_plus_plus_trainer/README.md @@ -0,0 +1,13 @@ +# REINFORCE++ + +REINFORCE++ is a simple, critic-free PG variant that extends REINFORCE with token-level KL penalties and advantage whitening. + +Reference: [REINFORCE++: A Simple and Efficient Approach for Aligning Large Language Models](https://arxiv.org/abs/2501.03262). + +## Canonical Scripts + +| Script | Infer | Train | Platform | +|--------------------------------------|-------|-------|----------| +| `run_qwen3_8b_fsdp.sh` | vLLM | FSDP | NVIDIA | + +Switch to the baseline variant by setting `ADV_ESTIMATOR=reinforce_plus_plus_baseline` when running the script. diff --git a/verl/examples/reinforce_plus_plus_trainer/run_qwen3_8b_fsdp.sh b/verl/examples/reinforce_plus_plus_trainer/run_qwen3_8b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..ddbeca8f5256e8ef13460bdefc8526f6c8ab76b5 --- /dev/null +++ b/verl/examples/reinforce_plus_plus_trainer/run_qwen3_8b_fsdp.sh @@ -0,0 +1,113 @@ +#!/usr/bin/env bash +# REINFORCE++ | text | vLLM rollout | FSDP training | NVIDIA GPUs +# Set ADV_ESTIMATOR=reinforce_plus_plus_baseline to use the baseline variant. + +set -xeuo pipefail + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-8B} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +adv_estimator=${ADV_ESTIMATOR:-reinforce_plus_plus} + +train_batch_size=${TRAIN_BATCH_SIZE:-1024} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-1024} +max_prompt_length=${MAX_PROMPT_LENGTH:-1024} +max_response_length=${MAX_RESPONSE_LENGTH:-2048} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-24576} + +actor_lr=${ACTOR_LR:-3e-6} +entropy_coeff=${ENTROPY_COEFF:-0} + +rollout_tp=${ROLLOUT_TP:-2} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.6} +rollout_n=${ROLLOUT_N:-8} + +total_epochs=${TOTAL_EPOCHS:-15} +save_freq=${SAVE_FREQ:-20} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_reinforce_plus_plus_gsm8k_math} +experiment_name=${EXPERIMENT_NAME:-qwen3_8b_vllm_fsdp} +# ---- end user-adjustable ---- + +gsm8k_train=$HOME/data/gsm8k/train.parquet +gsm8k_test=$HOME/data/gsm8k/test.parquet +math_train=$HOME/data/math/train.parquet +math_test=$HOME/data/math/test.parquet + +train_files="['$gsm8k_train', '$math_train']" +val_files="['$gsm8k_test', '$math_test']" +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=${adv_estimator} + algorithm.use_kl_in_reward=True + data.train_files="$train_files" + data.val_files="$val_files" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=False + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.fsdp_config.param_offload=False + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.fsdp_config.param_offload=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/remax_trainer/README.md b/verl/examples/remax_trainer/README.md new file mode 100644 index 0000000000000000000000000000000000000000..5319ff178430575e1ceb8e4ad02c6e626f07595c --- /dev/null +++ b/verl/examples/remax_trainer/README.md @@ -0,0 +1,19 @@ +# ReMax + +ReMax is a lightweight policy-gradient method that uses a single greedy-decoded baseline response per prompt to reduce variance without a critic. + +Reference: [ReMax: A Simple, Effective, and Efficient Reinforcement Learning Method for Aligning Large Language Models](https://arxiv.org/abs/2310.10505). + +## Canonical Scripts + +| Script | Infer | Train | Platform | +|-------------------------------------------|-------|-----------|----------| +| `run_qwen3_8b_fsdp.sh` | vLLM | FSDP | NVIDIA | +| `run_qwen2.5_math_7b_fsdp_sync.sh` | vLLM | FSDP+Sync | NVIDIA | + +Override any argument via env vars at the top of the script. + +## Key Flags + +- `algorithm.adv_estimator=remax` +- `actor_rollout_ref.actor.use_kl_loss=False` and `algorithm.use_kl_in_reward=True` diff --git a/verl/examples/remax_trainer/run_qwen2.5_math_7b_sync_fsdp.sh b/verl/examples/remax_trainer/run_qwen2.5_math_7b_sync_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..bce3596f9b0b28a9c2516f50f39f0af160147d82 --- /dev/null +++ b/verl/examples/remax_trainer/run_qwen2.5_math_7b_sync_fsdp.sh @@ -0,0 +1,114 @@ +#!/usr/bin/env bash +# ReMax | text | vLLM rollout | FSDP training | synchronous TransferQueue trainer | NVIDIA GPUs + +set -xeuo pipefail + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen2.5-Math-7B} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +train_batch_size=${TRAIN_BATCH_SIZE:-128} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-64} +max_prompt_length=${MAX_PROMPT_LENGTH:-1024} +max_response_length=${MAX_RESPONSE_LENGTH:-2048} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-24576} + +actor_lr=${ACTOR_LR:-1e-6} +entropy_coeff=${ENTROPY_COEFF:-0} + +rollout_tp=${ROLLOUT_TP:-2} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.7} +rollout_n=${ROLLOUT_N:-4} + +total_epochs=${TOTAL_EPOCHS:-10} +save_freq=${SAVE_FREQ:-20} +test_freq=${TEST_FREQ:-10} + +project_name=${PROJECT_NAME:-verl_remax_gsm8k_math} +experiment_name=${EXPERIMENT_NAME:-qwen2.5_math_7b_vllm_fsdp_sync} +# ---- end user-adjustable ---- + +gsm8k_train=$HOME/data/gsm8k/train.parquet +gsm8k_test=$HOME/data/gsm8k/test.parquet +math_train=$HOME/data/math/train.parquet +math_test=$HOME/data/math/test.parquet + +train_files="['$gsm8k_train', '$math_train']" +val_files="['$gsm8k_test', '$math_test']" +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=remax + algorithm.use_kl_in_reward=True + algorithm.kl_penalty=kl + algorithm.kl_ctrl.kl_coef=0.0 + data.train_files="$train_files" + data.val_files="$val_files" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=False + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.fsdp_config.param_offload=False + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False + actor_rollout_ref.actor.ulysses_sequence_parallel_size=2 +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.fsdp_config.param_offload=False +) + +TRAINER=( + trainer.balance_batch=True + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( + +ray_kwargs.ray_init.runtime_env.env_vars.TRANSFER_QUEUE_ENABLE=1 +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo_sync \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/remax_trainer/run_qwen3_8b_fsdp.sh b/verl/examples/remax_trainer/run_qwen3_8b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..b64ae04572c8f00445f63b4b90b909230781ff21 --- /dev/null +++ b/verl/examples/remax_trainer/run_qwen3_8b_fsdp.sh @@ -0,0 +1,112 @@ +#!/usr/bin/env bash +# ReMax | text | vLLM rollout | FSDP training | NVIDIA GPUs + +set -xeuo pipefail + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-8B} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +train_batch_size=${TRAIN_BATCH_SIZE:-1024} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-256} +max_prompt_length=${MAX_PROMPT_LENGTH:-1024} +max_response_length=${MAX_RESPONSE_LENGTH:-2048} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-24576} + +actor_lr=${ACTOR_LR:-1e-6} +entropy_coeff=${ENTROPY_COEFF:-0} + +rollout_tp=${ROLLOUT_TP:-2} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.6} +rollout_n=${ROLLOUT_N:-4} + +total_epochs=${TOTAL_EPOCHS:-10} +save_freq=${SAVE_FREQ:-20} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_remax_gsm8k_math} +experiment_name=${EXPERIMENT_NAME:-qwen3_8b_vllm_fsdp} +# ---- end user-adjustable ---- + +gsm8k_train=$HOME/data/gsm8k/train.parquet +gsm8k_test=$HOME/data/gsm8k/test.parquet +math_train=$HOME/data/math/train.parquet +math_test=$HOME/data/math/test.parquet + +train_files="['$gsm8k_train', '$math_train']" +val_files="['$gsm8k_test', '$math_test']" +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=remax + algorithm.use_kl_in_reward=True + algorithm.kl_penalty=kl + algorithm.kl_ctrl.kl_coef=0.001 + data.train_files="$train_files" + data.val_files="$val_files" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=False + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.fsdp_config.param_offload=False + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.fsdp_config.param_offload=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/rloo_trainer/README.md b/verl/examples/rloo_trainer/README.md new file mode 100644 index 0000000000000000000000000000000000000000..a855f556616652208a0dfdace654e604c608dfde --- /dev/null +++ b/verl/examples/rloo_trainer/README.md @@ -0,0 +1,19 @@ +# RLOO (REINFORCE Leave-One-Out) + +RLOO is a simple policy gradient baseline that avoids a critic. Advantage for each sample is computed against the average of its siblings (leave-one-out), which acts as a per-prompt variance-reduction baseline. + +Reference: [Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in LLMs](https://arxiv.org/abs/2402.14740). + +## Canonical Scripts + +| Script | Infer | Train | Platform | +|--------------------------------------|-------|-------|----------| +| `run_qwen3_8b_fsdp.sh` | vLLM | FSDP | NVIDIA | + +Override any argument via env vars at the top of the script (e.g. `MODEL_PATH=Qwen/Qwen3-14B bash run_qwen3_8b_fsdp.sh`). + +## Key Flags + +- `algorithm.adv_estimator=rloo` +- `actor_rollout_ref.rollout.n=5` — RLOO needs ≥2 samples per prompt; 5 is a common default. +- `actor_rollout_ref.actor.use_kl_loss=False` and `algorithm.use_kl_in_reward=True` — RLOO typically uses reward-side KL, not loss-side KL. diff --git a/verl/examples/rloo_trainer/run_qwen3_8b_fsdp.sh b/verl/examples/rloo_trainer/run_qwen3_8b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..db41415365a2b708504c18a4ba34918525541b48 --- /dev/null +++ b/verl/examples/rloo_trainer/run_qwen3_8b_fsdp.sh @@ -0,0 +1,112 @@ +#!/usr/bin/env bash +# RLOO | text | vLLM rollout | FSDP training | NVIDIA GPUs + +set -xeuo pipefail + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-8B} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +train_batch_size=${TRAIN_BATCH_SIZE:-1024} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-256} +max_prompt_length=${MAX_PROMPT_LENGTH:-1024} +max_response_length=${MAX_RESPONSE_LENGTH:-2048} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-24576} + +actor_lr=${ACTOR_LR:-1e-6} +entropy_coeff=${ENTROPY_COEFF:-0} + +rollout_tp=${ROLLOUT_TP:-2} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.6} +rollout_n=${ROLLOUT_N:-5} + +total_epochs=${TOTAL_EPOCHS:-15} +save_freq=${SAVE_FREQ:-20} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_rloo_gsm8k_math} +experiment_name=${EXPERIMENT_NAME:-qwen3_8b_vllm_fsdp} +# ---- end user-adjustable ---- + +gsm8k_train=$HOME/data/gsm8k/train.parquet +gsm8k_test=$HOME/data/gsm8k/test.parquet +math_train=$HOME/data/math/train.parquet +math_test=$HOME/data/math/test.parquet + +train_files="['$gsm8k_train', '$math_train']" +val_files="['$gsm8k_test', '$math_test']" +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=rloo + algorithm.use_kl_in_reward=True + algorithm.kl_penalty=kl + algorithm.kl_ctrl.kl_coef=0.001 + data.train_files="$train_files" + data.val_files="$val_files" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=False + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.fsdp_config.param_offload=False + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.fsdp_config.param_offload=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/rollout_correction/run_qwen2_5_7b_fsdp.sh b/verl/examples/rollout_correction/run_qwen2_5_7b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..69cedd59569c4f9fb7fa54b7c7f3ac43fdd8cafb --- /dev/null +++ b/verl/examples/rollout_correction/run_qwen2_5_7b_fsdp.sh @@ -0,0 +1,127 @@ +#!/usr/bin/env bash +# Example: RLOO (REINFORCE Leave-One-Out) with Rollout Correction +# This demonstrates self-normalized sequence-level IS with pure policy gradient +# +# References: +# - Rollout Correction Docs: https://github.com/verl-project/verl/blob/main/docs/algo/rollout_corr.md +# - Rollout Correction Math: https://github.com/verl-project/verl/blob/main/docs/algo/rollout_corr_math.md + +set -xeuo pipefail + +# ============================================================================== +# Rollout Correction Configuration (RLOO) +# ============================================================================== + +# Importance Sampling (IS) weights configuration +rollout_is="sequence" # Self-normalized sequence-level IS +rollout_is_threshold=2.0 # Upper threshold for IS weights +rollout_is_batch_normalize="true" # Self-normalization (mean=1.0) + +# Rejection Sampling (RS) configuration +rollout_rs="null" # No rejection sampling for basic RLOO +rollout_rs_threshold="null" # RS threshold spec (string or float) + +# Bypass mode with REINFORCE loss (no PPO clipping) +bypass_mode="true" # Skip old_log_prob computation +loss_type="reinforce" # REINFORCE with explicit IS weights (alternative: "ppo_clip") + +# ============================================================================== +# Model and Data Configuration +# ============================================================================== + +MODEL_PATH=${MODEL_PATH:-"Qwen/Qwen2.5-7B"} +TRAIN_FILE=${TRAIN_FILE:-"data/train.parquet"} +TEST_FILE=${TEST_FILE:-"data/test.parquet"} + +max_prompt_length=2048 +max_response_length=4096 + +# ============================================================================== +# Training Configuration +# ============================================================================== + +train_batch_size=128 +ppo_mini_batch_size=32 +ppo_epochs=1 +learning_rate=5e-7 + +# ============================================================================== +# Algorithm Configuration (RLOO) +# ============================================================================== + +adv_estimator=rloo # RLOO advantage estimator +gamma=1.0 + +# ============================================================================== +# Launch Training +# ============================================================================== +########################### parameter arrays ########################### + +DATA=( + data.train_files="${TRAIN_FILE}" + data.val_files="${TEST_FILE}" + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.train_batch_size=${train_batch_size} + data.truncation='left' + algorithm.adv_estimator=${adv_estimator} + algorithm.gamma=${gamma} + algorithm.rollout_correction.rollout_is=${rollout_is} + algorithm.rollout_correction.rollout_is_threshold=${rollout_is_threshold} + algorithm.rollout_correction.rollout_is_batch_normalize=${rollout_is_batch_normalize} + algorithm.rollout_correction.rollout_rs=${rollout_rs} + algorithm.rollout_correction.rollout_rs_threshold=${rollout_rs_threshold} + algorithm.rollout_correction.bypass_mode=${bypass_mode} + algorithm.rollout_correction.loss_type=${loss_type} +) + +MODEL=( + actor_rollout_ref.model.path="${MODEL_PATH}" +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${learning_rate} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=8 + actor_rollout_ref.actor.ppo_epochs=${ppo_epochs} +) + +ROLLOUT=( + actor_rollout_ref.rollout.calculate_log_probs=True + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 + actor_rollout_ref.rollout.name=vllm +) + +TRAINER=( + trainer.logger='["console","wandb"]' + trainer.project_name="rollout_corr_rloo_example" + trainer.experiment_name="rloo_seq_is_pure" + trainer.total_epochs=10 +) + +EXTRA=( +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" + +echo "Training completed!" +echo "" +echo "RLOO Configuration:" +echo " - Algorithm: RLOO (REINFORCE Leave-One-Out)" +echo " - Advantage estimator: ${adv_estimator}" +echo " - IS mode: ${rollout_is} (self-normalized: ${rollout_is_batch_normalize})" +echo " - IS threshold: ${rollout_is_threshold}" +echo " - Bypass mode: ${bypass_mode}, loss_type: ${loss_type}" +echo "" +echo "Monitor these key metrics in wandb:" +echo " - rollout_corr/rollout_is_mean (should be ~1.0 before batch norm)" +echo " - rollout_corr/rollout_is_batch_norm_factor (normalization factor applied)" +echo " - rollout_corr/rollout_is_eff_sample_size (should be >0.5)" diff --git a/verl/examples/rollout_correction/run_qwen2_5_7b_fsdp_multi_rs.sh b/verl/examples/rollout_correction/run_qwen2_5_7b_fsdp_multi_rs.sh new file mode 100644 index 0000000000000000000000000000000000000000..e0cead769e143d423c1cf1d0bf7c4cfe5ba94846 --- /dev/null +++ b/verl/examples/rollout_correction/run_qwen2_5_7b_fsdp_multi_rs.sh @@ -0,0 +1,129 @@ +#!/usr/bin/env bash +# Example: PPO-clip with Rollout Correction using multiple RS criteria +# Demonstrates chaining token-level and sequence-level rejection sampling +# (token_k1 + seq_max_k2) alongside optional IS metrics. +# +# References: +# - Rollout Correction Docs: https://github.com/verl-project/verl/blob/main/docs/algo/rollout_corr.md +# - Rollout Correction Math: https://github.com/verl-project/verl/blob/main/docs/algo/rollout_corr_math.md + +set -xeuo pipefail + +# ============================================================================== +# Rollout Correction Configuration (PPO-clip + multi RS) +# ============================================================================== + +# Importance Sampling (IS) weights configuration +rollout_is="token" # Token-level IS for metrics/analysis +rollout_is_threshold=2.0 # Upper threshold for IS weights +rollout_is_batch_normalize="false" # Keep raw truncated weights + +# Rejection Sampling (RS) configuration (multi-criteria) +# - token_k1 keeps per-token ratios inside [lower, upper] +# - seq_max_k2 rejects sequences with extreme chi-square spikes +rollout_rs="token_k1,seq_max_k2" +rollout_rs_threshold="0.6_1.6,2.5" + +# Bypass PPO mode (reuse rollout_log_prob) +bypass_mode="true" +loss_type="ppo_clip" + +# ============================================================================== +# Model and Data Configuration +# ============================================================================== + +MODEL_PATH=${MODEL_PATH:-"Qwen/Qwen2.5-7B"} +TRAIN_FILE=${TRAIN_FILE:-"data/train.parquet"} +TEST_FILE=${TEST_FILE:-"data/test.parquet"} + +max_prompt_length=2048 +max_response_length=4096 + +# ============================================================================== +# Training Configuration +# ============================================================================== + +train_batch_size=128 +ppo_mini_batch_size=32 +ppo_epochs=1 +learning_rate=3e-6 + +# ============================================================================== +# Algorithm Configuration +# ============================================================================== + +adv_estimator=grpo +gamma=1.0 + +# ============================================================================== +# Launch Training +# ============================================================================== +########################### parameter arrays ########################### + +DATA=( + data.train_files="${TRAIN_FILE}" + data.val_files="${TEST_FILE}" + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.train_batch_size=${train_batch_size} + data.truncation='left' + algorithm.adv_estimator=${adv_estimator} + algorithm.gamma=${gamma} + algorithm.rollout_correction.rollout_is=${rollout_is} + algorithm.rollout_correction.rollout_is_threshold=${rollout_is_threshold} + algorithm.rollout_correction.rollout_is_batch_normalize=${rollout_is_batch_normalize} + algorithm.rollout_correction.rollout_rs=\'${rollout_rs}\' + algorithm.rollout_correction.rollout_rs_threshold=\'${rollout_rs_threshold}\' + algorithm.rollout_correction.bypass_mode=${bypass_mode} + algorithm.rollout_correction.loss_type=${loss_type} +) + +MODEL=( + actor_rollout_ref.model.path="${MODEL_PATH}" +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${learning_rate} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=8 + actor_rollout_ref.actor.ppo_epochs=${ppo_epochs} +) + +ROLLOUT=( + actor_rollout_ref.rollout.calculate_log_probs=True + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 + actor_rollout_ref.rollout.name=vllm +) + +TRAINER=( + trainer.logger='["console","wandb"]' + trainer.project_name="rollout_corr_multi_rs_example" + trainer.experiment_name="ppo_clip_multi_rs" + trainer.total_epochs=5 +) + +EXTRA=( +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" + +echo "Training completed!" +echo "" +echo "Multi-RS Configuration:" +echo " - rollout_is: ${rollout_is} (threshold=${rollout_is_threshold}, batch_norm=${rollout_is_batch_normalize})" +echo " - rollout_rs: ${rollout_rs}" +echo " - rollout_rs_threshold: ${rollout_rs_threshold}" +echo " - bypass_mode: ${bypass_mode}, loss_type: ${loss_type}" +echo "" +echo "Track these metrics in wandb:" +echo " - rollout_corr/rollout_rs_token_k1_mean" +echo " - rollout_corr/rollout_rs_seq_max_k2_mean" +echo " - rollout_corr/rollout_rs_masked_fraction" diff --git a/verl/examples/router_replay/README.md b/verl/examples/router_replay/README.md new file mode 100644 index 0000000000000000000000000000000000000000..7ef50ca4e04b5e34b6f8359b9674cca5a9a4e41d --- /dev/null +++ b/verl/examples/router_replay/README.md @@ -0,0 +1,75 @@ +# Router Replay + +Router Replay is an advanced routing replay functionality within the Verl framework designed for Mixture of Experts (MoE) models. It enables deterministic training by recording and replaying routing decisions, ensuring consistent model behavior across training runs. + + +## Key Features + +### Multiple Operating Modes +- **`disabled`**: Router replay functionality is completely disabled +- **`R2`**: Standard router replay mode for recording and replaying routing decisions +- **`R3`**: Rollout-specific router replay mode optimized for reinforcement learning workflows + +### Core Capabilities +- **Seamless Integration**: Works with reinforcement learning pipelines including PPO +- **Distributed Training Support**: Compatible with multi-GPU and multi-node training environments +- **Flexible Configuration**: Easy to configure via YAML files or command-line parameters + +## Configuration + +### RouterReplayConfig Parameters + +```yaml +router_replay: + mode: "disabled" # Available options: disabled, R2, R3 + record_file: null # Path for recording routing decisions + replay_file: null # Path for replaying recorded decisions +``` + +## Quick Start Guide + +### Enabling R2 Mode + +#### Configuration File Method +Add the following to your training configuration: + +```yaml +actor: + megatron: + router_replay: + mode: "R2" +``` + +#### Command Line Method +Enable R2 mode via command-line parameters: + +```bash +actor_rollout_ref.actor.megatron.router_replay.mode="R2" +``` + +### Enabling R3 Mode + +#### Configuration File Method +Configure both actor and rollout settings: + +```yaml +# Actor configuration for megatron +actor: + megatron: + router_replay: + mode: "R3" + +# Rollout configuration +rollout: + enable_rollout_routing_replay: True +``` + +#### Command Line Method +Enable R3 mode via command-line parameters: + +```bash +actor_rollout_ref.actor.megatron.router_replay.mode="R3" +actor_rollout_ref.rollout.enable_rollout_routing_replay=True +``` + +R3 mode requires the rollout backend to support returning router selection results. Currently, this functionality is being tested based on the vllm implementation at https://github.com/vllm-project/vllm/pull/28284 as well as bug fix at https://github.com/vllm-project/vllm/pull/33013 and SGLang implementation at https://github.com/sgl-project/sglang/commit/bed301a5acaa9577c9aa706468bdf242f6a43051. diff --git a/verl/examples/router_replay/run_qwen3_30b_a3b_megatron.sh b/verl/examples/router_replay/run_qwen3_30b_a3b_megatron.sh new file mode 100644 index 0000000000000000000000000000000000000000..02a40ceb8bff1e1f5c636c22884e2f17a756ec6a --- /dev/null +++ b/verl/examples/router_replay/run_qwen3_30b_a3b_megatron.sh @@ -0,0 +1,177 @@ +#!/usr/bin/env bash +# GRPO | text | Megatron training | NVIDIA GPUs +# Router-replay example on Qwen3-30B-A3B (MoE). See README for R2 vs R3 modes. + +set -xeuo pipefail +export CUDA_DEVICE_MAX_CONNECTIONS=1 + +########################### user-adjustable ########################### +INFER_BACKEND=${INFER_BACKEND:-vllm} + +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-30B-A3B} +NNODES=${NNODES:-} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +ROUTING_REPLAY_MODE=${ROUTING_REPLAY_MODE:-R3} # R2 | R3 +ENABLE_ROLLOUT_ROUTING_REPLAY=${ENABLE_ROLLOUT_ROUTING_REPLAY:-} + +TRAIN_DATA_PATH=${TRAIN_DATA_PATH:-$HOME/data/gsm8k/train.parquet} +TEST_DATA_PATH=${TEST_DATA_PATH:-$HOME/data/gsm8k/test.parquet} + +train_batch_size=${TRAIN_BATCH_SIZE:-} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-} +micro_bs=${MICRO_BS:-3} +max_prompt_length=${MAX_PROMPT_LENGTH:-} +max_response_length=${MAX_RESPONSE_LENGTH:-} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-} + +actor_lr=${ACTOR_LR:-1e-6} + +actor_pp=${ACTOR_PP:-} +actor_tp=${ACTOR_TP:-} +actor_ep=${ACTOR_EP:-8} +actor_etp=${ACTOR_ETP:-1} +rollout_tp=${ROLLOUT_TP:-} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.65} +rollout_n=${ROLLOUT_N:-8} + +offload=${OFFLOAD:-True} + +total_training_steps=${TOTAL_TRAINING_STEPS:-50000} +project_name=${PROJECT_NAME:-verl_grpo_router_replay} +experiment_name=${EXPERIMENT_NAME:-} +########################### end user-adjustable ########################### + +########################### derived defaults ########################### +NNODES=${NNODES:-6} +train_batch_size=${train_batch_size:-3} +ppo_mini_batch_size=${ppo_mini_batch_size:-3} +max_prompt_length=${max_prompt_length:-512} +max_response_length=${max_response_length:-512} +actor_pp=${actor_pp:-6} +actor_tp=${actor_tp:-1} +rollout_tp=${rollout_tp:-4} +actor_use_dynamic_bsz=False +rollout_log_prob_use_dynamic_bsz=False +ref_log_prob_use_dynamic_bsz=False +moe_permute_fusion=False +trainer_balance_batch=False + +if [ -z "$ENABLE_ROLLOUT_ROUTING_REPLAY" ]; then + if [ "$ROUTING_REPLAY_MODE" = "R3" ]; then + ENABLE_ROLLOUT_ROUTING_REPLAY=True + else + ENABLE_ROLLOUT_ROUTING_REPLAY=False + fi +fi + +ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu:-$(((max_prompt_length + max_response_length) * 2))} +experiment_name=${experiment_name:-qwen3_30b_a3b_router_replay_${ROUTING_REPLAY_MODE}_${INFER_BACKEND}_megatron} + +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files="$TRAIN_DATA_PATH" + data.val_files="$TEST_DATA_PATH" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.use_fused_kernels=True +) + +ACTOR=( + actor_rollout_ref.actor.strategy=megatron + actor_rollout_ref.actor.model_engine=megatron + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${micro_bs} + actor_rollout_ref.actor.use_dynamic_bsz=${actor_use_dynamic_bsz} + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=False + actor_rollout_ref.actor.kl_loss_coef=0.001 + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=0 + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${actor_pp} + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.actor.megatron.expert_model_parallel_size=${actor_ep} + actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${actor_etp} + actor_rollout_ref.actor.megatron.param_offload=${offload} + actor_rollout_ref.actor.megatron.optimizer_offload=${offload} + actor_rollout_ref.actor.megatron.grad_offload=${offload} + actor_rollout_ref.actor.megatron.use_mbridge=True + actor_rollout_ref.actor.megatron.router_replay.mode=${ROUTING_REPLAY_MODE} + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_enable_deepep=True + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_token_dispatcher_type=flex + +actor_rollout_ref.actor.megatron.override_transformer_config.apply_rope_fusion=True + +actor_rollout_ref.actor.megatron.override_transformer_config.bias_activation_fusion=True + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32 + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 + +actor_rollout_ref.actor.megatron.override_transformer_config.gradient_accumulation_fusion=True + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=${moe_permute_fusion} +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=${INFER_BACKEND} + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.calculate_log_probs=True + actor_rollout_ref.rollout.enable_chunked_prefill=True + actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${micro_bs} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${rollout_log_prob_use_dynamic_bsz} + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.rollout.enable_rollout_routing_replay=${ENABLE_ROLLOUT_ROUTING_REPLAY} +) + +REF=( + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${micro_bs} + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${ref_log_prob_use_dynamic_bsz} + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${actor_pp} + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.ref.megatron.expert_model_parallel_size=${actor_ep} + actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=${actor_etp} + actor_rollout_ref.ref.megatron.param_offload=${offload} +) + +TRAINER=( + trainer.critic_warmup=0 + trainer.balance_batch=${trainer_balance_batch} + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=-1 + trainer.test_freq=10 + trainer.total_training_steps=${total_training_steps} + trainer.val_before_train=False +) + +# Conservative rollout extras shared by all inference backends. +EXTRA=( + actor_rollout_ref.rollout.skip_tokenizer_init=True +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/sapo_trainer/README.md b/verl/examples/sapo_trainer/README.md new file mode 100644 index 0000000000000000000000000000000000000000..d1b6a9c488aebe2bc0aadab2609fb37fb8139f3c --- /dev/null +++ b/verl/examples/sapo_trainer/README.md @@ -0,0 +1,20 @@ +# SAPO (Smooth Advantage Policy Optimization) + +SAPO replaces PPO's ratio clipping with a smooth, `tau`-parameterized surrogate objective. + +Reference: [Revisiting Policy Gradient Methods for Large Language Models](https://arxiv.org/pdf/2511.20347). + +## Canonical Scripts + +| Script | Infer | Train | Platform | +|---------------------------------------------|-------|-------|----------| +| `run_qwen3_8b_fsdp.sh` | vLLM | FSDP2 | Ascend | +| `run_qwen3_30b_a3b_fsdp.sh` | vLLM | FSDP2 | NVIDIA | + +## Key Flags + +- `actor_rollout_ref.actor.policy_loss.loss_mode=sapo` +- `+actor_rollout_ref.actor.policy_loss.tau_pos=1.0` +- `+actor_rollout_ref.actor.policy_loss.tau_neg=1.05` + +Note: SAPO disables ratio clipping; no `clip_ratio_low/high` needed. diff --git a/verl/examples/sapo_trainer/run_qwen3_30b_a3b_fsdp.sh b/verl/examples/sapo_trainer/run_qwen3_30b_a3b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..e5b6f4a78ce8cd5095c5ed417e6c0463b1029055 --- /dev/null +++ b/verl/examples/sapo_trainer/run_qwen3_30b_a3b_fsdp.sh @@ -0,0 +1,116 @@ +#!/usr/bin/env bash +# SAPO | MoE | vLLM rollout | FSDP training | NVIDIA GPUs +# SAPO (Smooth Advantage PO) replaces ratio clipping with a smooth tau-parameterized surrogate (arXiv:2511.20347). + +set -xeuo pipefail +export VLLM_USE_V1=1 + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-30B-A3B-Base} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +tau_pos=${TAU_POS:-1.0} +tau_neg=${TAU_NEG:-1.05} + +train_batch_size=${TRAIN_BATCH_SIZE:-256} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-32} +max_prompt_length=${MAX_PROMPT_LENGTH:-2048} +max_response_length=${MAX_RESPONSE_LENGTH:-8192} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-20480} + +actor_lr=${ACTOR_LR:-1e-6} +entropy_coeff=${ENTROPY_COEFF:-0} + +rollout_tp=${ROLLOUT_TP:-4} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.8} +rollout_n=${ROLLOUT_N:-16} + +total_epochs=${TOTAL_EPOCHS:-10} +save_freq=${SAVE_FREQ:-50} +test_freq=${TEST_FREQ:-10} + +project_name=${PROJECT_NAME:-verl_sapo_qwen3_moe} +experiment_name=${EXPERIMENT_NAME:-qwen3_30b_a3b_vllm_fsdp} +# ---- end user-adjustable ---- + +train_file=${TRAIN_FILE:-$HOME/data/dapo-math-17k/train.parquet} +val_file=${VAL_FILE:-$HOME/data/aime-2024/test.parquet} +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files="['$train_file']" + data.val_files="['$val_file']" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.policy_loss.loss_mode=sapo + +actor_rollout_ref.actor.policy_loss.tau_pos=${tau_pos} + +actor_rollout_ref.actor.policy_loss.tau_neg=${tau_neg} + actor_rollout_ref.actor.strategy=fsdp2 + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=False + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.fsdp_config.param_offload=True + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.strategy=fsdp2 + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.fsdp_config.param_offload=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.val_before_train=False + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/sapo_trainer/run_qwen3_8b_fsdp.sh b/verl/examples/sapo_trainer/run_qwen3_8b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..679cf8c80a5c3331ede3464a395760c6d5599c94 --- /dev/null +++ b/verl/examples/sapo_trainer/run_qwen3_8b_fsdp.sh @@ -0,0 +1,129 @@ +#!/usr/bin/env bash +# SAPO | text | vLLM rollout | FSDP training | Ascend NPU + +set -xeuo pipefail +ulimit -n 32768 + +export RAY_DEDUP_LOGS=0 +export HYDRA_FULL_ERROR=1 +export TASK_QUEUE_ENABLE=1 +export HCCL_EXEC_TIMEOUT=3600 +export HCCL_CONNECT_TIMEOUT=3600 +export HCCL_ASYNC_ERROR_HANDLING=0 +export CPU_AFFINITY_CONF=1 +export VLLM_USE_V1=1 + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-8B} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +tau_pos=${TAU_POS:-1.0} +tau_neg=${TAU_NEG:-1.05} + +train_batch_size=${TRAIN_BATCH_SIZE:-256} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-64} +max_prompt_length=${MAX_PROMPT_LENGTH:-2048} +max_response_length=${MAX_RESPONSE_LENGTH:-8192} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-20480} + +actor_lr=${ACTOR_LR:-1e-6} +entropy_coeff=${ENTROPY_COEFF:-0} + +rollout_tp=${ROLLOUT_TP:-2} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.7} +rollout_n=${ROLLOUT_N:-16} + +total_epochs=${TOTAL_EPOCHS:-10} +save_freq=${SAVE_FREQ:-100} +test_freq=${TEST_FREQ:--1} + +project_name=${PROJECT_NAME:-verl_sapo_gsm8k_math} +experiment_name=${EXPERIMENT_NAME:-qwen3_8b_vllm_fsdp} +# ---- end user-adjustable ---- + +gsm8k_train=$HOME/data/gsm8k/train.parquet +gsm8k_test=$HOME/data/gsm8k/test.parquet +math_train=$HOME/data/math/train.parquet +math_test=$HOME/data/math/test.parquet + +train_files="['$gsm8k_train', '$math_train']" +val_files="['$gsm8k_test', '$math_test']" +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files="$train_files" + data.val_files="$val_files" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.policy_loss.loss_mode=sapo + +actor_rollout_ref.actor.policy_loss.tau_pos=${tau_pos} + +actor_rollout_ref.actor.policy_loss.tau_neg=${tau_neg} + actor_rollout_ref.actor.strategy=fsdp2 + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=False + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.fsdp_config.param_offload=True + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.strategy=fsdp2 + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.fsdp_config.param_offload=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.val_before_train=False + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/sft/gsm8k/run_deepseek_coder_6_7b_fsdp.sh b/verl/examples/sft/gsm8k/run_deepseek_coder_6_7b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..963f884cc8eccc6a02f7309f9eceb01de4d5dfcc --- /dev/null +++ b/verl/examples/sft/gsm8k/run_deepseek_coder_6_7b_fsdp.sh @@ -0,0 +1,27 @@ +set -x + +if [ "$#" -lt 2 ]; then + echo "Usage: run_deepseek_coder_6_7b_fsdp.sh [other_configs...]" + exit 1 +fi + +nproc_per_node=$1 +save_path=$2 + +# Shift the arguments so $@ refers to the rest +shift 2 + +torchrun --standalone --nnodes=1 --nproc_per_node=$nproc_per_node \ + -m verl.trainer.sft_trainer \ + data.train_files=$HOME/data/gsm8k/train.parquet \ + data.val_files=$HOME/data/gsm8k/test.parquet \ + data.messages_key=messages \ + data.micro_batch_size_per_gpu=4 \ + optim.lr=1e-4 \ + engine=fsdp \ + model.path=deepseek-ai/deepseek-coder-6.7b-instruct \ + trainer.default_local_dir=$save_path \ + trainer.project_name=gsm8k-sft \ + trainer.experiment_name=gsm8k-sft-deepseek-coder-6.7b-instruct \ + trainer.total_epochs=4 \ + trainer.logger='["console","wandb"]' $@ \ No newline at end of file diff --git a/verl/examples/sft/gsm8k/run_gemma_2b_fsdp.sh b/verl/examples/sft/gsm8k/run_gemma_2b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..e0fd3b13938786293da03a770731548aaec00b36 --- /dev/null +++ b/verl/examples/sft/gsm8k/run_gemma_2b_fsdp.sh @@ -0,0 +1,29 @@ +# Tested with 2 & 4 GPUs + +set -x + +if [ "$#" -lt 2 ]; then + echo "Usage: run_gemma_2b_fsdp.sh [other_configs...]" + exit 1 +fi + +nproc_per_node=$1 +save_path=$2 + +# Shift the arguments so $@ refers to the rest +shift 2 + +torchrun --standalone --nnodes=1 --nproc_per_node=$nproc_per_node \ + -m verl.trainer.sft_trainer \ + data.train_files=$HOME/data/gsm8k/train.parquet \ + data.val_files=$HOME/data/gsm8k/test.parquet \ + data.messages_key=messages \ + data.micro_batch_size_per_gpu=4 \ + model.path=google/gemma-2b-it \ + optim.lr=1e-4 \ + engine=fsdp \ + trainer.default_local_dir=$save_path \ + trainer.project_name=gsm8k-sft \ + trainer.experiment_name=gsm8k-sft-gemma-2b-it \ + trainer.total_epochs=2 \ + trainer.logger='["console","wandb"]' $@ \ No newline at end of file diff --git a/verl/examples/sft/gsm8k/run_gemma_7b_fsdp.sh b/verl/examples/sft/gsm8k/run_gemma_7b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..30f881cdc9bcd83837a2b7e8438fe0fd54b07c0a --- /dev/null +++ b/verl/examples/sft/gsm8k/run_gemma_7b_fsdp.sh @@ -0,0 +1,27 @@ +set -x + +if [ "$#" -lt 2 ]; then + echo "Usage: run_gemma_7b_fsdp.sh [other_configs...]" + exit 1 +fi + +nproc_per_node=$1 +save_path=$2 + +# Shift the arguments so $@ refers to the rest +shift 2 + +torchrun --standalone --nnodes=1 --nproc_per_node=$nproc_per_node \ + -m verl.trainer.sft_trainer \ + data.train_files=$HOME/data/gsm8k/train.parquet \ + data.val_files=$HOME/data/gsm8k/test.parquet \ + data.messages_key=messages \ + data.micro_batch_size_per_gpu=4 \ + optim.lr=1e-4 \ + engine=fsdp \ + model.path=google/gemma-1.1-7b-it \ + trainer.default_local_dir=$save_path \ + trainer.project_name=gsm8k-sft \ + trainer.experiment_name=gsm8k-sft-gemma-1.1-7b-it \ + trainer.total_epochs=4 \ + trainer.logger='["console","wandb"]' $@ diff --git a/verl/examples/sft/gsm8k/run_mimo_7b_mtp_megatron.sh b/verl/examples/sft/gsm8k/run_mimo_7b_mtp_megatron.sh new file mode 100644 index 0000000000000000000000000000000000000000..68b4693fb3361ed0023815460e3d19b41d62ee60 --- /dev/null +++ b/verl/examples/sft/gsm8k/run_mimo_7b_mtp_megatron.sh @@ -0,0 +1,102 @@ +#!/usr/bin/env bash +set -xeuo pipefail + +NUM_GPUS=${NUM_GPUS:-8} +SP_SIZE=${SP_SIZE:-1} +TP_SIZE=${TP_SIZE:-1} +PP_SIZE=${PP_SIZE:-1} +VPP_SIZE=${VPP_SIZE:-null} +CP_SIZE=${CP_SIZE:-1} +PAD_MODE=${PAD_MODE:-no_padding} +USE_REMOVE_PADDING=${USE_REMOVE_PADDING:-False} +LR="1e-5" +MINLR="1e-6" + +export VERL_SFT_LOGGING_LEVEL=INFO + +backend=${BACKEND:-megatron} + +TENSORBOARD_DIR=~/tensorboard + +MASTER_ADDR=${MASTER_ADDR:-localhost} +MASTER_PORT=${MASTER_PORT:-29500} +NNODES=${NNODES:-1} +RANK=${RANK:-0} + +ENTRYPOINT=${ENTRYPOINT:-"-m verl.trainer.sft_trainer"} + +# Note the default MultiturnSFT Dataset requires all the sys/user/assistant in 'data.message_key' +DATASET_DIR=${DATASET_DIR:-~/dataset/rl/gsm8k} +TRAIN_FILES=${DATASET_DIR}/train.parquet +VAL_FILES=${DATASET_DIR}/eval.parquet + +project_name=verl_sft_test + +RESUME_MODE=disable + +MODEL_PATH="XiaomiMiMo/MiMo-7B-RL" +ckpts_home=${ckpts_home:-~/verl/test/gsm8k-sft-${backend}} + +# currently relies on these two commits that is not on master +PYPATH=$HOME/pythonpath +mkdir -p $PYPATH && cd $PYPATH +[ -d Megatron-LM ] || git clone https://github.com/NVIDIA/Megatron-LM -b dev && (cd Megatron-LM; git checkout 23e092f41ec8bc659020e401ddac9576c1cfed7e) +[ -d mbridge ] || git clone https://github.com/ArronHZG/mbridge -b feature/verl_mtp && (cd mbridge; git checkout 6bf2d45a15dc4fb52d2f0c38ff546bee33447d10) +cd - +export PYTHONPATH=$PYTHONPATH:$PYPATH/mbridge:$PYPATH/Megatron-LM + + +MEGATRON_ENGINE_CONFIG="\ + engine=${backend} \ + optim=${backend} \ + optim.lr=${LR} \ + optim.min_lr=${MINLR} \ + optim.lr_warmup_steps=10 \ + optim.weight_decay=0.1 \ + optim.betas='[0.9,0.95]' \ + optim.clip_grad=1.0 \ + optim.lr_warmup_init=0 \ + optim.lr_decay_style=cosine \ + engine.override_transformer_config.recompute_method=uniform \ + engine.override_transformer_config.recompute_granularity=full \ + engine.override_transformer_config.recompute_num_layers=1 \ + engine.use_dist_checkpointing=False \ + engine.tensor_model_parallel_size=${TP_SIZE} \ + engine.pipeline_model_parallel_size=${PP_SIZE} \ + engine.virtual_pipeline_model_parallel_size=${VPP_SIZE} \ + engine.context_parallel_size=${CP_SIZE} \ + engine.use_mbridge=True \ + " + +ENGINE_CONFIG="$MEGATRON_ENGINE_CONFIG" +echo "Using megatron engine" +exp_name=gsm8k-${backend}-tp${TP_SIZE}-pp${PP_SIZE}-vpp${VPP_SIZE}-cp${CP_SIZE}-lr-${MINLR}-${LR} + +mkdir -p "${ckpts_home}" + +$COMMAND \ + data.train_files="${TRAIN_FILES}" \ + data.val_files="${TRAIN_FILES}" \ + data.train_batch_size=64 \ + data.micro_batch_size_per_gpu=2 \ + data.pad_mode=${PAD_MODE} \ + data.truncation=error \ + data.max_length=1024 \ + data.use_dynamic_bsz=True \ + data.max_token_len_per_gpu=2048 \ + data.messages_key=prompt \ + data.num_workers=0 \ + model.path=$MODEL_PATH \ + model.use_remove_padding=${USE_REMOVE_PADDING} \ + model.trust_remote_code=True \ + model.mtp.enable=True \ + ${ENGINE_CONFIG} \ + trainer.test_freq=after_each_epoch \ + trainer.save_freq=-1 \ + trainer.logger='["console","wandb"]' \ + trainer.project_name="${project_name}" \ + trainer.experiment_name="${exp_name}" \ + trainer.total_epochs=1 \ + trainer.default_local_dir="${ckpts_home}" \ + trainer.resume_mode=${RESUME_MODE} + \ No newline at end of file diff --git a/verl/examples/sft/gsm8k/run_nemotron_nano_v3_megatron.sh b/verl/examples/sft/gsm8k/run_nemotron_nano_v3_megatron.sh new file mode 100644 index 0000000000000000000000000000000000000000..8aa277fc3f5124779af053c3f7d1ff03b619a51e --- /dev/null +++ b/verl/examples/sft/gsm8k/run_nemotron_nano_v3_megatron.sh @@ -0,0 +1,98 @@ +#!/usr/bin/env bash +set -xeuo pipefail +################################################### environment ################################################### +### # 1. use docker image `verlai/verl:vllm015.dev`` and install correct dependencies: +# pip install nvidia-modelopt +# MAX_JOBS=32 pip install git+https://github.com/Dao-AILab/causal-conv1d.git --no-build-isolation --no-cache-dir +# MAX_JOBS=32 pip install git+https://github.com/state-spaces/mamba.git --no-build-isolation --no-cache-dir +# pip install --no-deps git+https://github.com/NVIDIA-NeMo/Megatron-Bridge +# pip install --no-deps git+https://github.com/NVIDIA/Megatron-LM.git@core_dev_r0.16.0 +# unset ROCR_VISIBLE_DEVICES +# unset PYTORCH_CUDA_ALLOC_CONF + +ENTRYPOINT=${ENTRYPOINT:-"-m verl.trainer.sft_trainer"} + +TRAIN_FILES=$HOME/data/gsm8k/train.parquet +VAL_FILES=$HOME/data/gsm8k/eval.parquet + +backend=${BACKEND:-megatron} + +project_name=verl_sft_gsm8k + +RESUME_MODE=auto +MODEL_NAME=${MODEL_NAME:-NVIDIA-Nemotron-3-Nano-30B-A3B-BF16} +MODEL_PATH=${MODEL_PATH:-nvidia/${MODEL_NAME}} + +TP_SIZE=${TP_SIZE:-8} +PP_SIZE=${PP_SIZE:-1} +VPP_SIZE=${VPP_SIZE:-null} +CP_SIZE=${CP_SIZE:-1} +EP_SIZE=${EP_SIZE:-8} +ETP_SIZE=${ETP_SIZE:-1} + +PAD_MODE=${PAD_MODE:-no_padding} + +USE_REMOVE_PADDING=${USE_REMOVE_PADDING:-True} + +DTYPE=${DTYPE:-"bfloat16"} + + +MEGATRON_ENGINE_CONFIG=( + engine=${backend} + optim=${backend} + optim.lr=2e-5 + optim.lr_warmup_steps=5 + optim.weight_decay=0.1 + optim.betas="[0.9,0.95]" + optim.clip_grad=1.0 + optim.lr_warmup_init=0 + optim.lr_decay_style=cosine + optim.min_lr=2e-6 + +optim.override_optimizer_config.optimizer_offload_fraction=1 + +optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True + +optim.override_optimizer_config.use_precision_aware_optimizer=True + +optim.override_optimizer_config.optimizer_cpu_offload=True + engine.tensor_model_parallel_size=${TP_SIZE} + engine.pipeline_model_parallel_size=${PP_SIZE} + engine.virtual_pipeline_model_parallel_size=${VPP_SIZE} + engine.context_parallel_size=${CP_SIZE} + engine.use_mbridge=True + engine.dtype=${DTYPE} + engine.vanilla_mbridge=False + engine.expert_model_parallel_size=${EP_SIZE} + engine.expert_tensor_parallel_size=${ETP_SIZE} + engine.override_transformer_config.attention_backend=auto + +engine.override_transformer_config.recompute_method=uniform + +engine.override_transformer_config.recompute_granularity=full + +engine.override_transformer_config.recompute_num_layers=1 +) + +ENGINE_CONFIG="${MEGATRON_ENGINE_CONFIG[@]}" +echo "Using megatron engine" +exp_name=${MODEL_NAME}-${backend}-tp${TP_SIZE}-pp${PP_SIZE}-vpp${VPP_SIZE}-cp${CP_SIZE}-megatron-20260210 + +torchrun --nnodes=1 --nproc_per_node=8 ${ENTRYPOINT} \ + data.train_files="${TRAIN_FILES}" \ + data.val_files="${VAL_FILES}" \ + data.train_batch_size=96 \ + data.max_length=2048 \ + data.pad_mode=${PAD_MODE} \ + data.truncation=error \ + data.use_dynamic_bsz=True \ + data.max_token_len_per_gpu=2048 \ + data.messages_key=messages \ + data.ignore_input_ids_mismatch=True \ + model.path=$MODEL_PATH \ + model.use_remove_padding=${USE_REMOVE_PADDING} \ + model.trust_remote_code=True \ + ${ENGINE_CONFIG} \ + trainer.test_freq=-1 \ + trainer.save_freq=500 \ + trainer.logger=['console','wandb'] \ + trainer.project_name="${project_name}" \ + trainer.experiment_name="${exp_name}" \ + trainer.total_epochs=1 \ + trainer.default_local_dir="${ckpts_home}" \ + trainer.resume_mode=${RESUME_MODE} \ + trainer.max_ckpt_to_keep=10 \ + checkpoint.save_contents=[model,optimizer,extra] \ No newline at end of file diff --git a/verl/examples/sft/gsm8k/run_qwen2_5_0_5b_automodel.sh b/verl/examples/sft/gsm8k/run_qwen2_5_0_5b_automodel.sh new file mode 100644 index 0000000000000000000000000000000000000000..ab0d1a08eb59290fa3bdaab3c5099f706d83f7d7 --- /dev/null +++ b/verl/examples/sft/gsm8k/run_qwen2_5_0_5b_automodel.sh @@ -0,0 +1,55 @@ +# Requires: Automodel, transformers>=5.3.0, torchao +# MoE also requires: grouped_gemm (github.com/fanshiqing/grouped_gemm v1.1.4) + +set -x + +if [ "$#" -lt 2 ]; then + echo "Usage: run_qwen2_5_0_5b_automodel.sh [other_configs...]" + exit 1 +fi + +nproc_per_node=$1 +save_path=$2 + +# Shift the arguments so $@ refers to the rest +shift 2 + +torchrun --standalone --nnodes=1 --nproc_per_node=$nproc_per_node \ + -m verl.trainer.sft_trainer \ + data.train_files=$HOME/data/gsm8k_sft/train.parquet \ + data.val_files=$HOME/data/gsm8k_sft/test.parquet \ + data.train_batch_size=128 \ + data.pad_mode=no_padding \ + data.truncation=error \ + data.use_dynamic_bsz=True \ + data.max_token_len_per_gpu=2048 \ + data.messages_key=messages \ + data.ignore_input_ids_mismatch=True \ + model=hf_model \ + model.path=Qwen/Qwen2.5-0.5B-Instruct \ + model.use_remove_padding=True \ + engine=automodel \ + engine.distributed_strategy=fsdp2 \ + engine.tp_size=1 \ + engine.pp_size=1 \ + engine.cp_size=1 \ + engine.ep_size=1 \ + engine.use_torch_compile=False \ + optim=automodel \ + optim.lr=1e-5 \ + optim.lr_warmup_steps_ratio=0.2 \ + optim.weight_decay=0.1 \ + optim.betas='[0.9,0.95]' \ + optim.clip_grad=1.0 \ + optim.init_lr_ratio=0 \ + optim.min_lr_ratio=0.1 \ + optim.lr_scheduler_type=cosine \ + trainer.default_local_dir=$save_path \ + trainer.project_name=gsm8k-sft \ + trainer.experiment_name=gsm8k-sft-qwen-2.5-0.5b-automodel \ + trainer.total_epochs=2 \ + trainer.test_freq=-1 \ + trainer.save_freq=-1 \ + trainer.logger=console \ + trainer.seed=1111 \ + trainer.resume_mode=disable $@ diff --git a/verl/examples/sft/gsm8k/run_qwen2_5_0_5b_fsdp.sh b/verl/examples/sft/gsm8k/run_qwen2_5_0_5b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..f4388ea9302c73795269b1d53347de2ad839dae8 --- /dev/null +++ b/verl/examples/sft/gsm8k/run_qwen2_5_0_5b_fsdp.sh @@ -0,0 +1,73 @@ +#!/usr/bin/env bash +# SFT | GSM8K | FSDP engine | NVIDIA GPUs +# Covers plain SFT, Ulysses sequence parallel, Liger-kernel, and LoRA/PEFT via env vars. +# +# Examples: +# # plain SFT +# bash run_qwen2_5_0_5b_fsdp.sh 8 /tmp/sft-ckpt +# +# # sequence-parallel (Ulysses) = 2 +# SP_SIZE=2 bash run_qwen2_5_0_5b_fsdp.sh 8 /tmp/sft-ckpt +# +# # sequence-parallel + Liger kernel +# SP_SIZE=2 USE_LIGER=1 bash run_qwen2_5_0_5b_fsdp.sh 8 /tmp/sft-ckpt +# +# # LoRA (PEFT) +# USE_PEFT=1 bash run_qwen2_5_0_5b_fsdp.sh 8 /tmp/sft-ckpt + +set -xeuo pipefail + +if [ "$#" -lt 2 ]; then + echo "Usage: run_qwen2_5_0_5b_fsdp.sh [other_configs...]" + echo " Env: SP_SIZE (default 1), USE_LIGER (0|1), USE_PEFT (0|1)" + exit 1 +fi + +nproc_per_node=$1 +save_path=$2 +shift 2 + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen2.5-0.5B-Instruct} +SP_SIZE=${SP_SIZE:-1} +USE_LIGER=${USE_LIGER:-0} +USE_PEFT=${USE_PEFT:-0} +LORA_RANK=${LORA_RANK:-32} +LORA_ALPHA=${LORA_ALPHA:-16} +LORA_TARGETS=${LORA_TARGETS:-all-linear} +MICRO_BATCH_SIZE_PER_GPU=${MICRO_BATCH_SIZE_PER_GPU:-4} +LR=${LR:-1e-4} +TOTAL_EPOCHS=${TOTAL_EPOCHS:-1} +PROJECT_NAME=${PROJECT_NAME:-gsm8k-sft} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-gsm8k-sft-qwen2_5_0_5b} +# ---- end user-adjustable ---- + +extra_args=() +if [ "${USE_LIGER}" = "1" ]; then + extra_args+=("model.use_liger=True") +fi +if [ "${USE_PEFT}" = "1" ]; then + extra_args+=( + "model.lora_rank=${LORA_RANK}" + "model.lora_alpha=${LORA_ALPHA}" + "model.target_modules=${LORA_TARGETS}" + ) +fi + +torchrun --standalone --nnodes=1 --nproc_per_node=${nproc_per_node} \ + -m verl.trainer.sft_trainer \ + data.train_files=$HOME/data/gsm8k/train.parquet \ + data.val_files=$HOME/data/gsm8k/test.parquet \ + data.messages_key=messages \ + data.micro_batch_size_per_gpu=${MICRO_BATCH_SIZE_PER_GPU} \ + optim.lr=${LR} \ + engine=fsdp \ + engine.ulysses_sequence_parallel_size=${SP_SIZE} \ + model.path="${MODEL_PATH}" \ + model.use_remove_padding=true \ + trainer.default_local_dir="${save_path}" \ + trainer.project_name="${PROJECT_NAME}" \ + trainer.experiment_name="${EXPERIMENT_NAME}" \ + trainer.logger='["console","wandb"]' \ + trainer.total_epochs=${TOTAL_EPOCHS} \ + "${extra_args[@]}" "$@" diff --git a/verl/examples/sft/gsm8k/run_qwen3_30b_automodel.sh b/verl/examples/sft/gsm8k/run_qwen3_30b_automodel.sh new file mode 100644 index 0000000000000000000000000000000000000000..95d699d218afd86a29173c9b1d6f7d342179a800 --- /dev/null +++ b/verl/examples/sft/gsm8k/run_qwen3_30b_automodel.sh @@ -0,0 +1,75 @@ +# Requires: Automodel, transformers>=5.3.0, torchao +# MoE also requires: grouped_gemm (github.com/fanshiqing/grouped_gemm v1.1.4) + +set -x + +if [ "$#" -lt 2 ]; then + echo "Usage: run_qwen3_30b_automodel.sh [other_configs...]" + exit 1 +fi + +nproc_per_node=$1 +save_path=$2 + +# Shift the arguments so $@ refers to the rest +shift 2 + +torchrun --standalone --nnodes=1 --nproc_per_node=$nproc_per_node \ + -m verl.trainer.sft_trainer \ + data.train_files=$HOME/data/hellaswag_sft/hellaswag_sft.parquet \ + data.val_files=$HOME/data/hellaswag_sft/hellaswag_sft.parquet \ + data.train_batch_size=512 \ + data.max_length=2048 \ + data.truncation=left \ + data.use_dynamic_bsz=True \ + data.max_token_len_per_gpu=8192 \ + data.messages_key=messages \ + data.ignore_input_ids_mismatch=True \ + data.train_max_samples=-1 \ + data.val_max_samples=1024 \ + model=hf_model \ + model.path=Qwen/Qwen3-30B-A3B-Base \ + model.trust_remote_code=True \ + model.use_remove_padding=True \ + engine=automodel \ + engine.distributed_strategy=fsdp2 \ + engine.tp_size=1 \ + engine.pp_size=1 \ + engine.cp_size=1 \ + engine.ep_size=8 \ + engine.backend_config.dispatcher=deepep \ + engine.backend_config.attn=te \ + engine.backend_config.linear=te \ + engine.backend_config.rms_norm=torch_fp32 \ + engine.backend_config.enable_fsdp_optimizations=True \ + engine.backend_config.experts=torch_mm \ + engine.activation_checkpointing=True \ + engine.model_dtype=bf16 \ + engine.attn_implementation=te \ + engine.use_torch_compile=False \ + optim=automodel \ + optim.optimizer=FusedAdam \ + optim.optimizer_impl=transformer_engine.pytorch.optimizers.fused_adam \ + optim.lr=1e-5 \ + optim.lr_warmup_steps_ratio=0.1 \ + optim.weight_decay=0 \ + optim.betas='[0.9,0.95]' \ + optim.clip_grad=1.0 \ + optim.init_lr_ratio=0.1 \ + optim.min_lr_ratio=0.01 \ + optim.lr_scheduler_type=cosine \ + optim.master_weights=true \ + optim.store_param_remainders=true \ + optim.exp_avg_dtype=bf16 \ + optim.exp_avg_sq_dtype=bf16 \ + trainer.default_local_dir=$save_path \ + trainer.project_name=hellaswag-sft \ + trainer.experiment_name=hellaswag-sft-qwen3-30b-automodel \ + trainer.total_epochs=2 \ + trainer.total_training_steps=100 \ + trainer.save_freq=-1 \ + trainer.test_freq=10 \ + trainer.logger=console \ + trainer.seed=1111 \ + trainer.nnodes=1 \ + trainer.resume_mode=disable $@ diff --git a/verl/examples/sft/gsm8k/run_qwen3_5_397b_a17b_megatron.sh b/verl/examples/sft/gsm8k/run_qwen3_5_397b_a17b_megatron.sh new file mode 100644 index 0000000000000000000000000000000000000000..c6d9bfedd8b2791d998be7096676de461bfbb134 --- /dev/null +++ b/verl/examples/sft/gsm8k/run_qwen3_5_397b_a17b_megatron.sh @@ -0,0 +1,160 @@ +#!/usr/bin/env bash +# Qwen3.5-397B-A17B SFT with Megatron backend + mbridge +# +# Requirements: +# - 128+ GPUs (80GB each, e.g. 16x8 H100/H200) +# - Docker: verlai/verl:vllm015 (or equivalent) +# - Additional packages on top of the base image: +# pip install --upgrade transformers +# pip install flash-linear-attention +# pip install -U git+https://github.com/ISEEKYAN/mbridge.git +# - Megatron-LM==0.16.0 +# +# Qwen3.5 architecture notes: +# Qwen3.5 uses Gated Delta Net (GDN) linear attention which currently does +# NOT support packed sequences (THD format) in Megatron-LM. Therefore: +# - engine.use_remove_padding=False (forces bshd compute format) +# - data.use_dynamic_bsz=False (required for bshd mode) +# +# Once https://github.com/NVIDIA/Megatron-LM/pull/2644 is merged, THD +# format will be supported and engine.use_remove_padding can be set to True +# for better performance. +# +# MTP (Multi-Token Prediction) notes: +# - model.mtp.enable=True enables MTP module +# - model.mtp.enable_train=True enables MTP training loss +# - model.mtp.detach_encoder=True detaches encoder gradients for MTP +# - model.mtp.mtp_loss_scaling_factor weight of MTP auxiliary loss (e.g. 0.1) +# +# Tested parallelism config (128 GPUs / 16 nodes): +# TP=2 PP=4 EP=32 CP=1 + +set -xeuo pipefail + +# ============================================================ +# Distributed +# ============================================================ +NUM_GPUS=${NUM_GPUS:-8} +MASTER_ADDR=${MASTER_ADDR:-localhost} +MASTER_PORT=${MASTER_PORT:-29500} +NNODES=${NNODES:-16} +NODE_RANK=${NODE_RANK:-0} + +# ============================================================ +# Data +# ============================================================ +DATASET_DIR=${DATASET_DIR:-~/dataset} +TRAIN_FILES=${TRAIN_FILES:-${DATASET_DIR}/train.parquet} + +# ============================================================ +# Model +# ============================================================ +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3.5-397B-A17B} + +# ============================================================ +# Parallelism +# ============================================================ +TP_SIZE=${TP_SIZE:-2} +PP_SIZE=${PP_SIZE:-4} +VPP_SIZE=${VPP_SIZE:-null} +CP_SIZE=${CP_SIZE:-1} +EP_SIZE=${EP_SIZE:-32} +ETP_SIZE=${ETP_SIZE:-1} + +# ============================================================ +# Training +# ============================================================ +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-128} +MICRO_BATCH_SIZE=${MICRO_BATCH_SIZE:-2} +MAX_LENGTH=${MAX_LENGTH:-2048} +LR=${LR:-2e-5} +MIN_LR=${MIN_LR:-2e-6} +DTYPE=${DTYPE:-bfloat16} + +BACKEND=megatron +RESUME_MODE=${RESUME_MODE:-disable} + +project_name=verl_sft_qwen3_5 +exp_name=qwen3_5-${BACKEND}-tp${TP_SIZE}-pp${PP_SIZE}-cp${CP_SIZE}-ep${EP_SIZE} +ckpts_home=${ckpts_home:-~/verl/checkpoints/${project_name}/${exp_name}} +mkdir -p "${ckpts_home}" + +# ============================================================ +# MTP hyper-parameters +# ============================================================ +MTP_ENABLE=${MTP_ENABLE:-True} +MTP_ENABLE_TRAIN=${MTP_ENABLE_TRAIN:-True} +MTP_DETACH_ENCODER=${MTP_DETACH_ENCODER:-True} +MTP_LOSS_SCALING_FACTOR=${MTP_LOSS_SCALING_FACTOR:-0.1} + +# ============================================================ +# Engine config +# ============================================================ +# Key Qwen3.5 settings: +# engine.use_remove_padding=False - GDN requires bshd format (no THD) +# engine.vanilla_mbridge=True - use mbridge (not megatron-bridge) +ENGINE_CONFIG="\ + engine=${BACKEND} \ + optim=${BACKEND} \ + optim.lr=${LR} \ + optim.min_lr=${MIN_LR} \ + optim.lr_warmup_steps=10 \ + optim.weight_decay=0.1 \ + optim.betas='[0.9,0.95]' \ + optim.clip_grad=1.0 \ + optim.lr_warmup_init=0 \ + optim.lr_decay_style=cosine \ + +optim.override_optimizer_config.optimizer_offload_fraction=1 \ + +optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True \ + +optim.override_optimizer_config.use_precision_aware_optimizer=True \ + +optim.override_optimizer_config.optimizer_cpu_offload=True \ + engine.tensor_model_parallel_size=${TP_SIZE} \ + engine.pipeline_model_parallel_size=${PP_SIZE} \ + engine.virtual_pipeline_model_parallel_size=${VPP_SIZE} \ + engine.context_parallel_size=${CP_SIZE} \ + engine.expert_model_parallel_size=${EP_SIZE} \ + engine.expert_tensor_parallel_size=${ETP_SIZE} \ + engine.use_mbridge=True \ + engine.vanilla_mbridge=True \ + engine.dtype=${DTYPE} \ + engine.use_remove_padding=False \ + engine.override_transformer_config.attention_backend=auto \ + +engine.override_transformer_config.recompute_method=uniform \ + +engine.override_transformer_config.recompute_granularity=full \ + +engine.override_transformer_config.recompute_num_layers=1" + +# ============================================================ +# Launch +# ============================================================ +torchrun \ + --nproc_per_node=${NUM_GPUS} \ + --nnodes=${NNODES} \ + --node_rank=${NODE_RANK} \ + --master_addr=${MASTER_ADDR} \ + --master_port=${MASTER_PORT} \ + -m verl.trainer.sft_trainer \ + data.train_files="${TRAIN_FILES}" \ + data.train_batch_size=${TRAIN_BATCH_SIZE} \ + data.micro_batch_size_per_gpu=${MICRO_BATCH_SIZE} \ + data.max_length=${MAX_LENGTH} \ + data.pad_mode=no_padding \ + data.truncation=error \ + data.use_dynamic_bsz=False \ + data.max_token_len_per_gpu=${MAX_LENGTH} \ + data.messages_key=messages \ + model.path=${MODEL_PATH} \ + model.use_remove_padding=False \ + model.trust_remote_code=True \ + model.mtp.enable=${MTP_ENABLE} \ + model.mtp.enable_train=${MTP_ENABLE_TRAIN} \ + model.mtp.detach_encoder=${MTP_DETACH_ENCODER} \ + model.mtp.mtp_loss_scaling_factor=${MTP_LOSS_SCALING_FACTOR} \ + ${ENGINE_CONFIG} \ + trainer.test_freq=-1 \ + trainer.save_freq=500 \ + trainer.logger='["console","wandb"]' \ + trainer.project_name="${project_name}" \ + trainer.experiment_name="${exp_name}" \ + trainer.total_epochs=1 \ + trainer.default_local_dir="${ckpts_home}" \ + trainer.resume_mode=${RESUME_MODE} diff --git a/verl/examples/sft/gsm8k/run_qwen3_8b_fsdp.sh b/verl/examples/sft/gsm8k/run_qwen3_8b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..e08c0e882d60c093f1e72c7ec358b0d204f6e5cf --- /dev/null +++ b/verl/examples/sft/gsm8k/run_qwen3_8b_fsdp.sh @@ -0,0 +1,64 @@ +#!/usr/bin/env bash +# SFT | GSM8K | FSDP engine | Ascend NPU +# Toggle Ulysses sequence parallel and LoRA/PEFT via env vars. +# +# Examples: +# # plain SFT +# bash run_qwen3_8b_fsdp.sh 8 /tmp/sft-ckpt +# +# # sequence-parallel (Ulysses) = 2 + LoRA (the default demo) +# SP_SIZE=2 USE_PEFT=1 bash run_qwen3_8b_fsdp.sh 8 /tmp/sft-ckpt + +set -xeuo pipefail + +if [ "$#" -lt 2 ]; then + echo "Usage: run_qwen3_8b_fsdp.sh [other_configs...]" + echo " Env: SP_SIZE (default 2), USE_PEFT (0|1, default 1)" + exit 1 +fi + +nproc_per_node=$1 +save_path=$2 +shift 2 + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-8B} +SP_SIZE=${SP_SIZE:-2} +USE_PEFT=${USE_PEFT:-1} +LORA_RANK=${LORA_RANK:-32} +LORA_ALPHA=${LORA_ALPHA:-16} +LORA_TARGETS=${LORA_TARGETS:-all-linear} +MICRO_BATCH_SIZE_PER_GPU=${MICRO_BATCH_SIZE_PER_GPU:-64} +LR=${LR:-1e-4} +TOTAL_EPOCHS=${TOTAL_EPOCHS:-2} +PROJECT_NAME=${PROJECT_NAME:-gsm8k-sft} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-gsm8k-sft-qwen3-8b-instruct} +# ---- end user-adjustable ---- + +extra_args=() +if [ "${USE_PEFT}" = "1" ]; then + extra_args+=( + "model.lora_rank=${LORA_RANK}" + "model.lora_alpha=${LORA_ALPHA}" + "model.target_modules=${LORA_TARGETS}" + ) +fi + +torchrun --standalone --nnodes=1 --nproc_per_node=${nproc_per_node} \ + -m verl.trainer.sft_trainer \ + data.train_files=$HOME/data/gsm8k/train.parquet \ + data.val_files=$HOME/data/gsm8k/test.parquet \ + data.micro_batch_size_per_gpu=${MICRO_BATCH_SIZE_PER_GPU} \ + data.messages_key=messages \ + data.ignore_input_ids_mismatch=True \ + optim.lr=${LR} \ + engine=fsdp \ + engine.ulysses_sequence_parallel_size=${SP_SIZE} \ + model.path="${MODEL_PATH}" \ + model.use_remove_padding=true \ + trainer.default_local_dir="${save_path}" \ + trainer.project_name="${PROJECT_NAME}" \ + trainer.experiment_name="${EXPERIMENT_NAME}" \ + trainer.logger='["console","wandb"]' \ + trainer.total_epochs=${TOTAL_EPOCHS} \ + "${extra_args[@]}" "$@" diff --git a/verl/examples/sft/gsm8k/run_qwen_megatron_fsdp.sh b/verl/examples/sft/gsm8k/run_qwen_megatron_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..647aacffe211d038c67f6c0dc8db6afba05e4421 --- /dev/null +++ b/verl/examples/sft/gsm8k/run_qwen_megatron_fsdp.sh @@ -0,0 +1,80 @@ +#!/usr/bin/env bash +set -xeuo pipefail + +########################### Quick Config ########################### + +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen2.5-Math-7B} +SAVE_PATH=${SAVE_PATH:-/root/checkpoints/Qwen2.5-Math-7B} +TRAIN_FILES=${TRAIN_FILES:-$HOME/data/gsm8k_sft/train.parquet} +VAL_FILES=${VAL_FILES:-$HOME/data/gsm8k_sft/test.parquet} + +NPROC=${NPROC:-8} +TP=${TP:-4} +PP=${PP:-1} +EP=${EP:-1} + +export CUDA_DEVICE_MAX_CONNECTIONS=1 +export HYDRA_FULL_ERROR=1 +unset ROCR_VISIBLE_DEVICES + +########################### Parameter Arrays ########################### + +DATA=( + "data.train_files=${TRAIN_FILES}" + "data.val_files=${VAL_FILES}" + data.messages_key=messages + data.train_batch_size=8 + data.use_dynamic_bsz=True + data.max_token_len_per_gpu=1024 + data.pad_mode=no_padding + data.truncation=error +) + +MODEL=( + model=hf_model + "model.path=${MODEL_PATH}" + model.trust_remote_code=True + model.use_remove_padding=true +) + +OPTIM=( + optim=megatron + optim.lr=1e-5 + optim.lr_warmup_steps_ratio=0.2 + optim.weight_decay=0.1 + "optim.betas=[0.9,0.95]" + optim.clip_grad=1.0 + optim.lr_warmup_init=0 + optim.lr_decay_style=cosine + optim.min_lr=1e-6 +) + +ENGINE=( + engine=megatron + engine.tensor_model_parallel_size=${TP} + engine.pipeline_model_parallel_size=${PP} + engine.expert_model_parallel_size=${EP} + engine.use_mbridge=True + engine.vanilla_mbridge=False + engine.use_megatron_fsdp=True + +engine.override_transformer_config.gradient_accumulation_fusion=False +) + +TRAINER=( + "trainer.default_local_dir=${SAVE_PATH}" + trainer.project_name=gsm8k-sft + trainer.experiment_name=SFT-qwen2.5-7b-mfsdp + trainer.logger='["console","wandb","file"]' + trainer.total_epochs=4 +) + +########################### Launch ########################### + +torchrun --standalone --nnodes=1 --nproc_per_node=$NPROC \ + -m verl.trainer.sft_trainer \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${OPTIM[@]}" \ + "${ENGINE[@]}" \ + "${TRAINER[@]}" \ + "$@" diff --git a/verl/examples/sft/gsm8k/run_seed_oss_36b_fsdp.sh b/verl/examples/sft/gsm8k/run_seed_oss_36b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..96f972780c412eabba199f6f34b7b9c4a7de2726 --- /dev/null +++ b/verl/examples/sft/gsm8k/run_seed_oss_36b_fsdp.sh @@ -0,0 +1,28 @@ +set -x + +if [ "$#" -lt 2 ]; then + echo "Usage: run_seed_oss_36b_fsdp.sh [other_configs...]" + exit 1 +fi + +nproc_per_node=$1 +save_path=$2 + +# Shift the arguments so $@ refers to the rest +shift 2 + +torchrun --standalone --nnodes=1 --nproc_per_node=$nproc_per_node \ + -m verl.trainer.sft_trainer \ + data.train_files=$HOME/data/gsm8k/train.parquet \ + data.val_files=$HOME/data/gsm8k/test.parquet \ + data.micro_batch_size_per_gpu=4 \ + optim.lr=1e-4 \ + engine=fsdp \ + engine.ulysses_sequence_parallel_size=2 \ + model.path=ByteDance-Seed/Seed-OSS-36B-Base \ + model.use_remove_padding=true \ + trainer.default_local_dir=$save_path \ + trainer.project_name=gsm8k-sft \ + trainer.experiment_name=gsm8k-sft-seed-oss-36b \ + trainer.logger=console \ + trainer.total_training_steps=1 $@ diff --git a/verl/examples/sft/multiturn/run_qwen2_5_0_5b_fsdp.sh b/verl/examples/sft/multiturn/run_qwen2_5_0_5b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..1809f6e049586b950586be574b5c1f5aef793e04 --- /dev/null +++ b/verl/examples/sft/multiturn/run_qwen2_5_0_5b_fsdp.sh @@ -0,0 +1,40 @@ +#!/usr/bin/env bash +# SFT | multiturn | FSDP engine | NVIDIA GPUs +# Toggle Ulysses sequence parallel via SP_SIZE env var. + +set -xeuo pipefail + +if [ "$#" -lt 2 ]; then + echo "Usage: run_qwen2_5_0_5b_fsdp.sh [other_configs...]" + echo " Env: SP_SIZE (default 1)" + exit 1 +fi + +nproc_per_node=$1 +save_path=$2 +shift 2 + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen2.5-0.5B-Instruct} +SP_SIZE=${SP_SIZE:-1} +MICRO_BATCH_SIZE_PER_GPU=${MICRO_BATCH_SIZE_PER_GPU:-4} +TOTAL_TRAINING_STEPS=${TOTAL_TRAINING_STEPS:-1} +PROJECT_NAME=${PROJECT_NAME:-multiturn-sft} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-multiturn-sft-qwen2_5_0_5b} +# ---- end user-adjustable ---- + +torchrun --nnodes=1 --nproc_per_node=${nproc_per_node} \ + -m verl.trainer.sft_trainer \ + data.train_files=$HOME/data/multiturn/train.parquet \ + data.val_files=$HOME/data/multiturn/test.parquet \ + data.messages_key=messages \ + data.micro_batch_size_per_gpu=${MICRO_BATCH_SIZE_PER_GPU} \ + model.path="${MODEL_PATH}" \ + model.use_remove_padding=true \ + engine=fsdp \ + engine.ulysses_sequence_parallel_size=${SP_SIZE} \ + trainer.default_local_dir="${save_path}" \ + trainer.project_name="${PROJECT_NAME}" \ + trainer.experiment_name="${EXPERIMENT_NAME}" \ + trainer.logger='["console","wandb"]' \ + trainer.total_training_steps=${TOTAL_TRAINING_STEPS} "$@" diff --git a/verl/examples/sft/vlm/run_qwen3_vl_2b_fsdp.sh b/verl/examples/sft/vlm/run_qwen3_vl_2b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..28c21ffa0491234966d22f08a3d6ab0fc4e2b853 --- /dev/null +++ b/verl/examples/sft/vlm/run_qwen3_vl_2b_fsdp.sh @@ -0,0 +1,100 @@ +#!/usr/bin/env bash +# python examples/data_preprocess/pokemon.py +set -xeuo pipefail + +HDFS_ROOT=${HDFS_ROOT:-$PWD} +DATA_ROOT=${DATA_ROOT:-$PWD} + +ENTRYPOINT=${ENTRYPOINT:-"-m verl.trainer.sft_trainer"} + +TRAIN_FILES=${HOME}/data/pokemon-gpt4o-captions/train.parquet + +backend=${BACKEND:-fsdp} + +project_name=verl_sft_test + +RESUME_MODE=auto +MODEL_ID=${HDFS_ROOT}/model/Qwen3-VL-2B-Instruct +# MODEL_ID=${HDFS_ROOT}/model/Qwen3-VL-30B-A3B-Instruct + +SP_SIZE=${SP_SIZE:-2} +FSDP_SIZE=${FSDP_SIZE:--1} +FSDP_STRATEGY=${FSDP_STRATEGY:-"fsdp2"} + +TP_SIZE=${TP_SIZE:-2} +PP_SIZE=${PP_SIZE:-2} +VPP_SIZE=${VPP_SIZE:-null} +CP_SIZE=${CP_SIZE:-1} + +PAD_MODE=${PAD_MODE:-no_padding} + +USE_REMOVE_PADDING=${USE_REMOVE_PADDING:-True} + +FSDP_ENGINE_CONFIG="\ + engine=${backend} \ + optim=${backend} \ + optim.lr=2e-5 \ + optim.lr_warmup_steps_ratio=0.01 \ + optim.weight_decay=0.1 \ + optim.betas="[0.9,0.95]" \ + optim.clip_grad=1.0 \ + optim.min_lr_ratio=0.1 \ + optim.warmup_style=cosine \ + engine.ulysses_sequence_parallel_size=${SP_SIZE} \ + engine.strategy=${FSDP_STRATEGY} \ + engine.fsdp_size=${FSDP_SIZE}" + + +MEGATRON_ENGINE_CONFIG="\ + engine=${backend} \ + optim=${backend} \ + optim.lr=2e-5 \ + optim.lr_warmup_steps_ratio=0.01 \ + optim.weight_decay=0.1 \ + optim.betas="[0.9,0.95]" \ + optim.clip_grad=1.0 \ + optim.lr_warmup_init=0 \ + optim.lr_decay_style=cosine \ + optim.min_lr=2e-6 \ + engine.tensor_model_parallel_size=${TP_SIZE} \ + engine.pipeline_model_parallel_size=${PP_SIZE} \ + engine.virtual_pipeline_model_parallel_size=${VPP_SIZE} \ + engine.context_parallel_size=${CP_SIZE} \ + engine.use_mbridge=True \ + engine.vanilla_mbridge=True" + +if [ "$backend" = "fsdp" ]; then + ENGINE_CONFIG="$FSDP_ENGINE_CONFIG" + echo "Using fsdp engine" + exp_name=pokemon-qwen3-2b-${backend}-${FSDP_STRATEGY}-sp${SP_SIZE}-fsdp-1202a1 +else + ENGINE_CONFIG="$MEGATRON_ENGINE_CONFIG" + echo "Using megatron engine" + exp_name=pokemon-qwen3-2b-${backend}-tp${TP_SIZE}-pp${PP_SIZE}-vpp${VPP_SIZE}-cp${CP_SIZE}-megatron-1202a1 +fi + +CKPT_HOME=${CKPT_HOME:-$HOME/open_verl/sft/${project_name}/${exp_name}} +mkdir -p "${CKPT_HOME}" + +torchrun --standalone --nnodes=1 --nproc-per-node=${NUM_TRAINERS:-8} \ + ${ENTRYPOINT} \ + data.train_files="${TRAIN_FILES}" \ + data.train_batch_size=96 \ + data.max_length=2048 \ + data.pad_mode=${PAD_MODE} \ + data.truncation=error \ + data.use_dynamic_bsz=True \ + data.max_token_len_per_gpu=65536 \ + model.path=$MODEL_ID \ + model.use_remove_padding=${USE_REMOVE_PADDING} \ + ${ENGINE_CONFIG} \ + trainer.test_freq=-1 \ + trainer.save_freq=4000 \ + trainer.logger=['console','wandb'] \ + trainer.project_name="${project_name}" \ + trainer.experiment_name="${exp_name}" \ + trainer.total_epochs=10 \ + trainer.default_local_dir="${CKPT_HOME}" \ + trainer.resume_mode=${RESUME_MODE} \ + trainer.max_ckpt_to_keep=5 \ + checkpoint.save_contents=[model,optimizer,extra] \ No newline at end of file diff --git a/verl/examples/tuning/README.md b/verl/examples/tuning/README.md new file mode 100644 index 0000000000000000000000000000000000000000..cfdd13b3ce7f0811717798ee96ccf25e2fa13904 --- /dev/null +++ b/verl/examples/tuning/README.md @@ -0,0 +1,37 @@ +# Tuning Examples + +Examples that focus on *tuning-related* features rather than on a specific RL algorithm. Everything in here trains with `verl.trainer.main_ppo` and the current Hydra API. + +## Subdirectories + +### `lora/` — LoRA fine-tuning + +Canonical LoRA GRPO scripts (training only adapters, rollout still serves the adapter via `load_format=safetensors + layered_summon`). + +| Script | Model | Infer | Train | Notes | +|------------------------------------------------|-------------------|-------|----------|-----------------------------| +| `run_qwen3_8b_fsdp.sh` | Qwen3-8B | vLLM | FSDP | text, GSM8K | +| `run_qwen3_8b_from_adapter_fsdp.sh` | Qwen3-8B | vLLM | FSDP | start from existing adapter | +| `run_qwen3_8b_merge_fsdp.sh` | Qwen3-8B | vLLM | FSDP | merge adapter into base | +| `run_qwen2_5_vl_7b_fsdp.sh` | Qwen2.5-VL-7B | vLLM | FSDP | vision, Geo3K | +| `run_qwen3_30b_a3b_megatron.sh` | Qwen3-30B-A3B | vLLM | Megatron | MoE | + +Key flags: +- `actor_rollout_ref.model.lora_rank`, `actor_rollout_ref.model.lora_alpha` +- `actor_rollout_ref.rollout.load_format=safetensors` +- `actor_rollout_ref.rollout.layered_summon=True` + +### `scaling/` — Large-model scale demos + +Single/multi-node tuning recipes for large dense models; geared to practitioners trying to fit and run these models out of the box with GRPO + GSM8K/MATH. + +| Script | Model | Infer | Train | Hardware | +|-------------------------------------------|-----------------|-------|----------|--------------------------| +| `run_qwen2_5_32b_megatron.sh` | Qwen2.5-32B | vLLM | Megatron | 1×8 GPUs (TP=8) | +| `run_qwen2_5_72b_fsdp.sh` | Qwen2.5-72B | vLLM | FSDP | 4×8 GPUs (TP=16, offload)| + +## Conventions + +- All scripts expose `MODEL_PATH`, `NNODES`, `NGPUS_PER_NODE`, batch sizes, learning rates, `ROLLOUT_TP`, `ROLLOUT_N`, etc. via `VAR=${VAR:-default}`. +- Dynamic batch size and `trainer.balance_batch=True` are enabled by default. +- No deprecated knobs (`ppo_micro_batch_size`, `data.val_batch_size`, top-level `reward_model.*`, `actor.ulysses_sequence_parallel_size`, `ppo_megatron_trainer.yaml`). diff --git a/verl/examples/tuning/lora/run_qwen2_5_vl_7b_fsdp.sh b/verl/examples/tuning/lora/run_qwen2_5_vl_7b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..4541bd94257c4338ed53843bdf233d3151da4edd --- /dev/null +++ b/verl/examples/tuning/lora/run_qwen2_5_vl_7b_fsdp.sh @@ -0,0 +1,116 @@ +#!/usr/bin/env bash +# GRPO + LoRA | vision | vLLM rollout | FSDP training | NVIDIA GPUs +# LoRA freezes the vision tower by excluding `.*visual.*` modules. + +set -xeuo pipefail + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen2.5-VL-7B-Instruct} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +train_batch_size=${TRAIN_BATCH_SIZE:-512} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-128} +max_prompt_length=${MAX_PROMPT_LENGTH:-1024} +max_response_length=${MAX_RESPONSE_LENGTH:-2048} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-24576} + +actor_lr=${ACTOR_LR:-3e-6} +kl_loss_coef=${KL_LOSS_COEF:-0.01} +entropy_coeff=${ENTROPY_COEFF:-0} + +lora_rank=${LORA_RANK:-64} +lora_alpha=${LORA_ALPHA:-32} + +rollout_tp=${ROLLOUT_TP:-2} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.6} +rollout_n=${ROLLOUT_N:-5} + +total_epochs=${TOTAL_EPOCHS:-15} +save_freq=${SAVE_FREQ:-20} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_grpo_lora_geo3k} +experiment_name=${EXPERIMENT_NAME:-qwen2_5_vl_7b_lora_vllm_fsdp} +# ---- end user-adjustable ---- +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files=$HOME/data/geo3k/train.parquet + data.val_files=$HOME/data/geo3k/test.parquet + data.image_key=images + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True + actor_rollout_ref.model.lora_rank=${lora_rank} + actor_rollout_ref.model.lora_alpha=${lora_alpha} + actor_rollout_ref.model.target_modules=all-linear + actor_rollout_ref.model.exclude_modules='.*visual.*' +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.fsdp_config.param_offload=False + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.enable_chunked_prefill=False + actor_rollout_ref.rollout.enforce_eager=False + actor_rollout_ref.rollout.free_cache_engine=False + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.fsdp_config.param_offload=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/tuning/lora/run_qwen3_30b_a3b_megatron.sh b/verl/examples/tuning/lora/run_qwen3_30b_a3b_megatron.sh new file mode 100644 index 0000000000000000000000000000000000000000..903c748fcb4c26550e98891b0eb3fb01028ac969 --- /dev/null +++ b/verl/examples/tuning/lora/run_qwen3_30b_a3b_megatron.sh @@ -0,0 +1,138 @@ +#!/usr/bin/env bash +# GRPO + MoE-LoRA | MoE text | vLLM rollout | Megatron training | NVIDIA GPUs +# Requires Megatron-Bridge > 0.2.0 for proper MoE LoRA support. + +set -xeuo pipefail + +export CUDA_DEVICE_MAX_CONNECTIONS=1 + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-30B-A3B-Instruct-2507} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +train_batch_size=${TRAIN_BATCH_SIZE:-128} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-16} +max_prompt_length=${MAX_PROMPT_LENGTH:-1024} +max_response_length=${MAX_RESPONSE_LENGTH:-1024} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-$((max_prompt_length + max_response_length))} + +actor_lr=${ACTOR_LR:-3e-6} +kl_loss_coef=${KL_LOSS_COEF:-0.001} +entropy_coeff=${ENTROPY_COEFF:-0} + +lora_rank=${LORA_RANK:-32} +lora_alpha=${LORA_ALPHA:-64} + +actor_tp=${ACTOR_TP:-2} +actor_pp=${ACTOR_PP:-2} +actor_ep=${ACTOR_EP:-4} +actor_etp=${ACTOR_ETP:-1} +actor_cp=${ACTOR_CP:-2} +all_offload=${ALL_OFFLOAD:-True} + +rollout_tp=${ROLLOUT_TP:-8} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.25} +rollout_n=${ROLLOUT_N:-4} + +total_epochs=${TOTAL_EPOCHS:-15} +save_freq=${SAVE_FREQ:-20} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_grpo_lora_gsm8k} +experiment_name=${EXPERIMENT_NAME:-qwen3_30b_a3b_lora_vllm_megatron} +# ---- end user-adjustable ---- +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files=$HOME/data/gsm8k/train.parquet + data.val_files=$HOME/data/gsm8k/test.parquet + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_fused_kernels=True + actor_rollout_ref.model.lora.rank=${lora_rank} + actor_rollout_ref.model.lora.alpha=${lora_alpha} + actor_rollout_ref.model.lora.lora_A_init_method=kaiming +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.megatron.use_mbridge=True + actor_rollout_ref.actor.megatron.vanilla_mbridge=False + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${actor_pp} + actor_rollout_ref.actor.megatron.expert_model_parallel_size=${actor_ep} + actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${actor_etp} + actor_rollout_ref.actor.megatron.context_parallel_size=${actor_cp} + actor_rollout_ref.actor.megatron.param_offload=${all_offload} + actor_rollout_ref.actor.megatron.optimizer_offload=${all_offload} + actor_rollout_ref.actor.megatron.grad_offload=${all_offload} + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.enforce_eager=True + actor_rollout_ref.rollout.free_cache_engine=True + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${actor_pp} + actor_rollout_ref.ref.megatron.expert_model_parallel_size=${actor_ep} + actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=${actor_etp} + actor_rollout_ref.ref.megatron.context_parallel_size=${actor_cp} + actor_rollout_ref.ref.megatron.param_offload=${all_offload} +) + +TRAINER=( + trainer.balance_batch=True + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( + model_engine=megatron +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/tuning/lora/run_qwen3_8b_from_adapter_fsdp.sh b/verl/examples/tuning/lora/run_qwen3_8b_from_adapter_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..ea189c8cb19dec2e12504c048c5881b9482b788e --- /dev/null +++ b/verl/examples/tuning/lora/run_qwen3_8b_from_adapter_fsdp.sh @@ -0,0 +1,110 @@ +#!/usr/bin/env bash +# GRPO + LoRA resume from adapter | text | vLLM rollout | FSDP training | NVIDIA GPUs +# Loads a previously saved LoRA adapter (LORA_ADAPTER_PATH) and continues GRPO on it. + +set -xeuo pipefail + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-8B} +LORA_ADAPTER_PATH=${LORA_ADAPTER_PATH:?must set LORA_ADAPTER_PATH} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +train_batch_size=${TRAIN_BATCH_SIZE:-1024} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-256} +max_prompt_length=${MAX_PROMPT_LENGTH:-512} +max_response_length=${MAX_RESPONSE_LENGTH:-1024} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-12288} + +actor_lr=${ACTOR_LR:-3e-6} +kl_loss_coef=${KL_LOSS_COEF:-0.001} +entropy_coeff=${ENTROPY_COEFF:-0} + +rollout_tp=${ROLLOUT_TP:-2} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.6} +rollout_n=${ROLLOUT_N:-5} + +total_epochs=${TOTAL_EPOCHS:-15} +save_freq=${SAVE_FREQ:-20} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_grpo_lora_gsm8k} +experiment_name=${EXPERIMENT_NAME:-qwen3_8b_lora_from_adapter_vllm_fsdp} +# ---- end user-adjustable ---- +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files=$HOME/data/gsm8k/train.parquet + data.val_files=$HOME/data/gsm8k/test.parquet + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_shm=True + actor_rollout_ref.model.lora_adapter_path="$LORA_ADAPTER_PATH" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.fsdp_config.param_offload=False + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.load_format=safetensors + actor_rollout_ref.rollout.layered_summon=True + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.fsdp_config.param_offload=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/tuning/lora/run_qwen3_8b_fsdp.sh b/verl/examples/tuning/lora/run_qwen3_8b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..2f27cac9cb26d88d9a837fba2812b7eb23bc92a2 --- /dev/null +++ b/verl/examples/tuning/lora/run_qwen3_8b_fsdp.sh @@ -0,0 +1,113 @@ +#!/usr/bin/env bash +# GRPO + LoRA | text | vLLM rollout | FSDP training | NVIDIA GPUs +# Canonical LoRA fine-tuning baseline for dense text LLMs (GSM8K). + +set -xeuo pipefail + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-8B} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-2} + +train_batch_size=${TRAIN_BATCH_SIZE:-64} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-32} +max_prompt_length=${MAX_PROMPT_LENGTH:-512} +max_response_length=${MAX_RESPONSE_LENGTH:-1024} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-12288} + +actor_lr=${ACTOR_LR:-3e-6} +kl_loss_coef=${KL_LOSS_COEF:-0.001} +entropy_coeff=${ENTROPY_COEFF:-0} + +lora_rank=${LORA_RANK:-64} +lora_alpha=${LORA_ALPHA:-32} + +rollout_tp=${ROLLOUT_TP:-2} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.6} +rollout_n=${ROLLOUT_N:-5} + +total_epochs=${TOTAL_EPOCHS:-15} +save_freq=${SAVE_FREQ:-20} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_grpo_lora_gsm8k} +experiment_name=${EXPERIMENT_NAME:-qwen3_8b_lora_vllm_fsdp} +# ---- end user-adjustable ---- +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files=$HOME/data/gsm8k/train.parquet + data.val_files=$HOME/data/gsm8k/test.parquet + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.lora_rank=${lora_rank} + actor_rollout_ref.model.lora_alpha=${lora_alpha} + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.fsdp_config.param_offload=False + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.load_format=safetensors + actor_rollout_ref.rollout.layered_summon=True + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.fsdp_config.param_offload=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.val_before_train=False + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/tuning/lora/run_qwen3_8b_merge_fsdp.sh b/verl/examples/tuning/lora/run_qwen3_8b_merge_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..8ada1743ba145d4c00bfb1e05d25d3e88efbe1c7 --- /dev/null +++ b/verl/examples/tuning/lora/run_qwen3_8b_merge_fsdp.sh @@ -0,0 +1,118 @@ +#!/usr/bin/env bash +# GRPO + LoRA with merge | text | vLLM rollout | FSDP2 training | NVIDIA GPUs +# Trains LoRA adapters and merges them into the base model during rollout. + +set -xeuo pipefail + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-8B} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +train_batch_size=${TRAIN_BATCH_SIZE:-128} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-64} +max_prompt_length=${MAX_PROMPT_LENGTH:-1024} +max_response_length=${MAX_RESPONSE_LENGTH:-1024} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-24576} + +actor_lr=${ACTOR_LR:-1.0e-05} +kl_loss_coef=${KL_LOSS_COEF:-0.001} +entropy_coeff=${ENTROPY_COEFF:-0} + +lora_rank=${LORA_RANK:-32} +lora_alpha=${LORA_ALPHA:-64} + +rollout_tp=${ROLLOUT_TP:-2} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.6} +rollout_n=${ROLLOUT_N:-8} + +total_epochs=${TOTAL_EPOCHS:-1} +save_freq=${SAVE_FREQ:-20} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_grpo_lora_gsm8k} +experiment_name=${EXPERIMENT_NAME:-qwen3_8b_lora_merge_vllm_fsdp} +# ---- end user-adjustable ---- +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files=$HOME/data/gsm8k/train.parquet + data.val_files=$HOME/data/gsm8k/test.parquet + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True + +actor_rollout_ref.model.lora.merge=True + actor_rollout_ref.model.lora_rank=${lora_rank} + actor_rollout_ref.model.lora_alpha=${lora_alpha} +) + +ACTOR=( + actor_rollout_ref.actor.strategy=fsdp2 + actor_rollout_ref.actor.fsdp_config.model_dtype=bf16 + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.fsdp_config.param_offload=False + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.load_format=safetensors + actor_rollout_ref.rollout.layered_summon=True + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.strategy=fsdp2 + actor_rollout_ref.ref.fsdp_config.model_dtype=bf16 + actor_rollout_ref.ref.fsdp_config.param_offload=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.val_before_train=True + trainer.logger='["console","tensorboard","file"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/tuning/scaling/run_qwen2_5_32b_megatron.sh b/verl/examples/tuning/scaling/run_qwen2_5_32b_megatron.sh new file mode 100644 index 0000000000000000000000000000000000000000..10e2380aa0325a88592100cb14398a060b8ca919 --- /dev/null +++ b/verl/examples/tuning/scaling/run_qwen2_5_32b_megatron.sh @@ -0,0 +1,122 @@ +#!/usr/bin/env bash +# GRPO | text | vLLM rollout | Megatron training | NVIDIA GPUs (8xH20/H100/H200) +# Scaling demo: single-node 32B dense model, Megatron TP=8. + +set -xeuo pipefail + +export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:256 + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen2.5-32B} +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +train_batch_size=${TRAIN_BATCH_SIZE:-512} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-256} +max_prompt_length=${MAX_PROMPT_LENGTH:-2048} +max_response_length=${MAX_RESPONSE_LENGTH:-6144} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-32768} + +actor_lr=${ACTOR_LR:-1e-6} +kl_loss_coef=${KL_LOSS_COEF:-0.001} +entropy_coeff=${ENTROPY_COEFF:-0} + +actor_tp=${ACTOR_TP:-8} +actor_pp=${ACTOR_PP:-1} +rollout_tp=${ROLLOUT_TP:-8} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.6} +rollout_n=${ROLLOUT_N:-5} + +total_epochs=${TOTAL_EPOCHS:-15} +save_freq=${SAVE_FREQ:-20} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_grpo_scaling_gsm8k_math} +experiment_name=${EXPERIMENT_NAME:-qwen2_5_32b_vllm_megatron} +# ---- end user-adjustable ---- + +gsm8k_train=$HOME/data/gsm8k/train.parquet +gsm8k_test=$HOME/data/gsm8k/test.parquet +math_train=$HOME/data/math/train.parquet +math_test=$HOME/data/math/test.parquet + +train_files="['$gsm8k_train', '$math_train']" +val_files="['$gsm8k_test', '$math_test']" +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files="$train_files" + data.val_files="$val_files" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True +) + +ACTOR=( + actor_rollout_ref.actor.strategy=megatron + actor_rollout_ref.actor.model_engine=megatron + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${actor_pp} + actor_rollout_ref.actor.megatron.param_offload=True + actor_rollout_ref.actor.megatron.grad_offload=True + actor_rollout_ref.actor.megatron.optimizer_offload=True +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.megatron.param_offload=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/tuning/scaling/run_qwen2_5_72b_fsdp.sh b/verl/examples/tuning/scaling/run_qwen2_5_72b_fsdp.sh new file mode 100644 index 0000000000000000000000000000000000000000..75d6a79c075cb5b34166640a4306320b0ef1de89 --- /dev/null +++ b/verl/examples/tuning/scaling/run_qwen2_5_72b_fsdp.sh @@ -0,0 +1,114 @@ +#!/usr/bin/env bash +# GRPO | text | vLLM rollout | FSDP training | NVIDIA GPUs (32xH800/H200) +# Scaling demo: 72B dense model, FSDP with full offload, vLLM TP=16. + +set -xeuo pipefail + +# ---- user-adjustable ---- +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen2.5-72B-Instruct} +NNODES=${NNODES:-4} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +train_batch_size=${TRAIN_BATCH_SIZE:-1024} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-256} +max_prompt_length=${MAX_PROMPT_LENGTH:-1024} +max_response_length=${MAX_RESPONSE_LENGTH:-2048} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-24576} + +actor_lr=${ACTOR_LR:-1e-6} +kl_loss_coef=${KL_LOSS_COEF:-0.001} +entropy_coeff=${ENTROPY_COEFF:-0} + +rollout_tp=${ROLLOUT_TP:-16} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.6} +rollout_n=${ROLLOUT_N:-5} + +total_epochs=${TOTAL_EPOCHS:-1} +save_freq=${SAVE_FREQ:-20} +test_freq=${TEST_FREQ:-5} + +project_name=${PROJECT_NAME:-verl_grpo_scaling_gsm8k_math} +experiment_name=${EXPERIMENT_NAME:-qwen2_5_72b_vllm_fsdp} +# ---- end user-adjustable ---- + +gsm8k_train=$HOME/data/gsm8k/train.parquet +gsm8k_test=$HOME/data/gsm8k/test.parquet +math_train=$HOME/data/math/train.parquet +math_test=$HOME/data/math/test.parquet + +train_files="['$gsm8k_train', '$math_train']" +val_files="['$gsm8k_test', '$math_test']" +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files="$train_files" + data.val_files="$val_files" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.fsdp_config.param_offload=True + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.fsdp_config.param_offload=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.critic_warmup=0 + trainer.logger='["console","wandb"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} +) + +EXTRA=( +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/examples/tutorial/README.md b/verl/examples/tutorial/README.md new file mode 100644 index 0000000000000000000000000000000000000000..cd046fb8288fb67a8149e6a5e9354b985def237c --- /dev/null +++ b/verl/examples/tutorial/README.md @@ -0,0 +1,15 @@ +# verl Tutorials & Launchers + +Learning-oriented content that is not tied to a specific algorithm. + +## Contents + +| Subdir | What it is | +|----------------------------|----------------------------------------------------------------------------| +| `agent_loop_get_started/` | Notebook walkthrough of the agent-loop API (`agent_loop_tutorial.ipynb`). | +| `ray/` | Ray API crash-course notebook (`tutorial.ipynb`). | +| `slurm/` | SLURM job template for running Ray-on-SLURM (`ray_on_slurm.slurm`). | +| `skypilot/` | SkyPilot task specs for cloud / Kubernetes (`verl-ppo.yaml`, `verl-grpo.yaml`, `verl-multiturn-tools.yaml`). | + +Use these as starting points; the runnable training scripts live under the +trainer directories (`grpo_trainer/`, `ppo_trainer/`, `sft/`, ...). diff --git a/verl/examples/tutorial/agent_loop_get_started/agent_loop_tutorial.ipynb b/verl/examples/tutorial/agent_loop_get_started/agent_loop_tutorial.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..406446d37405df77961233d3939d658311ee7962 --- /dev/null +++ b/verl/examples/tutorial/agent_loop_get_started/agent_loop_tutorial.ipynb @@ -0,0 +1,929 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Train ReAct agent with code sandbox\n", + "\n", + "In this tutorial, we will demonstrate how to train a [ReAct](https://arxiv.org/abs/2210.03629) agent to solve math problem with code sandbox.\n", + "\n", + "The agent works as follows:\n", + "1. Given a math problem, the agent first query LLM to generate response and tool calls, which are python code to be executed in sandbox.\n", + "2. If there is a tool call, the agent execute the python code in code sandbox.\n", + "3. After code execution, the agent get the result from sandbox and append to chat history.\n", + "4. The agent query LLM again until no tool call or max context length reached.\n", + "\n", + "\n", + "
\n", + " \"ReAct\"\n", + "
\n", + " source: LangGraph\n", + "
\n", + "
" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Prerequisite\n", + "\n", + "To run the examples in this notebook, you need to install the verl package first.\n", + "```bash\n", + "git clone https://github.com/verl-project/verl\n", + "cd verl\n", + "pip install -e .\n", + "```" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "2025-10-16 23:20:11,956\tINFO worker.py:2004 -- Started a local Ray instance. View the dashboard at \u001b[1m\u001b[32mhttp://127.0.0.1:8265 \u001b[39m\u001b[22m\n", + "/usr/local/lib/python3.12/dist-packages/ray/_private/worker.py:2052: FutureWarning: Tip: In future versions of Ray, Ray will no longer override accelerator visible devices env var if num_gpus=0 or num_gpus=None (default). To enable this behavior and turn off this error message, set RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO=0\n", + " warnings.warn(\n" + ] + } + ], + "source": [ + "import asyncio\n", + "import sys\n", + "import tempfile\n", + "import os\n", + "import socket\n", + "import json\n", + "\n", + "import requests\n", + "import ray\n", + "import fastapi\n", + "import uvicorn\n", + "from starlette.requests import Request\n", + "from starlette.responses import JSONResponse\n", + "from pprint import pprint\n", + "\n", + "import verl\n", + "\n", + "ray.init()\n", + "verl_config_dir = os.path.join(os.path.dirname(verl.__file__), \"trainer/config\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For demo purpose, we will use Qwen/Qwen3-1.7B as the LLM. First, let's download required model and dataset used in this tutorial." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import pyarrow.parquet as pq\n", + "from huggingface_hub import snapshot_download\n", + "\n", + "snapshot_download(\n", + " repo_id=\"verl-team/lighteval-MATH-preprocessed\",\n", + " repo_type=\"dataset\",\n", + " local_dir=os.path.expanduser(\"~/verl-team/lighteval-MATH-preprocessed\"),\n", + ")\n", + "snapshot_download(\n", + " repo_id=\"Qwen/Qwen3-1.7B\",\n", + " repo_type=\"model\",\n", + " local_dir=os.path.expanduser(\"~/Qwen/Qwen3-1.7B\"),\n", + ")\n", + "\n", + "model_path = os.path.expanduser(\"~/Qwen/Qwen3-1.7B\")\n", + "train_file = os.path.expanduser(\"~/verl-team/lighteval-MATH-preprocessed/train.parquet\")\n", + "test_file = os.path.expanduser(\"~/verl-team/lighteval-MATH-preprocessed/test.parquet\")\n", + "\n", + "test = pq.read_table(test_file)\n", + "test_file = os.path.expanduser(\"~/verl-team/lighteval-MATH-preprocessed/test_100.parquet\")\n", + "pq.write_table(test[:100], test_file)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "verl support both vllm and sglang rollout server for high performance inference. This tutorial has been tested on both vllm and sglang, you can choose either of them to run the tutorial." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [], + "source": [ + "rollout_name = \"???\" # vllm or sglang" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Basic tool call\n", + "For beginning, let's see how we can do basic tool call in verl with example from [Transformer tool use](https://huggingface.co/docs/transformers/main/chat_extras#tool-use). To use tool in verl, we need to define a tool class that inherits from `BaseTool`, and implement the following methods:\n", + "- `get_openai_tool_schema`: return the schema of the tool in `OpenAIFunctionToolSchema` format.\n", + "- `execute`: execute the tool with the given parameters, and return the result in `ToolResponse` format." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "{\n", + " \"type\": \"function\",\n", + " \"function\": {\n", + " \"name\": \"get_current_temperature\",\n", + " \"description\": \"Get current temperature at a location.\",\n", + " \"parameters\": {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"location\": {\n", + " \"type\": \"string\",\n", + " \"description\": \"The location to get the temperature for, in the format \\\"City, State, Country\\\".\"\n", + " },\n", + " \"unit\": {\n", + " \"type\": \"string\",\n", + " \"description\": \"The unit to return the temperature in. Defaults to \\\"celsius\\\".\",\n", + " \"enum\": [\n", + " \"celsius\",\n", + " \"fahrenheit\"\n", + " ]\n", + " }\n", + " },\n", + " \"required\": [\n", + " \"location\"\n", + " ]\n", + " }\n", + " }\n", + "}\n" + ] + } + ], + "source": [ + "from transformers.utils import get_json_schema\n", + "from verl.tools.base_tool import BaseTool, OpenAIFunctionToolSchema, ToolResponse\n", + "\n", + "\n", + "class WeatherTool(BaseTool):\n", + " def get_current_temperature(self, location: str, unit: str = \"celsius\"):\n", + " \"\"\"Get current temperature at a location.\n", + "\n", + " Args:\n", + " location: The location to get the temperature for, in the format \"City, State, Country\".\n", + " unit: The unit to return the temperature in. Defaults to \"celsius\". (choices: [\"celsius\", \"fahrenheit\"])\n", + "\n", + " Returns:\n", + " the temperature, the location, and the unit in a dict\n", + " \"\"\"\n", + " return {\n", + " \"temperature\": 26.1,\n", + " \"location\": location,\n", + " \"unit\": unit,\n", + " }\n", + "\n", + " def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema:\n", + " schema = get_json_schema(self.get_current_temperature)\n", + " return OpenAIFunctionToolSchema(**schema)\n", + "\n", + " async def execute(self, instance_id: str, parameters: dict, **kwargs) -> tuple[ToolResponse, float, dict]:\n", + " try:\n", + " result = self.get_current_temperature(**parameters)\n", + " return ToolResponse(text=json.dumps(result)), 0, {}\n", + " except Exception as e:\n", + " return ToolResponse(text=str(e)), 0, {}\n", + "\n", + "\n", + "weather_tool = WeatherTool(config={}, tool_schema=None)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Next, let's launch a standalone rollout server without hybrid engine (which is more heavy to start) to test the basic tool call." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from hydra import compose, initialize_config_dir\n", + "from verl.workers.rollout.replica import get_rollout_replica_class\n", + "\n", + "with initialize_config_dir(config_dir=verl_config_dir):\n", + " config = compose(\n", + " config_name=\"ppo_trainer\",\n", + " overrides=[\n", + " \"actor_rollout_ref.rollout.name=\" + rollout_name,\n", + " \"actor_rollout_ref.rollout.mode=async\",\n", + " \"actor_rollout_ref.rollout.tensor_model_parallel_size=1\",\n", + " \"actor_rollout_ref.model.path=\" + model_path,\n", + " \"actor_rollout_ref.rollout.response_length=4096\",\n", + " \"actor_rollout_ref.rollout.skip_tokenizer_init=False\",\n", + " \"+actor_rollout_ref.rollout.engine_kwargs.vllm.enable_auto_tool_choice=True\",\n", + " \"+actor_rollout_ref.rollout.engine_kwargs.vllm.tool_call_parser=hermes\",\n", + " \"+actor_rollout_ref.rollout.engine_kwargs.sglang.tool_call_parser=qwen25\",\n", + " ],\n", + " )\n", + "\n", + "rollout_server_class = get_rollout_replica_class(config.actor_rollout_ref.rollout.name)\n", + "rollout_server = rollout_server_class(\n", + " replica_rank=0,\n", + " config=config.actor_rollout_ref.rollout,\n", + " model_config=config.actor_rollout_ref.model,\n", + ")\n", + "\n", + "await rollout_server.init_standalone()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Then, we can query LLM with openai client. Note that we need to pass the tool schema to server to guide LLM generating tool calls. We can see that the LLM correctly generates a tool call to get the temperature in Paris." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[{'content': \"Hey, what's the temperature in Paris right now?\", 'role': 'user'},\n", + " {'role': 'assistant',\n", + " 'tool_calls': [{'function': {'arguments': '{\"location\": \"Paris, France\"}',\n", + " 'name': 'get_current_temperature'},\n", + " 'id': 'call_b10bdde504a0411690e96b55',\n", + " 'index': -1,\n", + " 'type': 'function'}]}]\n" + ] + } + ], + "source": [ + "from openai import AsyncOpenAI\n", + "\n", + "client = AsyncOpenAI(\n", + " api_key=\"dummy\",\n", + " base_url=f\"http://{rollout_server._server_address}/v1\",\n", + ")\n", + "\n", + "messages = [{\"role\": \"user\", \"content\": \"Hey, what's the temperature in Paris right now?\"}]\n", + "completion = await client.chat.completions.create(\n", + " model=config.actor_rollout_ref.model.path,\n", + " messages=messages,\n", + " tools=[weather_tool.tool_schema.model_dump(exclude_unset=True, exclude_none=True)],\n", + " extra_body={\n", + " \"chat_template_kwargs\": {\"enable_thinking\": False},\n", + " },\n", + ")\n", + "\n", + "message = completion.choices[0].message.model_dump(exclude_unset=True, exclude_none=True)\n", + "messages.append(message)\n", + "pprint(messages)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can execute the tool call with arguments generated by LLM and get the temperature in Paris." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "text='{\"temperature\": 26.1, \"location\": \"Paris, France\", \"unit\": \"celsius\"}' image=None video=None\n" + ] + } + ], + "source": [ + "args = json.loads(message[\"tool_calls\"][0][\"function\"][\"arguments\"])\n", + "tool_response, _, _ = await weather_tool.execute(\"\", args)\n", + "print(tool_response)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Then, we can add the tool response to chat history and query LLM again. With the tool response, LLM can generate a final response to the user." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[{'content': \"Hey, what's the temperature in Paris right now?\", 'role': 'user'},\n", + " {'role': 'assistant',\n", + " 'tool_calls': [{'function': {'arguments': '{\"location\": \"Paris, France\"}',\n", + " 'name': 'get_current_temperature'},\n", + " 'id': 'call_b10bdde504a0411690e96b55',\n", + " 'index': -1,\n", + " 'type': 'function'}]},\n", + " {'content': '{\"temperature\": 26.1, \"location\": \"Paris, France\", \"unit\": '\n", + " '\"celsius\"}',\n", + " 'role': 'tool'},\n", + " {'content': 'The current temperature in Paris is 26.1°C.',\n", + " 'role': 'assistant'}]\n" + ] + } + ], + "source": [ + "messages.append({\"role\": \"tool\", \"content\": tool_response.text})\n", + "completion = await client.chat.completions.create(\n", + " model=config.actor_rollout_ref.model.path,\n", + " messages=messages,\n", + " tools=[weather_tool.tool_schema.model_dump(exclude_unset=True, exclude_none=True)],\n", + " extra_body={\n", + " \"chat_template_kwargs\": {\"enable_thinking\": False},\n", + " },\n", + ")\n", + "\n", + "message = completion.choices[0].message.model_dump(exclude_unset=True, exclude_none=True)\n", + "messages.append(message)\n", + "pprint(messages)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Advanced tool call with code sandbox\n", + "\n", + "Now, let's see a more realistic example of tool call with code sandbox, which is widely used in real-world applications.\n", + "\n", + "### 2.1 Implement a naive code sandbox\n", + "\n", + "To execute python code snippet generated by LLM, we need a code sandbox environment. In this tutorial, we will implement a very naive code sandbox, which is\n", + "a FastAPI http server with `/run_code` endpoint. The server works as follows:\n", + "1. Receive a http request, write the python code snippet to a temp file.\n", + "2. Spawn a subprocess to execute the code, and get stdout and stderr of the subprocess.\n", + "3. Return the stdout and stderr of the subprocess as http response.\n", + "\n", + "> 🚨 **WARNING:** This naive code sandbox is for demonstration purpose only, do not use it in production. Please use docker/kata container for stronger isolation and security restriction." + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [], + "source": [ + "@ray.remote(num_cpus=1)\n", + "class Sandbox:\n", + " \"\"\"Sandbox to execute python code.\"\"\"\n", + "\n", + " def __init__(self):\n", + " self.address = ray._private.services.get_node_ip_address()\n", + " self.port = self._get_free_port()\n", + " asyncio.create_task(self._start_fastapi_server())\n", + "\n", + " async def code_execution(self, request: Request):\n", + " request_json = await request.json()\n", + " code = request_json[\"code\"]\n", + " # print(f\"execute code:\\n{code}\")\n", + "\n", + " _, temp_file = tempfile.mkstemp(suffix=\".py\", prefix=\"temp_code\", dir=None, text=True)\n", + " with open(temp_file, \"w\") as f:\n", + " f.write(code)\n", + "\n", + " try:\n", + " process = await asyncio.create_subprocess_exec(\n", + " sys.executable, temp_file, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE\n", + " )\n", + "\n", + " stdout, stderr = await process.communicate()\n", + "\n", + " response = {\n", + " \"status\": \"Success\" if process.returncode == 0 else \"Failed\",\n", + " \"run_result\": {\n", + " \"status\": \"Finished\",\n", + " \"stdout\": stdout.decode(),\n", + " \"stderr\": stderr.decode(),\n", + " \"return_code\": process.returncode,\n", + " },\n", + " }\n", + " return JSONResponse(content=response)\n", + " finally:\n", + " try:\n", + " os.unlink(temp_file)\n", + " except Exception:\n", + " pass\n", + "\n", + " def _get_free_port(self):\n", + " with socket.socket() as sock:\n", + " sock.bind((\"\", 0))\n", + " return sock.getsockname()[1]\n", + "\n", + " async def _start_fastapi_server(self):\n", + " app = fastapi.FastAPI()\n", + " app.router.add_api_route(\"/run_code\", self.code_execution, methods=[\"POST\"])\n", + "\n", + " config = uvicorn.Config(app, host=[\"::\", \"0.0.0.0\"], port=self.port, log_level=\"warning\")\n", + " server = uvicorn.Server(config)\n", + " await server.serve()\n", + "\n", + " async def get_server_address(self) -> str:\n", + " \"\"\"Get FastAPI server address.\"\"\"\n", + " return f\"{self.address}:{self.port}\"" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [], + "source": [ + "sandbox = Sandbox.remote()\n", + "sandbox_address = ray.get(sandbox.get_server_address.remote())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2.2 Define sandbox tool\n", + "\n", + "As shown in the previous section, we also defined a tool for the code sandbox. In the `execute` method, we send the code snippet to code sandbox by http request and get the output." + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "{\n", + " \"type\": \"function\",\n", + " \"function\": {\n", + " \"name\": \"code_interpreter\",\n", + " \"description\": \"Execute the code in the sandbox.\",\n", + " \"parameters\": {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"code\": {\n", + " \"type\": \"string\",\n", + " \"description\": \"The code to be executed.\"\n", + " }\n", + " },\n", + " \"required\": [\n", + " \"code\"\n", + " ]\n", + " }\n", + " }\n", + "}\n" + ] + } + ], + "source": [ + "import re\n", + "import aiohttp\n", + "\n", + "\n", + "class SandboxTool(BaseTool):\n", + " def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema):\n", + " super().__init__(config, tool_schema)\n", + " # Different model may use different code pattern, e.g. python, py, etc.\n", + " self.code_pattern = re.compile(r\"```py(.*?)```\", re.DOTALL)\n", + "\n", + " async def code_interpreter(self, code: str) -> str:\n", + " \"\"\"Execute the code in the sandbox.\n", + "\n", + " Args:\n", + " code: The code to be executed.\n", + "\n", + " Returns:\n", + " str: The output of the code execution.\n", + " \"\"\"\n", + " async with aiohttp.ClientSession() as session:\n", + " async with session.post(\n", + " self.config.get(\"sandbox_fusion_url\"),\n", + " json={\"code\": code},\n", + " ) as resp:\n", + " resp.raise_for_status()\n", + " result = await resp.json()\n", + " stdout, stderr = result[\"run_result\"][\"stdout\"], result[\"run_result\"][\"stderr\"]\n", + " return stdout + stderr\n", + "\n", + " def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema:\n", + " schema = get_json_schema(self.code_interpreter)\n", + " return OpenAIFunctionToolSchema(**schema)\n", + "\n", + " async def execute(self, instance_id: str, parameters: dict, **kwargs) -> tuple[str, float, dict]:\n", + " code = parameters[\"code\"]\n", + " matches = self.code_pattern.findall(code)\n", + " if matches:\n", + " code = matches[0].strip()\n", + "\n", + " # NOTE: Some script may not explicitly print result, we need to add a print statement to the end of the script.\n", + " # More better way is to SFT the model to make it print result by default, we skip SFT stage in this tutorial.\n", + " lines = code.split(\"\\n\")\n", + " for i, line in reversed(list(enumerate(lines))):\n", + " if line == \"\":\n", + " continue\n", + " if not lines[i].startswith(\"print\"):\n", + " lines[i] = f\"print({line})\"\n", + " break\n", + " code = \"\\n\".join(lines)\n", + "\n", + " result = await self.code_interpreter(code)\n", + " return ToolResponse(text=result), 0.0, {}\n", + "\n", + "\n", + "sandbox_tool = SandboxTool(config={\"sandbox_fusion_url\": f\"http://{sandbox_address}/run_code\"}, tool_schema=None)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "First, let's try to execute a valid code and check the response with stdout." + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "(ToolResponse(text='sqrt(3)\\n', image=None, video=None), 0.0, {})\n" + ] + } + ], + "source": [ + "code = \"\"\"```py\n", + "import sympy\n", + "\n", + "print(sympy.sqrt(3))\n", + "```\"\"\"\n", + "\n", + "print(await sandbox_tool.execute(instance_id=\"\", parameters={\"code\": code}))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Then, let's try to execute an invalid code and check the response with stderr. The error message is important to inform LLM to fix code in next generation." + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "(ToolResponse(text='Traceback (most recent call last):\\n File \"/tmp/temp_code3e2f638_.py\", line 2, in \\n print(sympy.sqrt(3))\\n ^^^^^\\nNameError: name \\'sympy\\' is not defined\\n', image=None, video=None), 0.0, {})\n" + ] + } + ], + "source": [ + "code_invalid = \"\"\"\n", + "print(sympy.sqrt(3))\n", + "\"\"\"\n", + "\n", + "print(await sandbox_tool.execute(instance_id=\"\", parameters={\"code\": code_invalid}))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2.3 Test sandbox tool\n", + "\n", + "Now, we can test sandbox tool with real math problem. In this tutorial, we will use the [DigitalLearningGmbH/MATH-lighteval](https://huggingface.co/datasets/DigitalLearningGmbH/MATH-lighteval) dataset, which consists of problems from mathematics competitions, including the AMC 10, AMC 12, AIME, and more." + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": {}, + "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "ebd09c8816b140a59a879e5a5e218950", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Generating train split: 0 examples [00:00, ? examples/s]" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "from datasets import load_dataset\n", + "\n", + "dataset = load_dataset(\"parquet\", data_files=test_file)[\"train\"]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For debug purpose, we can implement ReAct agent as a simple loop. For RL training, there are more subtle issue and corner case to deal with, we provide a built-in ReAct agent loop which will be discussed in next section." + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "No tool calls, finish_reason: stop\n" + ] + } + ], + "source": [ + "messages = dataset[\"prompt\"][0]\n", + "\n", + "while True:\n", + " # 1. Chat with the model\n", + " completion = await client.chat.completions.create(\n", + " model=config.actor_rollout_ref.model.path,\n", + " messages=messages,\n", + " tools=[sandbox_tool.tool_schema.model_dump(exclude_unset=True, exclude_none=True)],\n", + " extra_body={\n", + " \"chat_template_kwargs\": {\"enable_thinking\": False},\n", + " },\n", + " )\n", + "\n", + " message = completion.choices[0].message.model_dump(exclude_unset=True, exclude_none=True)\n", + " messages.append(message)\n", + "\n", + " # 2. Call tools\n", + " finish_reason = completion.choices[0].finish_reason\n", + " if finish_reason != \"tool_calls\":\n", + " print(f\"No tool calls, finish_reason: {finish_reason}\")\n", + " break\n", + "\n", + " try:\n", + " tool_calls = completion.choices[0].message.tool_calls[0]\n", + " args = json.loads(tool_calls.function.arguments)\n", + " result, _, _ = await sandbox_tool.execute(\"\", args)\n", + " except Exception as e:\n", + " print(f\"Error: {e}\")\n", + "\n", + " # 3. Add tool response to messages\n", + " messages.append(\n", + " {\n", + " \"role\": \"tool\",\n", + " \"content\": result.text,\n", + " }\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "[{'content': \"How many vertical asymptotes does the graph of $y=\\\\frac{2}{x^2+x-6}$ have? Let's think step by step and output the final answer within \\\\boxed{}.\",\n", + " 'role': 'user'},\n", + " {'content': \"To determine the number of vertical asymptotes for the function $ y = \\\\frac{2}{x^2 + x - 6} $, we need to find the values of $ x $ where the denominator equals zero, as these points are where the function is undefined and potentially where it has vertical asymptotes.\\n\\nThe denominator is $ x^2 + x - 6 $. To find the vertical asymptotes, we need to solve the equation:\\n\\n$$ x^2 + x - 6 = 0 $$\\n\\nThis is a quadratic equation, and we can solve it using the quadratic formula:\\n\\n$$ x = \\\\frac{-b \\\\pm \\\\sqrt{b^2 - 4ac}}{2a} $$\\n\\nwhere $ a = 1 $, $ b = 1 $, and $ c = -6 $. Let's solve this equation to find the values of $ x $ where the denominator is zero, which will give us the vertical asymptotes.\",\n", + " 'role': 'assistant',\n", + " 'tool_calls': [{'id': 'call_4d873672ff8445159e4e5e45',\n", + " 'function': {'arguments': '{\"code\": \"from sympy import symbols, solve\\\\nx = symbols(\\'x\\')\\\\nroots = solve(x**2 + x - 6, x)\\\\nroots\"}',\n", + " 'name': 'code_interpreter'},\n", + " 'type': 'function',\n", + " 'index': -1}]},\n", + " {'role': 'tool', 'content': '[-3, 2]\\n'},\n", + " {'content': 'The roots of the equation $ x^2 + x - 6 = 0 $ are $ x = -3 $ and $ x = 2 $. These are the values of $ x $ where the denominator is zero, which means the function $ y = \\\\frac{2}{x^2 + x - 6} $ is undefined at these points. \\n\\nSince the denominator is zero at these values, the function has vertical asymptotes at $ x = -3 $ and $ x = 2 $. Therefore, the graph of the function has two vertical asymptotes.\\n\\nThe final answer is $\\\\boxed{2}$.',\n", + " 'role': 'assistant'}]" + ] + }, + "execution_count": 16, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "messages" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can see that the ReAct agent properly query LLM, execute sandbox tool call, finally generate the answer." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. End-to-end training with tool agent loop\n", + "\n", + "After tool has been implemented and tested, we can do end-to-end RL training to tune the model to properly use the tool. To simplify agentic RL training, verl provide [Agent Loop](https://verl.readthedocs.io/en/latest/advance/agent_loop.html) abstraction, which allow user to define custom agent loop:\n", + "- Search agent\n", + "- Math agent\n", + "- SWE agent\n", + "- GUI agent\n", + "- ...\n", + "\n", + "For ease of use, verl provide two pre-defined agent loop:\n", + "- SingleTurnAgentLoop: single-turn conversation without tool calling\n", + "- ToolAgentLoop: multi-turn conversation with tool calling, interaction\n", + "\n", + "To use ToolAgentLoop, user only need to provide tools configuration in json/yaml file. In the configuration file, user should specify following fields for each tool:\n", + "- class_name: fully qualified class name of the tool used to dynamically load the custom tool class\n", + "- config: key-word arguments used to initialize the tool instance\n", + "\n", + "Let's dump our sandbox tool configuration to a json file:" + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "2025-10-16 23:07:16,868\tINFO worker.py:2004 -- Started a local Ray instance. View the dashboard at \u001b[1m\u001b[32mhttp://127.0.0.1:8265 \u001b[39m\u001b[22m\n" + ] + } + ], + "source": [ + "ray.shutdown()\n", + "\n", + "sandbox = Sandbox.remote()\n", + "sandbox_address = ray.get(sandbox.get_server_address.remote())\n", + "\n", + "tool_config = {\n", + " \"tools\": [\n", + " {\n", + " \"class_name\": \"sandbox.SandboxTool\",\n", + " \"config\": {\n", + " \"type\": \"native\",\n", + " \"sandbox_fusion_url\": f\"http://{sandbox_address}/run_code\",\n", + " },\n", + " },\n", + " ],\n", + "}\n", + "\n", + "tool_config_path = \"tool_config.json\"\n", + "with open(tool_config_path, \"w\") as f:\n", + " json.dump(tool_config, f)" + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/tmp/ipykernel_174199/3963810189.py:3: UserWarning: \n", + "The version_base parameter is not specified.\n", + "Please specify a compatability version level, or None.\n", + "Will assume defaults for version 1.1\n", + " with initialize_config_dir(config_dir=verl_config_dir):\n" + ] + } + ], + "source": [ + "from hydra import compose, initialize_config_dir\n", + "\n", + "with initialize_config_dir(config_dir=verl_config_dir):\n", + " config = compose(\n", + " config_name=\"ppo_trainer\",\n", + " overrides=[\n", + " \"algorithm.adv_estimator=grpo\",\n", + " \"data.train_files=\" + train_file,\n", + " \"data.val_files=\" + test_file,\n", + " \"data.return_raw_chat=True\",\n", + " \"data.train_batch_size=32\",\n", + " \"data.max_prompt_length=1024\",\n", + " \"data.max_response_length=1024\",\n", + " \"+data.apply_chat_template_kwargs.enable_thinking=False\",\n", + " # actor related\n", + " \"actor_rollout_ref.model.path=\" + model_path,\n", + " \"actor_rollout_ref.actor.ppo_mini_batch_size=8\",\n", + " \"actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=8\",\n", + " \"actor_rollout_ref.actor.fsdp_config.param_offload=True\",\n", + " \"actor_rollout_ref.actor.fsdp_config.optimizer_offload=True\",\n", + " # rollout related\n", + " \"actor_rollout_ref.rollout.name=\" + rollout_name,\n", + " \"actor_rollout_ref.rollout.mode=async\",\n", + " \"actor_rollout_ref.rollout.tensor_model_parallel_size=1\",\n", + " \"actor_rollout_ref.rollout.n=8\",\n", + " \"actor_rollout_ref.rollout.multi_turn.tool_config_path=\" + tool_config_path,\n", + " \"actor_rollout_ref.rollout.agent.default_agent_loop=tool_agent\",\n", + " \"actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8\",\n", + " # trainer related\n", + " \"trainer.val_before_train=True\",\n", + " \"trainer.log_val_generations=10\",\n", + " \"trainer.n_gpus_per_node=8\",\n", + " \"trainer.test_freq=-1\",\n", + " \"trainer.total_training_steps=5\",\n", + " \"trainer.logger=['console','tensorboard', 'wandb']\",\n", + " \"trainer.project_name=verl\",\n", + " \"trainer.experiment_name=\" + os.path.basename(model_path),\n", + " ],\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from verl.trainer.main_ppo import main\n", + "\n", + "main(config)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For demo purpose, we only train 5 steps, you can verify the training process by checking wandb metrics:\n", + "- num_turns: min/max/mean chat conversation turns in each step.\n", + "- critic rewards: min/max/mean critic rewards in each step.\n", + "\n", + "For more realistic agentic RL training, please refer to our recipe:\n", + "- [retool](https://github.com/verl-project/verl-recipe/tree/main/retool): implementation of paper [ReTool: Reinforcement Learning for Strategic Tool Use in LLMs](https://arxiv.org/abs/2504.11536)\n", + "- [collabllm](https://github.com/verl-project/verl-recipe/tree/main/collabllm): implementation of paper [CollabLLM: From Passive Responders to Active Collaborators](https://arxiv.org/pdf/2502.00640)\n", + "- [deepeyes](https://github.com/verl-project/verl-recipe/tree/main/deepeyes): implementation of paper [DeepEyes: Incentivizing \"Thinking with Images\" via Reinforcement Learning](https://arxiv.org/abs/2505.14362)" + ] + } + ], + "metadata": { + "fileId": "398ea641-8a51-4a0b-b64e-6b7cd6b72164", + "filePath": "/opt/tiger/open_verl/examples/agent_loop_tutorial.ipynb", + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.3" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} \ No newline at end of file diff --git a/verl/examples/tutorial/agent_loop_get_started/sandbox.py b/verl/examples/tutorial/agent_loop_get_started/sandbox.py new file mode 100644 index 0000000000000000000000000000000000000000..6478173431796c24575d17a4808a64223cfd876e --- /dev/null +++ b/verl/examples/tutorial/agent_loop_get_started/sandbox.py @@ -0,0 +1,69 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import re + +import aiohttp +from transformers.utils import get_json_schema + +from verl.tools.base_tool import BaseTool, OpenAIFunctionToolSchema, ToolResponse + + +class SandboxTool(BaseTool): + def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema): + super().__init__(config, tool_schema) + # Different model may use different code pattern, e.g. python, py, etc. + self.code_pattern = re.compile(r"```py(.*?)```", re.DOTALL) + + async def code_interpreter(self, code: str) -> str: + """Execute the code in the sandbox. + + Args: + code: The code to be executed. + + Returns: + str: The output of the code execution. + """ + async with aiohttp.ClientSession() as session: + async with session.post( + self.config.get("sandbox_fusion_url"), + json={"code": code}, + ) as resp: + resp.raise_for_status() + result = await resp.json() + stdout, stderr = result["run_result"]["stdout"], result["run_result"]["stderr"] + return stdout + stderr + + def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema: + schema = get_json_schema(self.code_interpreter) + return OpenAIFunctionToolSchema(**schema) + + async def execute(self, instance_id: str, parameters: dict, **kwargs) -> tuple[str, float, dict]: + code = parameters["code"] + matches = self.code_pattern.findall(code) + if matches: + code = matches[0].strip() + + # NOTE: Some script may not explicitly print result, we need to add a print statement to the end of the script. + # More better way is to SFT the model to make it print result by default, we skip SFT stage in this tutorial. + lines = code.split("\n") + for i, line in reversed(list(enumerate(lines))): + if line == "": + continue + if not lines[i].startswith("print"): + lines[i] = f"print({line})" + break + code = "\n".join(lines) + + result = await self.code_interpreter(code) + return ToolResponse(text=result), 0.0, {} diff --git a/verl/examples/tutorial/ray/tutorial.ipynb b/verl/examples/tutorial/ray/tutorial.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..ca176af0f7940f705281de7ce707d1fa27238c02 --- /dev/null +++ b/verl/examples/tutorial/ray/tutorial.ipynb @@ -0,0 +1,963 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0ddc582b", + "metadata": {}, + "source": [ + "# VeRL Ray API Tutorial" + ] + }, + { + "cell_type": "markdown", + "id": "71fe3b94", + "metadata": {}, + "source": [ + "## Chapter 1: Ray Basics" + ] + }, + { + "cell_type": "code", + "execution_count": 144, + "id": "1347d381", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "import os" + ] + }, + { + "cell_type": "code", + "execution_count": 145, + "id": "e75b9d44", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "import warnings\n", + "\n", + "import ray\n", + "import torch\n", + "\n", + "warnings.filterwarnings(\"ignore\")" + ] + }, + { + "cell_type": "code", + "execution_count": 146, + "id": "2e90ae00", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "2024-11-01 17:27:19,132\tINFO worker.py:1752 -- Started a local Ray instance.\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "9cc9d2ccbdfb48918c8fd6cd13a0807a", + "version_major": 2, + "version_minor": 0 + }, + "text/html": [ + "
\n", + "
\n", + "
\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
Python version:3.9.2
Ray version:2.10.0
\n", + "\n", + "
\n", + "
\n" + ], + "text/plain": [ + "RayContext(dashboard_url='', python_version='3.9.2', ray_version='2.10.0', ray_commit='09abba26b5bf2707639bb637c208d062a47b46f6')" + ] + }, + "execution_count": 146, + "metadata": {}, + "output_type": "execute_result" + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\u001b[36m(GPUAccumulator pid=224400)\u001b[0m rank 0, value: tensor([1.], device='cuda:0')\n", + "\u001b[36m(GPUAccumulator pid=225234)\u001b[0m rank 2, value: tensor([3.], device='cuda:0')\n", + "\u001b[36m(GPUAccumulator pid=225607)\u001b[0m rank 0, value: tensor([2.], device='cuda:0')\n", + "\u001b[36m(GPUAccumulator pid=226423)\u001b[0m rank 1, value: tensor([3.], device='cuda:0')\n", + "\u001b[36m(GPUAccumulator pid=226857)\u001b[0m rank 3, value: tensor([6.], device='cuda:0')\n", + "\u001b[36m(GPUAccumulatorDecorator pid=227475)\u001b[0m 10\n", + "\u001b[36m(GPUAccumulatorDecorator pid=227475)\u001b[0m rank 0, value: tensor([10.], device='cuda:0')\n", + "\u001b[36m(GPUAccumulatorDecorator pid=227655)\u001b[0m rank 1, value: tensor([11.], device='cuda:0')\n" + ] + } + ], + "source": [ + "# Build a local ray cluster. The head node and worker node are on this machine\n", + "ray.init()" + ] + }, + { + "cell_type": "markdown", + "id": "a127e4e4", + "metadata": {}, + "source": [ + "Implement an Accumulator class." + ] + }, + { + "cell_type": "code", + "execution_count": 147, + "id": "20e7b9a3", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "@ray.remote\n", + "class Accumulator:\n", + " def __init__(self):\n", + " self.value = 0\n", + "\n", + " def add(self, x):\n", + " self.value += x\n", + "\n", + " def get_value(self):\n", + " return self.value" + ] + }, + { + "cell_type": "code", + "execution_count": 148, + "id": "3b80098c", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "# Instantiate an accumulator. Accumulator can be viewed as a process, acting as an RPC service.\n", + "accumulator = Accumulator.remote()" + ] + }, + { + "cell_type": "code", + "execution_count": 149, + "id": "b14b1009", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "0\n" + ] + } + ], + "source": [ + "value_ref = accumulator.get_value.remote() # Check the current value. Note that this function returns immediately and does not actually wait for the remote execution to complete.\n", + "# Get the value\n", + "value = ray.get(value_ref)\n", + "print(value)" + ] + }, + { + "cell_type": "code", + "execution_count": 150, + "id": "513a84b3", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "10\n" + ] + } + ], + "source": [ + "# Accumulate, then check the result.\n", + "accumulator.add.remote(10) # Similarly, the 'add' here will return immediately.\n", + "new_value = ray.get(accumulator.get_value.remote())\n", + "print(new_value)" + ] + }, + { + "cell_type": "markdown", + "id": "3c332fe0", + "metadata": {}, + "source": [ + "## Chapter 2: Resource Pool and RayWorkerGroup\n", + "In the previous example, it was a simple single-process worker. \n", + "In this example, we implement a worker with a GPU and form a RayWorkerGroup. Within this RayWorkerGroup, we implement a simple operation of an accumulator." + ] + }, + { + "cell_type": "code", + "execution_count": 151, + "id": "04229afb", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "from verl.single_controller.base import Worker\n", + "from verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup, merge_resource_pool" + ] + }, + { + "cell_type": "code", + "execution_count": 152, + "id": "0d0dbd58", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "resource_pool = RayResourcePool([4], use_gpu=True)" + ] + }, + { + "cell_type": "code", + "execution_count": 153, + "id": "68f6838a", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "@ray.remote\n", + "class GPUAccumulator(Worker):\n", + " def __init__(self) -> None:\n", + " super().__init__()\n", + " # The initial value of each rank is the same as the rank\n", + " self.value = torch.zeros(size=(1,), device=\"cuda\") + self.rank\n", + "\n", + " def add(self, x):\n", + " self.value += x\n", + " print(f\"rank {self.rank}, value: {self.value}\")\n", + " return self.value.cpu()" + ] + }, + { + "cell_type": "code", + "execution_count": 154, + "id": "23aad8fe", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[tensor([1.]), tensor([2.]), tensor([3.]), tensor([4.])]\n" + ] + } + ], + "source": [ + "# Each worker's initial value is its rank, and then each rank's value is incremented by 1, so the values obtained on each rank are [1, 2, 3, 4]\n", + "class_with_args = RayClassWithInitArgs(cls=GPUAccumulator)\n", + "worker_group = RayWorkerGroup(resource_pool, class_with_args)\n", + "print(worker_group.execute_all_sync(\"add\", x=[1, 1, 1, 1]))" + ] + }, + { + "cell_type": "markdown", + "id": "e6705284", + "metadata": {}, + "source": [ + "The principle of parameter passing: The input parameter is a list of length world_size, where each element in the list is dispatched respectively to each worker in the RayWorkerGroup. \n", + "The return parameter is also a list, corresponding to the return value of each worker." + ] + }, + { + "cell_type": "markdown", + "id": "d25c2412", + "metadata": {}, + "source": [ + "### GPU Resource Sharing" + ] + }, + { + "cell_type": "markdown", + "id": "f74f6d24", + "metadata": {}, + "source": [ + "RayWorkerGroups mapped to the same resource pool share the GPU. In this example, we implement three resource pools: the first occupies 4 GPUs, the second also occupies 4 GPUs, and the last occupies all 8 GPUs. Among them, the first resource pool reuses the resource pool mentioned above." + ] + }, + { + "cell_type": "code", + "execution_count": 155, + "id": "49f9c06f", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "# Create a new resource pool and then merge the newly created resource pool with the previous one.\n", + "resource_pool_1 = RayResourcePool([4], use_gpu=True, name_prefix=\"a\")\n", + "resource_pool_merge = merge_resource_pool(resource_pool, resource_pool_1)" + ] + }, + { + "cell_type": "code", + "execution_count": 156, + "id": "05c2e305", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "# Establish a RayWorkerGroup on the newly created resource pool.\n", + "worker_group_1 = RayWorkerGroup(resource_pool_1, class_with_args)\n", + "worker_group_merge = RayWorkerGroup(resource_pool_merge, class_with_args)" + ] + }, + { + "cell_type": "code", + "execution_count": 157, + "id": "6b9b13f4", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[tensor([2.]), tensor([3.]), tensor([4.]), tensor([5.])]\n" + ] + } + ], + "source": [ + "# Run 'add' on the second set of 4 GPUs; the result should be [2, 3, 4, 5].\n", + "output_1 = worker_group_1.execute_all_sync(\"add\", x=[2, 2, 2, 2])\n", + "print(output_1)" + ] + }, + { + "cell_type": "code", + "execution_count": 158, + "id": "d856d030", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[tensor([3.]), tensor([4.]), tensor([5.]), tensor([6.]), tensor([7.]), tensor([8.]), tensor([9.]), tensor([10.])]\n" + ] + } + ], + "source": [ + "# Run 'add' on the merged set of 8 GPUs; the result should be [3, 4, 5, 6, 7, 8, 9, 10].\n", + "output_merge = worker_group_merge.execute_all_sync(\"add\", x=[3, 3, 3, 3, 3, 3, 3, 3])\n", + "print(output_merge)" + ] + }, + { + "cell_type": "code", + "execution_count": 159, + "id": "33a4628c", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "4 4 8\n" + ] + } + ], + "source": [ + "print(worker_group.world_size, worker_group_1.world_size, worker_group_merge.world_size)" + ] + }, + { + "cell_type": "markdown", + "id": "3df19d13", + "metadata": {}, + "source": [ + "## Chapter 3: Data Dispatch, Execution and Collection" + ] + }, + { + "cell_type": "markdown", + "id": "acb22d9d", + "metadata": {}, + "source": [ + "In the above example, we used the `execute_all_sync` function in the RayWorkerGroup to dispatch data from the driver to each worker. This is very inconvenient for coding. \n", + "In this chapter, we use the form of function decorators to allow RayWorkerGroup to directly call functions written in the Worker, and to greatly simplify parameter passing." + ] + }, + { + "cell_type": "code", + "execution_count": 160, + "id": "35237432", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "from verl.single_controller.base.decorator import Dispatch, Execute, register" + ] + }, + { + "cell_type": "code", + "execution_count": 161, + "id": "88b8ba3b", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "@ray.remote\n", + "class GPUAccumulatorDecorator(Worker):\n", + " def __init__(self) -> None:\n", + " super().__init__()\n", + " # The initial value of each rank is the same as the rank\n", + " self.value = torch.zeros(size=(1,), device=\"cuda\") + self.rank\n", + "\n", + " # map from a single input to all the worker\n", + " @register(Dispatch.ONE_TO_ALL)\n", + " def add(self, x):\n", + " print(x)\n", + " self.value = self.value + x\n", + " print(f\"rank {self.rank}, value: {self.value}\")\n", + " return self.value.cpu()" + ] + }, + { + "cell_type": "code", + "execution_count": 162, + "id": "eddaa043", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "class_with_args = RayClassWithInitArgs(cls=GPUAccumulatorDecorator)\n", + "gpu_accumulator_decorator = RayWorkerGroup(resource_pool_merge, class_with_args)" + ] + }, + { + "cell_type": "code", + "execution_count": 163, + "id": "10087c91", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[tensor([10.]), tensor([11.]), tensor([12.]), tensor([13.]), tensor([14.]), tensor([15.]), tensor([16.]), tensor([17.])]\n" + ] + } + ], + "source": [ + "# As we can see, 10 is automatically dispatched to each Worker in this RayWorkerGroup.\n", + "print(gpu_accumulator_decorator.add(x=10))" + ] + }, + { + "cell_type": "markdown", + "id": "540ee6ad", + "metadata": {}, + "source": [ + "### Custom Dispatch, Collection\n", + "Users can customize `dispatch` and `collection` function. You only need to write the `dispatch_fn` and `collect_fn` functions yourself. We also support executing RPC only on rank_zero, with specific examples provided below." + ] + }, + { + "cell_type": "code", + "execution_count": 164, + "id": "8e041270", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "from verl.single_controller.base.decorator import Dispatch, collect_all_to_all, register" + ] + }, + { + "cell_type": "code", + "execution_count": 165, + "id": "43b5be31", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "def two_to_all_dispatch_fn(worker_group, *args, **kwargs):\n", + " \"\"\"\n", + " Assume the input is a list of 2. Duplicate the input interleaved and pass to each worker.\n", + " \"\"\"\n", + " for arg in args:\n", + " assert len(arg) == 2\n", + " for i in range(worker_group.world_size - 2):\n", + " arg.append(arg[i % 2])\n", + " for k, v in kwargs.items():\n", + " assert len(v) == 2\n", + " for i in range(worker_group.world_size - 2):\n", + " v.append(v[i % 2])\n", + " return args, kwargs\n", + "\n", + "\n", + "@ray.remote\n", + "class TestActor(Worker):\n", + " # TODO: pass *args and **kwargs is bug prone and not very convincing\n", + " def __init__(self, x) -> None:\n", + " super().__init__()\n", + " self._x = x\n", + "\n", + " def foo(self, y):\n", + " return self._x + y\n", + "\n", + " @register(dispatch_mode=Dispatch.ALL_TO_ALL, execute_mode=Execute.RANK_ZERO)\n", + " def foo_rank_zero(self, x, y):\n", + " return self._x + y + x\n", + "\n", + " @register(dispatch_mode={\"dispatch_fn\": two_to_all_dispatch_fn, \"collect_fn\": collect_all_to_all})\n", + " def foo_custom(self, x, y):\n", + " return self._x + y + x" + ] + }, + { + "cell_type": "code", + "execution_count": 166, + "id": "83ec6609", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "class_with_args = RayClassWithInitArgs(cls=TestActor, x=2)\n", + "worker_group = RayWorkerGroup(resource_pool, class_with_args)" + ] + }, + { + "cell_type": "code", + "execution_count": 167, + "id": "62c58d8a", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "output_ref = worker_group.foo_custom(x=[1, 2], y=[5, 6])\n", + "assert output_ref == [8, 10, 8, 10]\n", + "\n", + "output_ref = worker_group.foo_rank_zero(x=1, y=2)\n", + "assert output_ref == 5" + ] + }, + { + "cell_type": "code", + "execution_count": 168, + "id": "14689353", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "8\n" + ] + } + ], + "source": [ + "print(gpu_accumulator_decorator.world_size)" + ] + }, + { + "cell_type": "code", + "execution_count": 169, + "id": "2c80bbf4", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "# Shutdown ray cluster\n", + "ray.shutdown()" + ] + }, + { + "cell_type": "markdown", + "id": "a5c8151c", + "metadata": {}, + "source": [ + "## Chapter 4: NVMegatronRayWorkerGroup" + ] + }, + { + "cell_type": "markdown", + "id": "cd5680e9", + "metadata": {}, + "source": [ + "Due to the Ray issue, we can only support max_colocate_count=1 in RayResourcePool for now. \n", + "This means that each GPU can only have one process.\n", + "We can support max_colocate > 1 when applying this pull request: https://github.com/ray-project/ray/pull/44385" + ] + }, + { + "cell_type": "markdown", + "id": "92724419", + "metadata": {}, + "source": [ + "Therefore, we need to restart the ray and initialize a new resource_pool to demonstrate the **NVMegatronRayWorkerGroup**" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9b038538", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "# Build a local ray cluster. The head node and worker node are on this machine\n", + "ray.init()" + ] + }, + { + "cell_type": "markdown", + "id": "ebfd8798", + "metadata": {}, + "source": [ + "Finally, we implement a `NVMegatronRayWorkerGroup`, within which we create a Megatron and then run a tensor parallel (tp) split Llama mlp layer. Here, we use a complex dispatch mode, `Megatron_COMPUTE`. This dispatch mode assumes that user passes the data partitioned by DP dimension. The data is dispatched to all tp/pp ranks within the same dp group, and ultimately only collects output data from tp=0 and the last pp. In this way, for users that only write code on the driver, the Megatron behind the RPC becomes transparent." + ] + }, + { + "cell_type": "code", + "execution_count": 171, + "id": "5a032154", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "/opt/tiger/Megatron-LM\n", + "/opt/tiger/Megatron-LM/megatron/__init__.py\n" + ] + } + ], + "source": [ + "import sys\n", + "\n", + "current_pythonpath = os.environ.get(\"PYTHONPATH\", \"\")\n", + "\n", + "new_path = \"/opt/tiger/Megatron-LM\"\n", + "\n", + "new_pythonpath = f\"{new_path}:{current_pythonpath}\" if current_pythonpath else new_path\n", + "\n", + "os.environ[\"PYTHONPATH\"] = new_pythonpath\n", + "\n", + "print(new_path)\n", + "sys.path.append(new_path)\n", + "\n", + "import megatron\n", + "\n", + "print(megatron.__file__)" + ] + }, + { + "cell_type": "code", + "execution_count": 172, + "id": "8c84cd5a", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "from megatron.core import parallel_state as mpu\n", + "from omegaconf import OmegaConf\n", + "\n", + "from verl.single_controller.base.decorator import Dispatch, Execute, register\n", + "from verl.single_controller.base.megatron.worker import MegatronWorker\n", + "from verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup\n", + "from verl.single_controller.ray.megatron import NVMegatronRayWorkerGroup" + ] + }, + { + "cell_type": "code", + "execution_count": 173, + "id": "1b1debcc", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "resource_pool = RayResourcePool([4], use_gpu=True, max_colocate_count=1)" + ] + }, + { + "cell_type": "code", + "execution_count": 174, + "id": "bccbe081", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "@ray.remote\n", + "class MLPLayerWorker(MegatronWorker):\n", + " def __init__(self):\n", + " super().__init__()\n", + " rank = int(os.environ[\"LOCAL_RANK\"])\n", + " torch.distributed.init_process_group(backend=\"nccl\")\n", + " torch.cuda.set_device(rank)\n", + "\n", + " mpu.initialize_model_parallel(\n", + " tensor_model_parallel_size=4,\n", + " pipeline_model_parallel_size=1,\n", + " virtual_pipeline_model_parallel_size=None,\n", + " pipeline_model_parallel_split_rank=None,\n", + " use_sharp=False,\n", + " context_parallel_size=1,\n", + " expert_model_parallel_size=1,\n", + " nccl_communicator_config_path=None,\n", + " )\n", + " from megatron.core import tensor_parallel\n", + "\n", + " tensor_parallel.model_parallel_cuda_manual_seed(10)\n", + "\n", + " @register(Dispatch.ONE_TO_ALL)\n", + " def init_model(self, config):\n", + " from omegaconf import OmegaConf\n", + "\n", + " from verl.models.llama.megatron.layers import ParallelLlamaMLP\n", + " from verl.utils.megatron_utils import init_model_parallel_config\n", + "\n", + " megatron_config = OmegaConf.create(\n", + " {\n", + " \"sequence_parallel\": False,\n", + " \"param_dtype\": \"fp32\",\n", + " \"tensor_model_parallel_size\": mpu.get_tensor_model_parallel_world_size(),\n", + " \"pipeline_model_parallel_rank\": mpu.get_pipeline_model_parallel_rank(),\n", + " \"pipeline_model_parallel_size\": mpu.get_pipeline_model_parallel_world_size(),\n", + " \"virtual_pipeline_model_parallel_rank\": mpu.get_virtual_pipeline_model_parallel_rank(),\n", + " \"virtual_pipeline_model_parallel_size\": mpu.get_virtual_pipeline_model_parallel_world_size(),\n", + " }\n", + " )\n", + "\n", + " megatron_config = init_model_parallel_config(megatron_config)\n", + " self.parallel_layer = ParallelLlamaMLP(config=config, megatron_config=megatron_config)\n", + "\n", + " @register(Dispatch.ONE_TO_ALL)\n", + " def get_weights(self):\n", + " output = {}\n", + " for key, val in self.parallel_layer.named_parameters():\n", + " output[key] = val\n", + " return output\n", + "\n", + " @register(Dispatch.MEGATRON_COMPUTE)\n", + " def run_layer(self, x):\n", + " x = x.to(\"cuda\")\n", + " y = self.parallel_layer(x)\n", + " return y" + ] + }, + { + "cell_type": "code", + "execution_count": 175, + "id": "a655271d", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "layer_cls = RayClassWithInitArgs(cls=MLPLayerWorker)\n", + "layer_worker_group = NVMegatronRayWorkerGroup(\n", + " resource_pool=resource_pool,\n", + " ray_cls_with_init=layer_cls,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 176, + "id": "f105ebee", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "4 4 1 1\n" + ] + } + ], + "source": [ + "print(layer_worker_group.world_size, layer_worker_group.tp_size, layer_worker_group.pp_size, layer_worker_group.dp_size)" + ] + }, + { + "cell_type": "code", + "execution_count": 177, + "id": "38655091", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "ffn_hidden_size = 11008\n", + "batch_size = 16\n", + "seq_len = 2048\n", + "hidden_size = 4096\n", + "\n", + "config = OmegaConf.create(\n", + " {\n", + " \"hidden_size\": hidden_size,\n", + " \"intermediate_size\": ffn_hidden_size,\n", + " \"hidden_act\": \"silu\",\n", + " \"pretraining_tp\": 1,\n", + " \"tp\": layer_worker_group.tp_size,\n", + " }\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 178, + "id": "a026efca", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "x = torch.rand(size=(seq_len, batch_size, hidden_size), dtype=torch.float32)" + ] + }, + { + "cell_type": "code", + "execution_count": 179, + "id": "f5fcaf13", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "data": { + "text/plain": [ + "[None, None, None, None]" + ] + }, + "execution_count": 179, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "layer_worker_group.init_model(config)" + ] + }, + { + "cell_type": "code", + "execution_count": 180, + "id": "3f5cc9b4", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "torch.Size([2048, 16, 4096])\n" + ] + } + ], + "source": [ + "output = layer_worker_group.run_layer(\n", + " [x]\n", + ") # This must be a list of size 1, ensuring that the input equals the data parallel (dp).\n", + "print(output[0].shape)" + ] + }, + { + "cell_type": "code", + "execution_count": 181, + "id": "49792210", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "# Shutdown ray cluster\n", + "ray.shutdown()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.9.2" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/verl/examples/tutorial/skypilot/README.md b/verl/examples/tutorial/skypilot/README.md new file mode 100644 index 0000000000000000000000000000000000000000..78bd8458a83914a75c096dda8ef6e81e519981f1 --- /dev/null +++ b/verl/examples/tutorial/skypilot/README.md @@ -0,0 +1,107 @@ +# verl with SkyPilot + +Run verl reinforcement learning training jobs on Kubernetes clusters or cloud platforms with GPU nodes using [SkyPilot](https://github.com/skypilot-org/skypilot). + +## Installation and Configuration + +### Step 1: Install SkyPilot + +Choose the installation based on your target platform: + +```bash +# For Kubernetes only +pip install "skypilot[kubernetes]" + +# For AWS +pip install "skypilot[aws]" + +# For Google Cloud Platform +pip install "skypilot[gcp]" + +# For Azure +pip install "skypilot[azure]" + +# For multiple platforms +pip install "skypilot[kubernetes,aws,gcp,azure]" +``` + +### Step 2: Configure Your Platform + +See https://docs.skypilot.co/en/latest/getting-started/installation.html + +### Step 3: Set Up Environment Variables + +Export necessary API keys for experiment tracking: + +```bash +# For Weights & Biases tracking +export WANDB_API_KEY="your-wandb-api-key" + +# For HuggingFace gated models (if needed) +export HF_TOKEN="your-huggingface-token" +``` + +## Examples + +### PPO Training +```bash +sky launch -c verl-ppo verl-ppo.yaml --secret WANDB_API_KEY -y +``` +Runs PPO training on GSM8K dataset using Qwen2.5-0.5B-Instruct model across 2 nodes with H100 GPUs. Based on examples in [`../ppo_trainer/`](../ppo_trainer/). + +### GRPO Training +```bash +sky launch -c verl-grpo verl-grpo.yaml --secret WANDB_API_KEY -y +``` +Runs GRPO (Group Relative Policy Optimization) training on MATH dataset using Qwen2.5-7B-Instruct model. Memory-optimized configuration for 2 nodes. Based on examples in [`../grpo_trainer/`](../grpo_trainer/). + +### Multi-turn Tool Usage Training +```bash +sky launch -c verl-multiturn verl-multiturn-tools.yaml --secret WANDB_API_KEY --secret HF_TOKEN -y +``` +Single-node training with 8xH100 GPUs for multi-turn tool usage with Qwen2.5-3B-Instruct. Includes tool and interaction configurations for GSM8K. Based on examples in [`../sglang_multiturn/`](../sglang_multiturn/) but uses vLLM instead of sglang. + +## Configuration + +The example YAML files are pre-configured with: + +- **Infrastructure**: Kubernetes clusters (`infra: k8s`) - can be changed to `infra: aws` or `infra: gcp`, etc. +- **Docker Image**: verl's official Docker image with CUDA 12.6 support +- **Setup**: Automatically clones and installs verl from source +- **Datasets**: Downloads required datasets during setup phase +- **Ray Cluster**: Configures distributed training across nodes +- **Logging**: Supports Weights & Biases via `--secret WANDB_API_KEY` +- **Models**: Supports gated HuggingFace models via `--secret HF_TOKEN` + +## Launch Command Options + +- `-c `: Cluster name for managing the job +- `--secret KEY`: Pass secrets for API keys (can be used multiple times) +- `-y`: Skip confirmation prompt + +## Monitoring Your Jobs + +### Check cluster status +```bash +sky status +``` + +### View logs +```bash +sky logs verl-ppo # View logs for the PPO job +``` + +### SSH into head node +```bash +ssh verl-ppo +``` + +### Access Ray dashboard +```bash +sky status --endpoint 8265 verl-ppo # Get dashboard URL +``` + +### Stop a cluster +```bash +sky down verl-ppo +``` diff --git a/verl/examples/tutorial/skypilot/verl-grpo.yaml b/verl/examples/tutorial/skypilot/verl-grpo.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7172a4f0480e778324e6f5bfa4b144ba295546cc --- /dev/null +++ b/verl/examples/tutorial/skypilot/verl-grpo.yaml @@ -0,0 +1,99 @@ +resources: + infra: k8s + accelerators: H100:1 + memory: 128+ + image_id: docker:verlai/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.0-fa2.7.4 + ports: 8265 + +num_nodes: 2 + +secrets: + WANDB_API_KEY: + +setup: | + rm -rf verl + git clone https://github.com/verl-project/verl.git + cd verl + pip3 install -v -e .[vllm] + pip3 install flashinfer-python + echo "Downloading Math dataset..." + mkdir -p ~/data/math + python3 "$(pwd)/examples/data_preprocess/math_dataset.py" --local_dir ~/data/math + echo "Math dataset download completed" + +run: | + HEAD_IP=$(echo "$SKYPILOT_NODE_IPS" | head -n1) + NUM_NODES=$SKYPILOT_NUM_NODES + NUM_GPUS_PER_NODE=$SKYPILOT_NUM_GPUS_PER_NODE + + if [ "$SKYPILOT_NODE_RANK" == "0" ]; then + echo "Starting Ray head node..." + ps aux | grep ray | grep 6379 &> /dev/null || ray start --head --disable-usage-stats \ + --port=6379 \ + --dashboard-host=0.0.0.0 \ + --dashboard-port=8265 + + # Wait for all worker nodes to join + retry_count=0 + max_retries=30 + while [ $retry_count -lt $max_retries ]; do + connected_nodes=$(ray status 2>/dev/null | grep -c "node_" || echo "0") + echo "Connected nodes: $connected_nodes/$NUM_NODES (attempt $((retry_count+1))/$max_retries)" + + if [ "$connected_nodes" -ge "$NUM_NODES" ]; then + echo "All nodes connected to Ray cluster" + break + fi + + retry_count=$((retry_count+1)) + sleep 10 + done + + python3 -m verl.trainer.main_ppo \ + algorithm.adv_estimator=grpo \ + data.train_files=$HOME/data/math/train.parquet \ + data.val_files=$HOME/data/math/test.parquet \ + data.train_batch_size=32 \ + data.max_prompt_length=256 \ + data.max_response_length=256 \ + data.filter_overlong_prompts=True \ + data.truncation='error' \ + actor_rollout_ref.model.path=Qwen/Qwen2.5-7B-Instruct \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.model.use_remove_padding=True \ + actor_rollout_ref.actor.ppo_mini_batch_size=16 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \ + actor_rollout_ref.actor.ppo_epochs=1 \ + actor_rollout_ref.actor.use_kl_loss=False \ + actor_rollout_ref.actor.entropy_coeff=0 \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.fsdp_config.param_offload=True \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=1 \ + actor_rollout_ref.rollout.name=vllm \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \ + actor_rollout_ref.rollout.n=1 \ + actor_rollout_ref.rollout.enable_chunked_prefill=True \ + actor_rollout_ref.rollout.max_num_batched_tokens=2048 \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \ + actor_rollout_ref.ref.fsdp_config.param_offload=True \ + algorithm.use_kl_in_reward=False \ + trainer.critic_warmup=0 \ + trainer.logger=[console,wandb] \ + trainer.project_name=verl_math_grpo_demo \ + trainer.experiment_name=qwen25_7b_grpo \ + trainer.n_gpus_per_node=$NUM_GPUS_PER_NODE \ + trainer.nnodes=$NUM_NODES \ + trainer.save_freq=-1 \ + trainer.test_freq=-1 \ + trainer.total_epochs=1 + + else + sleep 15 + echo "Starting Ray worker node..." + ps aux | grep ray | grep $HEAD_IP:6379 &> /dev/null || ray start --address $HEAD_IP:6379 --disable-usage-stats + sleep 10 + fi + + echo "Node setup and Ray start script finished for rank $SKYPILOT_NODE_RANK." \ No newline at end of file diff --git a/verl/examples/tutorial/skypilot/verl-multiturn-tools.yaml b/verl/examples/tutorial/skypilot/verl-multiturn-tools.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d4aef0e45abe37424f56285e55550efda08ac6d6 --- /dev/null +++ b/verl/examples/tutorial/skypilot/verl-multiturn-tools.yaml @@ -0,0 +1,90 @@ +resources: + infra: k8s + accelerators: H100:8 + memory: 128+ + image_id: docker:verlai/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.0-fa2.7.4 + ports: 8265 + +num_nodes: 1 + +secrets: + WANDB_API_KEY: + HF_TOKEN: # in case you're using gated models from the HF hub + +setup: | + rm -rf verl + git clone https://github.com/verl-project/verl.git + cd verl + pip3 install -v -e .[vllm] + pip3 install flashinfer-python + pip install "transformers<4.54.0" # https://github.com/vllm-project/vllm-ascend/issues/2046 + # Download GSM8K dataset for multiturn tool training + echo "Downloading GSM8K dataset..." + mkdir -p ~/data/gsm8k + python3 "$(pwd)/examples/data_preprocess/gsm8k.py" --local_dir ~/data/gsm8k + echo "GSM8K dataset download completed" + +run: | + NUM_GPUS_PER_NODE=$SKYPILOT_NUM_GPUS_PER_NODE + PROJECT_DIR="$(pwd)/verl" + CONFIG_PATH="$PROJECT_DIR/examples/sglang_multiturn/config" + + # Single node setup - no worker coordination needed + echo "Starting Ray head node..." + ps aux | grep ray | grep 6379 &> /dev/null || ray start --head --disable-usage-stats \ + --port=6379 \ + --dashboard-host=0.0.0.0 \ + --dashboard-port=8265 + + cd verl + + python3 -m verl.trainer.main_ppo \ + --config-path="$CONFIG_PATH" \ + --config-name='gsm8k_multiturn_grpo' \ + algorithm.adv_estimator=grpo \ + data.train_batch_size=512 \ + data.max_prompt_length=1024 \ + data.max_response_length=1024 \ + data.filter_overlong_prompts=True \ + data.truncation='error' \ + data.return_raw_chat=True \ + data.train_files=$HOME/data/gsm8k/train.parquet \ + data.val_files=$HOME/data/gsm8k/test.parquet \ + actor_rollout_ref.model.path=Qwen/Qwen2.5-3B-Instruct \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.model.use_remove_padding=True \ + actor_rollout_ref.actor.ppo_mini_batch_size=512 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=32 \ + actor_rollout_ref.actor.use_kl_loss=True \ + actor_rollout_ref.actor.kl_loss_coef=0.001 \ + actor_rollout_ref.actor.kl_loss_type=low_var_kl \ + actor_rollout_ref.actor.entropy_coeff=0 \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.fsdp_config.param_offload=False \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=64 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=4 \ + actor_rollout_ref.rollout.name=vllm \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \ + actor_rollout_ref.rollout.n=16 \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=64 \ + actor_rollout_ref.ref.fsdp_config.param_offload=True \ + algorithm.use_kl_in_reward=False \ + trainer.critic_warmup=0 \ + trainer.logger=[console,wandb] \ + trainer.project_name=verl_multiturn_tools \ + trainer.experiment_name=qwen25_7b_gsm8k_multiturn_tools \ + trainer.n_gpus_per_node=$NUM_GPUS_PER_NODE \ + trainer.nnodes=1 \ + trainer.save_freq=10 \ + trainer.test_freq=5 \ + trainer.total_epochs=10 \ + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=8192 \ + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=8192 \ + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=8192 \ + critic.ppo_max_token_len_per_gpu=8192 \ + critic.forward_max_token_len_per_gpu=8192 \ + actor_rollout_ref.rollout.multi_turn.tool_config_path="$PROJECT_DIR/examples/sglang_multiturn/config/tool_config/gsm8k_tool_config.yaml" \ + actor_rollout_ref.rollout.multi_turn.max_user_turns=1 + + echo "Node setup and Ray start script finished for rank $SKYPILOT_NODE_RANK." \ No newline at end of file diff --git a/verl/examples/tutorial/skypilot/verl-ppo.yaml b/verl/examples/tutorial/skypilot/verl-ppo.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5f6879aa4854557ad570972f3da118a9e0bf1a69 --- /dev/null +++ b/verl/examples/tutorial/skypilot/verl-ppo.yaml @@ -0,0 +1,109 @@ +resources: + infra: k8s + accelerators: H100:1 + memory: 128+ + image_id: docker:verlai/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.0-fa2.7.4 + ports: 8265 + +num_nodes: 2 + +secrets: + WANDB_API_KEY: + +setup: | + rm -rf verl + git clone https://github.com/verl-project/verl.git + cd verl + pip3 install -v -e .[vllm] + pip3 install flashinfer-python + # Download GSM8K dataset - alternative approach + echo "Downloading GSM8K dataset..." + mkdir -p ~/data/gsm8k + # Check if the script exists and use absolute path + if [ -f "$(pwd)/examples/data_preprocess/gsm8k.py" ]; then + python3 "$(pwd)/examples/data_preprocess/gsm8k.py" --local_dir ~/data/gsm8k + else + echo "Warning: gsm8k.py script not found, skipping dataset download" + # You might want to download the dataset manually or use a different approach + fi + echo "GSM8K dataset download completed" + +run: | + # Get the Head node's IP and total number of nodes + HEAD_IP=$(echo "$SKYPILOT_NODE_IPS" | head -n1) + NUM_NODES=$SKYPILOT_NUM_NODES + + # login wandb + # python3 -c "import wandb; wandb.login(relogin=True, key='$WANDB_API_KEY')" + + if [ "$SKYPILOT_NODE_RANK" == "0" ]; then + # Head node starts Ray Head + echo "Starting Ray head node..." + ps aux | grep ray | grep 6379 &> /dev/null || ray start --head --disable-usage-stats \ + --port=6379 \ + --dashboard-host=0.0.0.0 \ + --dashboard-port=8265 + + # Wait for all worker nodes to join the cluster with better checking + echo "Waiting for all nodes to join Ray cluster..." + retry_count=0 + max_retries=30 + while [ $retry_count -lt $max_retries ]; do + connected_nodes=$(ray status 2>/dev/null | grep -c "node_" || echo "0") + echo "Connected nodes: $connected_nodes/$NUM_NODES (attempt $((retry_count+1))/$max_retries)" + + if [ "$connected_nodes" -ge "$NUM_NODES" ]; then + echo "All nodes connected to Ray cluster" + break + fi + + retry_count=$((retry_count+1)) + sleep 10 + done + + if [ $retry_count -eq $max_retries ]; then + echo "WARNING: Not all nodes connected to Ray cluster after $max_retries attempts" + echo "Current Ray status:" + ray status + fi + + python3 -m verl.trainer.main_ppo \ + data.train_files=$HOME/data/gsm8k/train.parquet \ + data.val_files=$HOME/data/gsm8k/test.parquet \ + data.train_batch_size=256 \ + data.max_prompt_length=512 \ + data.max_response_length=256 \ + actor_rollout_ref.model.path=Qwen/Qwen2.5-0.5B-Instruct \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.actor.ppo_mini_batch_size=64 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=1 \ + actor_rollout_ref.rollout.name=vllm \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \ + critic.optim.lr=1e-5 \ + critic.model.path=Qwen/Qwen2.5-0.5B-Instruct \ + critic.ppo_micro_batch_size_per_gpu=4 \ + algorithm.kl_ctrl.kl_coef=0.001 \ + trainer.logger=[console,wandb] \ + trainer.val_before_train=False \ + trainer.default_hdfs_dir=null \ + trainer.n_gpus_per_node=1 \ + trainer.nnodes=2 \ + trainer.save_freq=20 \ + trainer.test_freq=20 \ + trainer.total_epochs=2 \ + trainer.project_name=verl_examples \ + trainer.experiment_name=experiment_name_gsm8k + + else + # Wait for Ray Head to start + sleep 15 + # Worker node starts Ray Worker + echo "Starting Ray worker node..." + ps aux | grep ray | grep $HEAD_IP:6379 &> /dev/null || ray start --address $HEAD_IP:6379 --disable-usage-stats + sleep 10 + fi + + echo "Node setup and Ray start script finished for rank $SKYPILOT_NODE_RANK." \ No newline at end of file diff --git a/verl/examples/tutorial/slurm/ray_on_slurm.slurm b/verl/examples/tutorial/slurm/ray_on_slurm.slurm new file mode 100644 index 0000000000000000000000000000000000000000..86567d811be50e583dd715a3a60cf0053451e891 --- /dev/null +++ b/verl/examples/tutorial/slurm/ray_on_slurm.slurm @@ -0,0 +1,98 @@ +#!/bin/bash +#SBATCH --job-name=verl-ray-on-slurm +#SBATCH --nodes=2 +#SBATCH --ntasks-per-node=1 +#SBATCH --mem=200G +#SBATCH --partition=your-partition +#SBATCH --time=01:00:00 +#SBATCH --account=your-account +#SBATCH --gpus-per-node=4 +#SBATCH --cpus-per-task=64 +#SBATCH --output=slurm-%j.out +#SBATCH --error=slurm-%j.err + +# load necessary modules + +# replace these information with your own +verl_workdir=/path/to/verl +train_files=/path/to/gsm8k/train.parquet +val_files=/path/to/gsm8k/test.parquet +apptainer_image_path=/path/to/verl-ngc.sif +# replace these information with your own + +# Getting the node names +nodes=$(scontrol show hostnames "$SLURM_JOB_NODELIST") +nodes_array=("$nodes") + +head_node=${nodes_array[0]} +head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address) + +# if we detect a space character in the head node IP, we'll +# convert it to an ipv4 address. This step is optional. +if [[ "$head_node_ip" == *" "* ]]; then +IFS=' ' read -ra ADDR <<<"$head_node_ip" +if [[ ${#ADDR[0]} -gt 16 ]]; then + head_node_ip=${ADDR[1]} +else + head_node_ip=${ADDR[0]} +fi +echo "IPV6 address detected. We split the IPV4 address as $head_node_ip" +fi + +port=6379 +ip_head=$head_node_ip:$port +export ip_head +echo "IP Head: $ip_head" + +# make sure we set environment variables before Ray initialization + +printenv + +echo "Starting HEAD at $head_node" +srun --nodes=1 --ntasks=1 -w "$head_node" \ + apptainer run --nv --bind $verl_workdir $apptainer_image_path \ + ray start --head --node-ip-address="$head_node_ip" --port=$port \ + --num-cpus "${SLURM_CPUS_PER_TASK}" --num-gpus "${SLURM_GPUS_PER_NODE}" --block & +# optional, though may be useful in certain versions of Ray < 1.0. +sleep 10 + +# number of nodes other than the head node +worker_num=$((SLURM_JOB_NUM_NODES - 1)) + +for ((i = 1; i <= worker_num; i++)); do + node_i=${nodes_array[$i]} + echo "Starting WORKER $i at $node_i" + srun --nodes=1 --ntasks=1 -w "$node_i" \ + apptainer run --nv --bind $verl_workdir $apptainer_image_path \ + ray start --address "$ip_head" --num-cpus "${SLURM_CPUS_PER_TASK}" --num-gpus "${SLURM_GPUS_PER_NODE}" --block & + sleep 5 +done + +PYTHONUNBUFFERED=1 srun --overlap --nodes=1 --ntasks=1 -w "$head_node" \ + apptainer run --nv --bind $verl_workdir $apptainer_image_path \ + python3 -m verl.trainer.main_ppo \ + algorithm.adv_estimator=gae \ + data.train_files=$train_files \ + data.val_files=$val_files \ + data.train_batch_size=256 \ + data.max_prompt_length=512 \ + data.max_response_length=256 \ + actor_rollout_ref.model.path=Qwen/Qwen2.5-0.5B-Instruct \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.actor.ppo_mini_batch_size=64 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \ + actor_rollout_ref.actor.use_kl_loss=False \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=1 \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \ + critic.optim.lr=1e-5 \ + critic.model.path=Qwen/Qwen2.5-0.5B-Instruct \ + critic.ppo_micro_batch_size_per_gpu=4 \ + algorithm.use_kl_in_reward=False \ + trainer.logger=console \ + trainer.val_before_train=False \ + trainer.n_gpus_per_node="${SLURM_GPUS_PER_NODE}" \ + trainer.nnodes="${SLURM_NNODES}" \ + trainer.save_freq=10 \ + trainer.test_freq=10 \ + trainer.total_epochs=15 2>&1 | tee verl_demo_slurm.log diff --git a/verl/pyproject.toml b/verl/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..a597421e920335f2f177a49bf55016201ad25873 --- /dev/null +++ b/verl/pyproject.toml @@ -0,0 +1,114 @@ +# ------------------------------- +# build-system +# ------------------------------- +[build-system] +requires = [ + "setuptools>=61.0", + "wheel" +] +build-backend = "setuptools.build_meta" + +# ------------------------------- +# project (PEP 621 metadata) +# ------------------------------- +[project] +name = "verl" +# We'll mark the version as "dynamic" because it's read from the file "verl/version/version" +# (PEP 621 calls this "dynamic version"). +# The actual version is specified in the [tool.setuptools.dynamic] section below. +dynamic = ["version", "dependencies", "optional-dependencies", "authors", "urls"] + +description = "verl: Volcano Engine Reinforcement Learning for LLM" +license = {text = "Apache-2.0"} # Changed from file to text format +readme = {file = "README.md", content-type = "text/markdown"} +requires-python = ">=3.10" + +# ------------------------------- +# tool.ruff - Linting configuration +# ------------------------------- +[tool.ruff] +# Note: While the formatter will attempt to format lines such that they remain within the line-length, +# it isn't a hard upper bound, and formatted lines may exceed the line-length. +line-length = 120 +exclude = ["scripts/legacy_model_merger.py"] + +[tool.ruff.lint] +isort = {known-first-party = ["verl"]} +# c.f. https://github.com/vllm-project/vllm/blob/ce8d6b75fc0586045df75ee1568a5b5f9957251b/pyproject.toml +select = [ + # pycodestyle + "E", + # Pyflakes + "F", + # pyupgrade + "UP", + # flake8-bugbear + "B", + # isort + "I", + "G", +] +ignore = [ + # star imports + "F405", "F403", + # lambda expression assignment + "E731", + # Loop control variable not used within loop body + "B007", + # f-string format + "UP032", + # `.log()` statement uses f-string + "G004", + # X | None for type annotations + "UP045", + # deprecated import + "UP035", +] + +# ------------------------------- +# tool.mypy - typechecking config +# ------------------------------- +[tool.mypy] +pretty = true +ignore_missing_imports = true +explicit_package_bases = true +follow_imports = "skip" + +# Blanket silence +ignore_errors = true + +[[tool.mypy.overrides]] +module = [ +"verl.trainer.config.algorithm", +"verl.trainer.ppo.core_algos", +"verl.trainer.ppo.reward", +"verl.workers.reward_manager", +"verl.workers.reward_manager.*", +] +ignore_errors = false + +# ------------------------------- +# tool.setuptools - Additional config +# ------------------------------- +[tool.setuptools] +# True means `setuptools` will attempt to include all relevant files in package_data automatically. +# This corresponds to `include_package_data=True` in setup.py. +include-package-data = true + +# We read the version from a file in 'verl/version/version' +[tool.setuptools.dynamic] +version = {file = "verl/version/version"} + +# If you need to mimic `package_dir={'': '.'}`: +[tool.setuptools.package-dir] +"" = "." + +# If you need to include specific non-Python data (like YAML files or version file): +# This is the rough equivalent of package_data={'': ['version/*'], 'verl': ['trainer/config/*.yaml']} +[tool.setuptools.package-data] +verl = [ + "version/*", + "trainer/config/*.yaml", + "trainer/config/*/*.yaml", + "experimental/*/config/*.yaml", +] diff --git a/verl/requirements-npu.txt b/verl/requirements-npu.txt new file mode 100644 index 0000000000000000000000000000000000000000..2ec0485509210c9e7c9c079dc321a42548c019c6 --- /dev/null +++ b/verl/requirements-npu.txt @@ -0,0 +1,23 @@ +# requirements.txt records the full set of dependencies for development +accelerate +bytecode +codetiming +datasets +dill +hydra-core +numpy<2.0.0 +pandas +peft>=0.15.2 +pyarrow>=15.0.0 +pybind11 +pylatexenc +tensordict>=0.8.0,<=0.10.0,!=0.9.0 +ray[default] +wandb +mathruler +torchdata +einops +qwen_vl_utils +hf_transfer +triton-ascend==3.2.0 +TransferQueue==0.1.7 \ No newline at end of file diff --git a/verl/requirements-test.txt b/verl/requirements-test.txt new file mode 100644 index 0000000000000000000000000000000000000000..92b4996eeb393ac21a203c8b3bc256abedaee87d --- /dev/null +++ b/verl/requirements-test.txt @@ -0,0 +1,5 @@ +pytest +pre-commit +py-spy +pytest-asyncio +pytest-rerunfailures diff --git a/verl/requirements.txt b/verl/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..5769c9c38357525cb50b2676c4d520fb911c9373 --- /dev/null +++ b/verl/requirements.txt @@ -0,0 +1,27 @@ +# requirements.txt records the full set of dependencies for development +accelerate +codetiming +datasets +dill +hydra-core +liger-kernel +numpy<2.0.0 +pandas +peft +pyarrow>=19.0.0 +pybind11 +pylatexenc +pre-commit +ray[default] +tensordict>=0.8.0,<=0.10.0,!=0.9.0 +torchdata +transformers +# vllm==0.8.4 +wandb +packaging>=20.0 +uvicorn +fastapi +latex2sympy2_extended +math_verify +tensorboard +TransferQueue==0.1.7 diff --git a/verl/scripts/__init__.py b/verl/scripts/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1ce90c5eb352d85c59105c0dc85b5f1dd576f095 --- /dev/null +++ b/verl/scripts/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/verl/scripts/converter_hf_to_mcore.py b/verl/scripts/converter_hf_to_mcore.py new file mode 100644 index 0000000000000000000000000000000000000000..6e7cdf2b5ab16240787c1455bdb4b1c12c2ecd8a --- /dev/null +++ b/verl/scripts/converter_hf_to_mcore.py @@ -0,0 +1,610 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import os +import warnings +from contextlib import contextmanager +from importlib.metadata import version +from typing import Any, Callable, ContextManager, Optional + +import numpy as np +import torch +import torch.distributed as dist + +try: + # NPU patch + import mindspeed.megatron_adaptor # noqa: F401 + from mindspeed.megatron_adaptor import repatch +except ImportError: + repatch = None + pass + +from accelerate import init_empty_weights +from megatron.core import dist_checkpointing +from megatron.core import parallel_state as mpu +from megatron.core.dist_checkpointing.mapping import ShardedTensor +from megatron.core.dist_checkpointing.serialization import StrictHandling +from megatron.core.models.gpt.gpt_model import ModelType +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from packaging.version import Version +from transformers import AutoConfig + +from verl.model_merger.megatron_model_merger import get_dynamic_pipeline_shards +from verl.models.mcore import hf_to_mcore_config +from verl.utils.device import get_device_name, get_torch_device +from verl.utils.megatron_utils import get_model + + +def _init_args(): + """ + Examples: + + 1. single rank conversion for any model: + > python converter_hf_to_mcore.py --hf_model_path %{hf_model} --output_path ${output_path} + 2. distributed conversion for DeepseekV3 671B: + > torchrun --nproc_per_node 1 --nnodes 4 --node_rank ${RANK} converter_hf_to_mcore.py \ + --hf_model_path %{hf_model} --output_path ${output_path} + """ + parser = argparse.ArgumentParser() + parser.add_argument("--hf_model_path", type=str, required=True, help="The path for the huggingface model") + parser.add_argument("--output_path", type=str, required=True, help="The path for the output mcore model") + parser.add_argument("--pp_size", type=int, default=1, help="pipeline model parallel size") + parser.add_argument("--ep_size", type=int, default=1, help="expert model parallel size") + parser.add_argument("--use_cpu_initialization", action="store_true", help="Whether to use cpu initialization") + parser.add_argument("--test", action="store_true", help="Whether to test the conversion") + parser.add_argument("--trust_remote_code", action="store_true", help="Whether to trust remote code") + args = parser.parse_args() + return args + + +def test_conversion(megatron_model_provider, tfconfig, output_path, model): + ########### test ########### + # load model + model_test = get_model( + model_provider_func=megatron_model_provider, + model_type=ModelType.encoder_or_decoder, + wrap_with_ddp=True, + transformer_config=tfconfig, + ) + ref_state_dict = model_test[0].module.sharded_state_dict() + dist_checkpointing.load(ref_state_dict, output_path, strict=StrictHandling.ASSUME_OK_UNEXPECTED) + + dut_state_dict = model[0].module.state_dict() + for name in dut_state_dict.keys(): + if dut_state_dict[name] is None: + print(f"[Warning] {name} is none in dut_state_dict") + continue + dut_data = dut_state_dict[name].data + if name in ref_state_dict: + ref_data = ref_state_dict[name] + if isinstance(ref_data, ShardedTensor): + ref_data = ref_data.data.view(ref_data.local_shape) + else: + ref_data = ref_data.data + assert dut_data.shape == ref_data.shape, f"{name=} {dut_data.shape=} {ref_data.shape=}" + assert (dut_data == ref_data).all(), f"{name} is not equal" + print(f"{name} is equal") + else: + print(f"[Warning] {name} is not in ref_state_dict") + for name in ref_state_dict.keys(): + if ref_state_dict[name] is None: + print(f"[Warning] {name} is none in ref_state_dict") + continue + ref_data = ref_state_dict[name] + if isinstance(ref_data, ShardedTensor): + ref_data = ref_data.data.view(ref_data.local_shape) + else: + ref_data = ref_data.data + if name in dut_state_dict: + dut_data = dut_state_dict[name].data + assert dut_data.shape == ref_data.shape, f"{name=} {dut_data.shape=} {ref_data.shape=}" + assert (dut_data == ref_data).all(), f"{name} is not equal" + print(f"{name} is equal") + else: + print(f"[Warning] {name} is not in dut_state_dict") + print("Conversion test passed!") + + +@torch.inference_mode() +def convert_checkpoint_from_transformers_to_megatron( + hf_model, model, hf_config, layer_start_end: Optional[tuple[int, int]] = None +): + if layer_start_end is None: + layer_start_end = (0, len(model.decoder.layers)) + layer_start, layer_end = layer_start_end + pp_rank = mpu.get_pipeline_model_parallel_rank() + pp_size = mpu.get_pipeline_model_parallel_world_size() + ep_rank = mpu.get_expert_model_parallel_rank() + ep_size = mpu.get_expert_model_parallel_world_size() + numel = 0 + + num_attention_heads = hf_config.num_attention_heads + num_key_value_heads = hf_config.num_key_value_heads + hidden_dim = hf_config.hidden_size + head_dim = getattr(hf_config, "head_dim", hidden_dim // num_attention_heads) + if num_attention_heads != num_key_value_heads: + print("[WARNING] Converting GQA model") + has_qkv_bias = getattr(hf_config, "qkv_bias", False) or getattr(hf_config, "attention_bias", False) + has_share_expert = getattr(hf_config, "shared_expert_intermediate_size", None) + if pp_rank == 0: + numel += safe_copy(hf_model.model.embed_tokens.weight, model.embedding.word_embeddings.weight) + + assert len(model.decoder.layers) == (layer_end - layer_start), ( + f"Expected {len(model.decoder.layers)} layers, but got {layer_end - layer_start}" + ) + for layer_idx, (layer, hf_layer) in enumerate( + zip(model.decoder.layers, hf_model.model.layers[layer_start:layer_end], strict=True) + ): + global_layer_idx = layer_idx + layer_start + numel_cur = numel + numel += safe_copy(hf_layer.input_layernorm.weight, layer.self_attention.linear_qkv.layer_norm_weight) + + q = hf_layer.self_attn.q_proj.weight.view( + [num_key_value_heads, head_dim * num_attention_heads // num_key_value_heads, -1] + ) + k = hf_layer.self_attn.k_proj.weight.view([num_key_value_heads, head_dim, -1]) + v = hf_layer.self_attn.v_proj.weight.view([num_key_value_heads, head_dim, -1]) + qkv = torch.cat([q, k, v], dim=1).view(-1, hidden_dim).contiguous() + numel += safe_copy(qkv, layer.self_attention.linear_qkv.weight) + + if has_qkv_bias: + q_bias = hf_layer.self_attn.q_proj.bias.view([num_key_value_heads, -1]) + k_bias = hf_layer.self_attn.k_proj.bias.view([num_key_value_heads, -1]) + v_bias = hf_layer.self_attn.v_proj.bias.view([num_key_value_heads, -1]) + qkv_bias = torch.cat([q_bias, k_bias, v_bias], dim=1).view(-1).contiguous() + numel += safe_copy(qkv_bias, layer.self_attention.linear_qkv.bias) + + if hasattr(hf_layer.self_attn, "q_norm"): + numel += safe_copy(hf_layer.self_attn.q_norm.weight.data, layer.self_attention.q_layernorm.weight) + numel += safe_copy(hf_layer.self_attn.k_norm.weight.data, layer.self_attention.k_layernorm.weight) + + numel += safe_copy(hf_layer.self_attn.o_proj.weight, layer.self_attention.linear_proj.weight) + numel += safe_copy(hf_layer.post_attention_layernorm.weight, layer.pre_mlp_layernorm.weight) + + numel += safe_copy(hf_layer.mlp.gate.weight, layer.mlp.router.weight) + + for idx, hf_expert in enumerate(hf_layer.mlp.experts): + num_experts = len(hf_layer.mlp.experts) + num_local_experts = num_experts // ep_size + expert_idx_start = ep_rank * num_local_experts + expert_idx_end = (ep_rank + 1) * num_local_experts + if idx < expert_idx_start or idx >= expert_idx_end: + continue + local_expert_idx = idx - expert_idx_start + + fc1_weight = torch.cat([hf_expert.gate_proj.weight, hf_expert.up_proj.weight]) + numel += safe_copy(fc1_weight, layer.mlp.experts.linear_fc1._parameters[f"weight{local_expert_idx}"]) + numel += safe_copy( + hf_expert.down_proj.weight, layer.mlp.experts.linear_fc2._parameters[f"weight{local_expert_idx}"] + ) + + if has_share_expert: + numel += safe_copy(hf_layer.mlp.shared_expert_gate.weight, layer.mlp.shared_experts.gate_weight) + shared_fc1_weight = torch.cat( + [hf_layer.mlp.shared_expert.gate_proj.weight, hf_layer.mlp.shared_expert.up_proj.weight] + ) + numel += safe_copy(shared_fc1_weight, layer.mlp.shared_experts.linear_fc1.weight) + numel += safe_copy(hf_layer.mlp.shared_expert.down_proj.weight, layer.mlp.shared_experts.linear_fc2.weight) + print(f"{pp_rank=} {global_layer_idx=} {layer_idx=} {numel=} numel this layer={numel - numel_cur}") + + if pp_rank == pp_size - 1: + numel += safe_copy(hf_model.model.norm.weight, model.decoder.final_layernorm.weight) + numel += safe_copy(hf_model.lm_head.weight, model.output_layer.weight) + return numel + + +def safe_copy( + src_tensor: torch.Tensor, + dst_tensor: torch.Tensor, + skip_dtype_assert: bool = False, +): + if not skip_dtype_assert: + if src_tensor.dtype != dst_tensor.dtype: + raise ValueError(f"Get source dtype {src_tensor.dtype}, but target dtype {dst_tensor.dtype}") + assert src_tensor.shape == dst_tensor.shape + dst_tensor.data.copy_(src_tensor.data) + return src_tensor.numel() + + +@torch.inference_mode() +def convert_checkpoint_from_transformers_to_megatron_qwen2_5_vl(hfmodel, mgmodel, hf_config): + mgmodel = mgmodel.bfloat16() + hfmodel = hfmodel.bfloat16() + num_attention_heads = hf_config.num_attention_heads + num_query_groups = hf_config.num_key_value_heads + hidden_size = hf_config.hidden_size + head_dim = hidden_size // num_attention_heads + + # 1. vision model + if Version(version("transformers")) < Version("4.52.0"): + print("Using transformers < 4.52 API to load vision model") + hfvision = hfmodel.visual + else: + hfvision = hfmodel.model.visual + mgvision = mgmodel.vision_model + vision_hidden_size = mgvision.config.hidden_size + vision_num_query_groups = mgvision.config.num_query_groups + vision_head_dim = vision_hidden_size // mgvision.config.num_attention_heads + copied_numel = 0 + safe_copy(hfvision.rotary_pos_emb.inv_freq, mgvision.rotary_pos_emb.inv_freq) + copied_numel += safe_copy(hfvision.patch_embed.proj.weight, mgvision.patch_embed.proj.weight) + for hfblock, mgblock in zip(hfvision.blocks, mgvision.decoder.layers, strict=True): + # norm1 --> linear_qkv.norm + copied_numel += safe_copy(hfblock.norm1.weight, mgblock.self_attention.linear_qkv.layer_norm_weight) + # norm2 --> mlp.linear_fc1.norm + copied_numel += safe_copy(hfblock.norm2.weight, mgblock.mlp.linear_fc1.layer_norm_weight) + # qkv --> self_attention.linear_qkv + converted_weight = ( + hfblock.attn.qkv.weight.view(3, vision_num_query_groups, -1, vision_head_dim, vision_hidden_size) + .transpose(0, 1) + .flatten(1, 2) + .reshape(-1, vision_hidden_size) + .contiguous() + ) + copied_numel += safe_copy(converted_weight, mgblock.self_attention.linear_qkv.weight) + converted_bias = ( + hfblock.attn.qkv.bias.view(3, vision_num_query_groups, -1) + .transpose(0, 1) + .flatten(1, 2) + .view(-1) + .contiguous() + ) + copied_numel += safe_copy(converted_bias, mgblock.self_attention.linear_qkv.bias) + # proj --> self_attention.linear_proj + copied_numel += safe_copy(hfblock.attn.proj.weight, mgblock.self_attention.linear_proj.weight) + copied_numel += safe_copy(hfblock.attn.proj.bias, mgblock.self_attention.linear_proj.bias) + # mlp --> mlp: gate + fc1_weight = torch.cat([hfblock.mlp.gate_proj.weight, hfblock.mlp.up_proj.weight]) + fc1_bias = torch.cat([hfblock.mlp.gate_proj.bias, hfblock.mlp.up_proj.bias]) + copied_numel += safe_copy(fc1_weight, mgblock.mlp.linear_fc1.weight) + copied_numel += safe_copy(fc1_bias, mgblock.mlp.linear_fc1.bias) + copied_numel += safe_copy(hfblock.mlp.down_proj.weight, mgblock.mlp.linear_fc2.weight) + copied_numel += safe_copy(hfblock.mlp.down_proj.bias, mgblock.mlp.linear_fc2.bias) + + # 2. vision projector + hfprojector = hfvision.merger + mgprojector = mgvision.projection + copied_numel += safe_copy(hfprojector.ln_q.weight, mgvision.decoder.final_layernorm.weight) + + copied_numel += safe_copy(hfprojector.mlp[0].weight, mgprojector.encoder.linear_fc1.weight) + copied_numel += safe_copy(hfprojector.mlp[0].bias, mgprojector.encoder.linear_fc1.bias) + copied_numel += safe_copy(hfprojector.mlp[2].weight, mgprojector.encoder.linear_fc2.weight) + copied_numel += safe_copy(hfprojector.mlp[2].bias, mgprojector.encoder.linear_fc2.bias) + n_params = sum([t.numel() for t in hfvision.state_dict().values()]) + assert n_params == copied_numel, f"n_params={n_params} != copied_numel={copied_numel}" + # 3. llm [just Qwen2] + if Version(version("transformers")) < Version("4.52.0"): + print("Using transformers < 4.52 API to load llm") + hfllm = hfmodel.model + else: + hfllm = hfmodel.model.language_model + mgllm = mgmodel.language_model + copied_numel = 0 + copied_numel += safe_copy(hfllm.embed_tokens.weight, mgllm.embedding.word_embeddings.weight) + layermaps = zip(mgllm.decoder.layers, hfllm.layers, strict=True) + for mglayer, hflayer in layermaps: + copied_numel += safe_copy(hflayer.input_layernorm.weight, mglayer.self_attention.linear_qkv.layer_norm_weight) + + q_proj_weight = hflayer.self_attn.q_proj.weight.view(num_query_groups, -1, head_dim, hidden_size) + k_proj_weight = hflayer.self_attn.k_proj.weight.view(num_query_groups, -1, head_dim, hidden_size) + v_proj_weight = hflayer.self_attn.v_proj.weight.view(num_query_groups, -1, head_dim, hidden_size) + qkv_proj = torch.cat([q_proj_weight, k_proj_weight, v_proj_weight], dim=1).view(-1, hidden_size).contiguous() + copied_numel += safe_copy(qkv_proj, mglayer.self_attention.linear_qkv.weight) + + q_proj_bias = hflayer.self_attn.q_proj.bias.view(num_query_groups, -1) + k_proj_bias = hflayer.self_attn.k_proj.bias.view(num_query_groups, -1) + v_proj_bias = hflayer.self_attn.v_proj.bias.view(num_query_groups, -1) + qkv_bias = torch.cat([q_proj_bias, k_proj_bias, v_proj_bias], dim=1).view(-1).contiguous() + copied_numel += safe_copy(qkv_bias, mglayer.self_attention.linear_qkv.bias) + copied_numel += safe_copy(hflayer.self_attn.o_proj.weight, mglayer.self_attention.linear_proj.weight) + + fc1_weight = torch.cat([hflayer.mlp.gate_proj.weight, hflayer.mlp.up_proj.weight]) + copied_numel += safe_copy(fc1_weight, mglayer.mlp.linear_fc1.weight) + + copied_numel += safe_copy(hflayer.mlp.down_proj.weight, mglayer.mlp.linear_fc2.weight) + copied_numel += safe_copy(hflayer.post_attention_layernorm.weight, mglayer.mlp.linear_fc1.layer_norm_weight) + + copied_numel += safe_copy(hfllm.norm.weight, mgllm.decoder.final_layernorm.weight) + if not hf_config.tie_word_embeddings: + safe_copy(hfmodel.lm_head.weight, mgllm.output_layer.weight) + + n_params = sum([t.numel() for t in hfllm.state_dict().values()]) + + assert n_params == copied_numel, f"n_params={n_params} != copied_numel={copied_numel}" + + +@torch.inference_mode() +def convert_checkpoint_from_transformers_to_megatron_dpskv3( + hf_model, + model, + hf_config, + tfconfig, + layer_start_end: Optional[tuple[int, int]] = None, +): + warnings.warn("MTP model is not supported yet", stacklevel=2) + if layer_start_end is None: + layer_start_end = (0, len(model.decoder.layers)) + layer_start, layer_end = layer_start_end + numel: int = 0 + pp_rank = mpu.get_pipeline_model_parallel_rank() + pp_size = mpu.get_pipeline_model_parallel_world_size() + ep_rank = mpu.get_expert_model_parallel_rank() + ep_size = mpu.get_expert_model_parallel_world_size() + + if pp_rank == 0: + numel += safe_copy(hf_model.model.embed_tokens.weight, model.embedding.word_embeddings.weight) + + assert len(model.decoder.layers) == (layer_end - layer_start), ( + f"Expected {len(model.decoder.layers)} layers, but got {layer_end - layer_start}" + ) + for layer_idx, (layer, hf_layer) in enumerate( + zip(model.decoder.layers, hf_model.model.layers[layer_start:layer_end], strict=True) + ): + global_layer_idx = layer_idx + layer_start + numel_cur: int = numel + numel += safe_copy(hf_layer.input_layernorm.weight, layer.input_layernorm.weight) + + if hf_config.q_lora_rank is None: + numel += safe_copy(hf_layer.self_attn.q_proj.weight, layer.self_attention.linear_q_proj.weight) + else: + numel += safe_copy(hf_layer.self_attn.q_a_proj.weight, layer.self_attention.linear_q_down_proj.weight) + numel += safe_copy(hf_layer.self_attn.q_b_proj.weight, layer.self_attention.linear_q_up_proj.weight) + numel += safe_copy( + hf_layer.self_attn.q_a_layernorm.weight, layer.self_attention.linear_q_up_proj.layer_norm_weight + ) + + numel += safe_copy( + hf_layer.self_attn.kv_a_proj_with_mqa.weight, layer.self_attention.linear_kv_down_proj.weight + ) + numel += safe_copy(hf_layer.self_attn.kv_b_proj.weight, layer.self_attention.linear_kv_up_proj.weight) + numel += safe_copy( + hf_layer.self_attn.kv_a_layernorm.weight, layer.self_attention.linear_kv_up_proj.layer_norm_weight + ) + numel += safe_copy(hf_layer.self_attn.o_proj.weight, layer.self_attention.linear_proj.weight) + + if not hasattr(layer.mlp, "router"): + numel += safe_copy(hf_layer.post_attention_layernorm.weight, layer.mlp.linear_fc1.layer_norm_weight) + numel += safe_copy( + torch.cat([hf_layer.mlp.gate_proj.weight, hf_layer.mlp.up_proj.weight]), layer.mlp.linear_fc1.weight + ) + numel += safe_copy(hf_layer.mlp.down_proj.weight, layer.mlp.linear_fc2.weight) + else: + numel += safe_copy(hf_layer.mlp.gate.weight, layer.mlp.router.weight) + # NOTE: the e_score_correction_bias in mcore model will be initialized with bfloat16 and \ + # recover to fp32 in the first forward. There is always a diff in the bias between two models (~0.3%) + numel += safe_copy( + hf_layer.mlp.gate.e_score_correction_bias, layer.mlp.router.expert_bias, skip_dtype_assert=True + ) + if tfconfig.moe_grouped_gemm: + for i, hf_expert in enumerate(hf_layer.mlp.experts): + num_experts = len(hf_layer.mlp.experts) + num_local_experts = num_experts // ep_size + expert_idx_start = ep_rank * num_local_experts + expert_idx_end = (ep_rank + 1) * num_local_experts + if i < expert_idx_start or i >= expert_idx_end: + continue + local_expert_idx = i - expert_idx_start + + fc1_weight = torch.cat([hf_expert.gate_proj.weight, hf_expert.up_proj.weight]) + linear_fc1_weighti = getattr(layer.mlp.experts.linear_fc1, "weight" + str(local_expert_idx)) + numel += safe_copy(fc1_weight, linear_fc1_weighti) + linear_fc2_weighti = getattr(layer.mlp.experts.linear_fc2, "weight" + str(local_expert_idx)) + numel_w2 = safe_copy(hf_expert.down_proj.weight, linear_fc2_weighti) + numel += numel_w2 + else: + for i, hf_expert in enumerate(hf_layer.mlp.experts): + expert = layer.mlp.experts.local_experts[i] + fc1_weight = torch.cat([hf_expert.gate_proj.weight, hf_expert.up_proj.weight]) + numel += safe_copy(fc1_weight, expert.linear_fc1.weight) + numel += safe_copy(hf_expert.down_proj.weight, expert.linear_fc2.weight) + numel += safe_copy(hf_layer.post_attention_layernorm.weight, layer.pre_mlp_layernorm.weight) + shared_fc1_weight = torch.cat( + [hf_layer.mlp.shared_experts.gate_proj.weight, hf_layer.mlp.shared_experts.up_proj.weight] + ) + numel += safe_copy(shared_fc1_weight, layer.mlp.shared_experts.linear_fc1.weight) + numel += safe_copy(hf_layer.mlp.shared_experts.down_proj.weight, layer.mlp.shared_experts.linear_fc2.weight) + print(f"{pp_rank=} {global_layer_idx=} {layer_idx=} {numel=} numel this layer={numel - numel_cur}") + numel_hf_one_layer = sum([i.numel() for i in hf_layer.state_dict().values()]) + if hasattr(layer.mlp, "router"): + numel_hf_one_layer -= numel_w2 * 3 * len(hf_layer.mlp.experts) // ep_size * (ep_size - 1) + assert numel - numel_cur == numel_hf_one_layer, "numel mismatch" + + if pp_rank == pp_size - 1: + numel += safe_copy(hf_model.model.norm.weight, model.decoder.final_layernorm.weight) + if not hf_config.tie_word_embeddings: + numel += safe_copy(hf_model.lm_head.weight, model.output_layer.weight) + print(f"{pp_rank=} {numel=}") + return numel + + +@contextmanager +def noop_context() -> Any: + yield + + +def support_distributed_convert(hf_config: AutoConfig) -> bool: + for arch in ["DeepseekV3ForCausalLM", "Qwen3MoeForCausalLM", "Qwen2MoeForCausalLM"]: + if arch in hf_config.architectures: + return True + return False + + +def convert_hf_to_mcore( + hf_model_path, output_path, pp_size=1, ep_size=1, use_cpu_initialization=False, test=False, trust_remote_code=False +): + os.makedirs(output_path, exist_ok=True) + if len(os.listdir(output_path)) > 0 and not test: + print(f"Output path {output_path} is not empty, skipping conversion") + return + + # init torch distributed and mpu + if "WORLD_SIZE" not in os.environ: + os.environ["RANK"] = "0" + os.environ["WORLD_SIZE"] = "1" + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = "12355" + + torch.distributed.init_process_group("nccl") + + local_rank = os.getenv("LOCAL_RANK", 0) + world_size = dist.get_world_size() + get_torch_device().set_device(f"{get_device_name()}:{local_rank}") + if ep_size * pp_size != world_size: + pp_size = world_size + print(f"pp_size is set to {pp_size}") + + mpu.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=pp_size, + virtual_pipeline_model_parallel_size=None, + context_parallel_size=1, + expert_model_parallel_size=ep_size, + ) + model_parallel_cuda_manual_seed(0) + + # init hf config + hf_config = AutoConfig.from_pretrained(hf_model_path, trust_remote_code=trust_remote_code) + print(hf_config, flush=True) + + if repatch: + if hf_config.architectures[0] == "DeepseekV3ForCausalLM": + config_repatch = dict(multi_head_latent_attention=True) + repatch(config_repatch) + + if world_size > 1 and not support_distributed_convert(hf_config): + raise NotImplementedError(f"distributed conversion is not supported for {hf_config.architectures} yet.") + + pipeline_shards = get_dynamic_pipeline_shards(hf_config.num_hidden_layers, pp_size) + print(f"Pipeline shards: {pipeline_shards}", flush=True) + + tfconfig = hf_to_mcore_config( + hf_config, + torch.bfloat16, + num_layers_in_first_pipeline_stage=pipeline_shards[0] if len(pipeline_shards) > 1 else None, + num_layers_in_last_pipeline_stage=pipeline_shards[-1] if len(pipeline_shards) > 2 else None, + ) + tfconfig.use_cpu_initialization = use_cpu_initialization + tie_word_embeddings = getattr(hf_config, "tie_word_embeddings", False) + + # init megatron model + def megatron_model_provider(pre_process, post_process): + from verl.models.mcore import init_mcore_model + + parallel_model = init_mcore_model( + tfconfig, + hf_config, + pre_process, + post_process, + share_embeddings_and_output_weights=tie_word_embeddings, + value=False, + ) + return parallel_model + + context: Callable[..., ContextManager] = init_empty_weights if use_cpu_initialization else noop_context + with context(): + model = get_model( + model_provider_func=megatron_model_provider, + model_type=ModelType.encoder_or_decoder, + wrap_with_ddp=False, + transformer_config=tfconfig, + ) + + if use_cpu_initialization: + # convert meta device to empty tensor so it can use `copy_` function + model[0].module = model[0].module.to_empty(device="cpu") + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + from transformers import AutoModelForCausalLM, AutoModelForImageTextToText + + # init hf model + if "Qwen2_5_VLForConditionalGeneration" in hf_config.architectures: + hf_model = AutoModelForImageTextToText.from_pretrained( + hf_model_path, torch_dtype=torch.bfloat16, trust_remote_code=trust_remote_code + ) + else: + hf_model = AutoModelForCausalLM.from_pretrained( + hf_model_path, torch_dtype=torch.bfloat16, trust_remote_code=trust_remote_code + ) + hf_state_dict = hf_model.state_dict() + + pp_rank = mpu.get_pipeline_model_parallel_rank() + + # distributed convert + if world_size > 1 and support_distributed_convert(hf_config): + pipeline_cumsum = np.cumsum(pipeline_shards) + layer_start = 0 if pp_rank == 0 else pipeline_cumsum[pp_rank - 1] + layer_end = pipeline_cumsum[pp_rank] + if "DeepseekV3ForCausalLM" in hf_config.architectures: + numel_partial: int = convert_checkpoint_from_transformers_to_megatron_dpskv3( + hf_model, model[0].module, hf_config, tfconfig=tfconfig, layer_start_end=(layer_start, layer_end) + ) + elif "Qwen3MoeForCausalLM" in hf_config.architectures or "Qwen2MoeForCausalLM" in hf_config.architectures: + numel_partial: int = convert_checkpoint_from_transformers_to_megatron( + hf_model, model[0].module, hf_config, layer_start_end=(layer_start, layer_end) + ) + else: + raise NotImplementedError(f"Distributed conversion is not supported for {hf_config.architectures} yet.") + + numel_tensor = torch.tensor([numel_partial]).to(get_device_name()) + dist.all_reduce(numel_tensor, op=dist.ReduceOp.SUM) + numel = int(numel_tensor.cpu().item()) + print(f"total numel={numel} vs {hf_model.num_parameters()=}") + if numel != hf_model.num_parameters(): + warnings.warn(f"numel mismatch: {numel=} != {hf_model.num_parameters()=}", stacklevel=1) + + # load hf state dict to megatron model + elif "Qwen2MoeForCausalLM" in hf_config.architectures: + convert_checkpoint_from_transformers_to_megatron(hf_model, model[0].module, hf_config) + elif "Qwen2_5_VLForConditionalGeneration" in hf_config.architectures: + convert_checkpoint_from_transformers_to_megatron_qwen2_5_vl(hf_model, model[0].module, hf_config) + elif "DeepseekV3ForCausalLM" in hf_config.architectures: + convert_checkpoint_from_transformers_to_megatron_dpskv3(hf_model, model[0].module, hf_config, tfconfig=tfconfig) + elif "Qwen3MoeForCausalLM" in hf_config.architectures: + convert_checkpoint_from_transformers_to_megatron(hf_model, model[0].module, hf_config) + else: + assert not use_cpu_initialization, "use_cpu_initialization is only supported for MoE model" + from verl.models.mcore.loader import load_state_dict_to_megatron_gptmodel + + load_state_dict_to_megatron_gptmodel( + state_dict=hf_state_dict, + wrapped_models=model, + config=hf_config, + params_dtype=torch.bfloat16, + is_value_model=False, + ) + + megatron_state_dict = model[0].module.sharded_state_dict() + del hf_state_dict, hf_model + + # save megatron model + if len(os.listdir(output_path)) == 0: + dist_checkpointing.save(megatron_state_dict, output_path, sharded_strategy=None, async_sharded_save=False) + if test: + test_conversion(megatron_model_provider, tfconfig, output_path, model) + + +if __name__ == "__main__": + args = _init_args() + convert_hf_to_mcore( + args.hf_model_path, + args.output_path, + args.pp_size, + args.ep_size, + args.use_cpu_initialization, + args.test, + args.trust_remote_code, + ) diff --git a/verl/scripts/diagnose.py b/verl/scripts/diagnose.py new file mode 100644 index 0000000000000000000000000000000000000000..cb78f9e5c6297a8ba8e84262253ff385f49e0d2a --- /dev/null +++ b/verl/scripts/diagnose.py @@ -0,0 +1,312 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Diagnose script for checking OS/hardware/python/pip/verl/network. +The output of this script can be a very good hint to issue/problem. +""" + +import os +import platform +import socket +import subprocess +import sys +import time + +import psutil + +try: + from urllib.parse import urlparse + from urllib.request import urlopen +except ImportError: + from urllib2 import urlopen + from urlparse import urlparse +import argparse +import importlib.metadata + +import torch + +URLS = { + "PYPI": "https://pypi.python.org/pypi/pip", +} + +REGIONAL_URLS = { + "cn": { + "PYPI(douban)": "https://pypi.douban.com/", + "Conda(tsinghua)": "https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free/", + } +} + + +def test_connection(name, url, timeout=10): + """Simple connection test""" + urlinfo = urlparse(url) + start = time.time() + try: + socket.gethostbyname(urlinfo.netloc) + except Exception as e: + print("Error resolving DNS for {}: {}, {}".format(name, url, e)) + return + dns_elapsed = time.time() - start + start = time.time() + try: + _ = urlopen(url, timeout=timeout) + except Exception as e: + print("Error open {}: {}, {}, DNS finished in {} sec.".format(name, url, e, dns_elapsed)) + return + load_elapsed = time.time() - start + print("Timing for {}: {}, DNS: {:.4f} sec, LOAD: {:.4f} sec.".format(name, url, dns_elapsed, load_elapsed)) + + +def check_python(): + print("----------Python Info----------") + print("Version :", platform.python_version()) + print("Compiler :", platform.python_compiler()) + print("Build :", platform.python_build()) + print("Arch :", platform.architecture()) + + +def check_pip(): + print("------------Pip Info-----------") + try: + import pip + + print("Version :", pip.__version__) + print("Directory :", os.path.dirname(pip.__file__)) + except ImportError: + print("No corresponding pip install for current python.") + + +def _get_current_git_commit(): + try: + result = subprocess.run(["git", "rev-parse", "HEAD"], capture_output=True, text=True, check=True) + return result.stdout.strip() + except subprocess.CalledProcessError as e: + print(f"Error running git command: {e.stderr.strip()}") + return None + except FileNotFoundError: + print("Did not find command: git") + return None + + +def check_verl(): + print("----------verl Info-----------") + try: + sys.path.insert(0, os.getcwd()) + import verl + + print("Version :", verl.__version__) + verl_dir = os.path.dirname(verl.__file__) + print("Directory :", verl_dir) + try: + commit_hash = _get_current_git_commit() + print("Commit Hash :", commit_hash) + except AttributeError: + print("Commit hash not found. ") + except ImportError as e: + print(f"No verl installed: {e}") + except Exception as e: + import traceback + + if not isinstance(e, IOError): + print("An error occurred trying to import verl.") + print("This is very likely due to missing or incompatible library files.") + print(traceback.format_exc()) + + +def check_os(): + print("----------Platform Info----------") + print("Platform :", platform.platform()) + print("system :", platform.system()) + print("node :", platform.node()) + print("release :", platform.release()) + print("version :", platform.version()) + + +def check_hardware(): + print("----------Hardware Info----------") + print("machine :", platform.machine()) + print("processor :", platform.processor()) + if sys.platform.startswith("darwin"): + pipe = subprocess.Popen(("sysctl", "-a"), stdout=subprocess.PIPE) + output = pipe.communicate()[0] + for line in output.split(b"\n"): + if b"brand_string" in line or b"features" in line: + print(line.strip()) + elif sys.platform.startswith("linux"): + subprocess.call(["lscpu"]) + elif sys.platform.startswith("win32"): + subprocess.call(["wmic", "cpu", "get", "name"]) + + +def check_network(args): + print("----------Network Test----------") + if args.timeout > 0: + print("Setting timeout: {}".format(args.timeout)) + socket.setdefaulttimeout(10) + for region in args.region.strip().split(","): + r = region.strip().lower() + if not r: + continue + if r in REGIONAL_URLS: + URLS.update(REGIONAL_URLS[r]) + else: + import warnings + + warnings.warn("Region {} do not need specific test, please refer to global sites.".format(r), stacklevel=2) + for name, url in URLS.items(): + test_connection(name, url, args.timeout) + + +def check_environment(): + print("----------Environment----------") + for k, v in os.environ.items(): + if k.startswith("VERL_") or k.startswith("OMP_") or k.startswith("KMP_") or k == "CC" or k == "CXX": + print('{}="{}"'.format(k, v)) + + +def check_pip_package_versions(): + packages = ["vllm", "sglang", "ray", "torch"] + for package in packages: + try: + version = importlib.metadata.version(package) + print(f"{package}\t : {version}") + except importlib.metadata.PackageNotFoundError: + print(f"{package}\t : not found.") + + +def check_cuda_versions(): + if torch.cuda.is_available(): + try: + cuda_runtime_version = torch.version.cuda + print(f"CUDA Runtime : {cuda_runtime_version}") + import subprocess + + nvcc_output = subprocess.check_output(["nvcc", "--version"]).decode("utf-8") + cuda_compiler_version = next((line for line in nvcc_output.splitlines() if "release" in line), None) + if cuda_compiler_version: + print(f"CUDA Compiler : {cuda_compiler_version.strip()}") + else: + print("Could not determine CUDA compiler version.") + except FileNotFoundError as e: + print(f"CUDA compiler : Not found: {e}") + except Exception as e: + print(f"An error occurred while checking CUDA versions: {e}") + else: + print("CUDA is not available.") + + +def _get_cpu_memory(): + """ + Get the total CPU memory capacity in GB. + """ + memory = psutil.virtual_memory() + return memory.total / (1024**3) + + +def _get_gpu_info(): + """ + Get GPU type, GPU memory, and GPU count using nvidia-smi command. + """ + try: + result = subprocess.run( + ["nvidia-smi", "--query-gpu=gpu_name,memory.total", "--format=csv,noheader,nounits"], + capture_output=True, + text=True, + check=True, + ) + gpu_lines = result.stdout.strip().split("\n") + gpu_count = len(gpu_lines) + gpu_info = [] + for line in gpu_lines: + gpu_name, gpu_memory = line.split(", ") + gpu_info.append( + { + "type": gpu_name, + "memory": float(gpu_memory) / 1024, # Convert to GB + } + ) + return gpu_count, gpu_info + except (subprocess.CalledProcessError, FileNotFoundError): + print("Failed to execute nvidia-smi command.") + return 0, [] + + +def _get_system_info(): + """ + Get CPU memory capacity, GPU type, GPU memory, and GPU count. + """ + cpu_memory = _get_cpu_memory() + gpu_count, gpu_info = _get_gpu_info() + return {"cpu_memory": cpu_memory, "gpu_count": gpu_count, "gpu_info": gpu_info} + + +def check_system_info(): + print("----------System Info----------") + system_info = _get_system_info() + print(f"CPU Memory\t: {system_info['cpu_memory']:.2f} GB") + print(f"GPU Count\t: {system_info['gpu_count']}") + for i, gpu in enumerate(system_info["gpu_info"]): + print(f"GPU {i + 1}\tType : {gpu['type']}") + print(f"GPU {i + 1}\tMemory : {gpu['memory']:.2f} GB") + + +def parse_args(): + """Parse arguments.""" + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + description="Diagnose script for checking the current system.", + ) + choices = ["python", "pip", "verl", "system", "os", "environment"] + for choice in choices: + parser.add_argument("--" + choice, default=1, type=int, help="Diagnose {}.".format(choice)) + parser.add_argument("--network", default=0, type=int, help="Diagnose network.") + parser.add_argument("--hardware", default=0, type=int, help="Diagnose hardware.") + parser.add_argument( + "--region", + default="", + type=str, + help="Additional sites in which region(s) to test. \ + Specify 'cn' for example to test mirror sites in China.", + ) + parser.add_argument("--timeout", default=10, type=int, help="Connection test timeout threshold, 0 to disable.") + args = parser.parse_args() + return args + + +if __name__ == "__main__": + args = parse_args() + if args.python: + check_python() + + if args.pip: + check_pip() + check_pip_package_versions() + + if args.verl: + check_verl() + + if args.os: + check_os() + + if args.hardware: + check_hardware() + + if args.network: + check_network(args) + + if args.environment: + check_environment() + check_cuda_versions() + + if args.system: + check_system_info() diff --git a/verl/scripts/generate_trainer_config.sh b/verl/scripts/generate_trainer_config.sh new file mode 100644 index 0000000000000000000000000000000000000000..e8283ceadd9bdb10663e18be06c4eeb37d09230d --- /dev/null +++ b/verl/scripts/generate_trainer_config.sh @@ -0,0 +1,53 @@ +#!/usr/bin/env bash +set -euox pipefail + + +# Define config specifications: "config_name:output_file:config_arg" +CONFIG_SPECS=( + "ppo_trainer:_generated_ppo_trainer.yaml:" + "ppo_trainer:_generated_ppo_megatron_trainer.yaml:model_engine=megatron" + "ppo_trainer:_generated_ppo_veomni_trainer.yaml:model_engine=veomni" + "ppo_trainer:_generated_ppo_torchtitan_trainer.yaml:model_engine=torchtitan" +) + +generate_config() { + local config_name="$1" + local output_file="$2" + local config_arg="$3" + + local target_cfg="verl/trainer/config/${output_file}" + local tmp_header=$(mktemp) + local tmp_cfg=$(mktemp) + + echo "# This reference configration yaml is automatically generated via 'scripts/generate_trainer_config.sh'" > "$tmp_header" + echo "# in which it invokes 'python3 scripts/print_cfg.py --cfg job ${config_arg}' to flatten the 'verl/trainer/config/${config_name}.yaml' config fields into a single file." >> "$tmp_header" + echo "# Do not modify this file directly." >> "$tmp_header" + echo "# The file is usually only for reference and never used." >> "$tmp_header" + echo "" >> "$tmp_header" + + python3 scripts/print_cfg.py --cfg job ${config_arg} > "$tmp_cfg" + + cat "$tmp_header" > "$target_cfg" + sed -n '/^actor_rollout_ref/,$p' "$tmp_cfg" >> "$target_cfg" + + rm "$tmp_cfg" "$tmp_header" + + echo "Generated: $target_cfg" +} + +for spec in "${CONFIG_SPECS[@]}"; do + IFS=':' read -r config_name output_file config_arg <<< "$spec" + generate_config "$config_name" "$output_file" "$config_arg" +done + +for spec in "${CONFIG_SPECS[@]}"; do + IFS=':' read -r config_name output_file config_arg <<< "$spec" + target_cfg="verl/trainer/config/${output_file}" + if ! git diff --exit-code -- "$target_cfg" >/dev/null; then + echo "✖ $target_cfg is out of date. Please regenerate via 'scripts/generate_trainer_config.sh' and commit the changes." + exit 1 + fi +done + +echo "All good" +exit 0 diff --git a/verl/scripts/init_random_model.py b/verl/scripts/init_random_model.py new file mode 100644 index 0000000000000000000000000000000000000000..432fffe2b30755a779a5b3fb41d945d81a89f152 --- /dev/null +++ b/verl/scripts/init_random_model.py @@ -0,0 +1,108 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +This script override a model with custom config and random weights, mainly for create small models for +debugging purposes. + +Usage: + python scripts/init_random_model.py \ + --hf_model_path \ + --new_config_path \ + --output_path + +""" + +import argparse +import json +import os +import warnings +from typing import Any + +from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, PretrainedConfig + + +def _init_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--hf_model_path", type=str, required=True, help="The path for the huggingface model") + parser.add_argument("--new_config_path", type=str, required=True, help="The path for the new config file") + parser.add_argument("--output_path", type=str, required=True, help="The path for the output random model") + parser.add_argument( + "--trust_remote_code", + action="store_true", + help="Whether to trust remote code when loading HF model. Disabled by default for security.", + ) + args = parser.parse_args() + return args + + +def check_output_path(output_path: str): + if os.path.exists(output_path): + warnings.warn(f"Output path '{output_path}' already exists. Will do nothing.", stacklevel=2) + exit() + else: + os.makedirs(output_path, exist_ok=True) + print(f"Output path '{output_path}' created.") + + +def check_configs(original_config: dict[str, Any], new_config: dict[str, Any]) -> bool: + """ + Check if the original config and new config are compatible. + This is a placeholder function; actual implementation may vary based on requirements. + """ + # Example check: ensure 'model_type' is the same + if new_config.get("model_type", None) is not None and original_config.get("model_type") != new_config.get( + "model_type" + ): + raise RuntimeError("Model types do not match.") + for key in new_config: + if key not in original_config: + warnings.warn( + f"Key '{key}' in new config does not exist in original config, may not take effect.", stacklevel=2 + ) + + +def init_random_model(hf_model_path, new_config_path, output_path, trust_remote_code: bool = False): + config = AutoConfig.from_pretrained(hf_model_path, trust_remote_code=trust_remote_code) + tokenizer = AutoTokenizer.from_pretrained(hf_model_path, trust_remote_code=trust_remote_code) + config_dict = PretrainedConfig.get_config_dict(hf_model_path)[0] + print(config_dict) + with open(new_config_path) as f: + new_config_dict = json.load(f) + check_configs(config_dict, new_config_dict) + config_dict.update(new_config_dict) + new_confg = config.from_dict(config_dict) + print(f"new_config: {new_confg}") + if trust_remote_code: + model = AutoModelForCausalLM.from_pretrained( + hf_model_path, config=new_confg, trust_remote_code=trust_remote_code, torch_dtype=new_confg.torch_dtype + ) + else: + model = AutoModelForCausalLM.from_config(new_confg, torch_dtype=new_confg.torch_dtype) + model.save_pretrained(output_path) + tokenizer.save_pretrained(output_path) + new_confg.save_pretrained(output_path) + print(f"Random model initialized and saved to {output_path}") + + +if __name__ == "__main__": + args = _init_args() + check_output_path(args.output_path) + init_random_model( + hf_model_path=args.hf_model_path, + new_config_path=args.new_config_path, + output_path=args.output_path, + trust_remote_code=args.trust_remote_code, + ) diff --git a/verl/scripts/install_sglang_mcore_npu.sh b/verl/scripts/install_sglang_mcore_npu.sh new file mode 100644 index 0000000000000000000000000000000000000000..86678faed562d490d65ecb55c8e520db06810f2d --- /dev/null +++ b/verl/scripts/install_sglang_mcore_npu.sh @@ -0,0 +1,58 @@ +#!/bin/bash +set -e +NPU_DEVICE=${NPU_DEVICE:=A3} +USE_MEGATRON=${USE_MEGATRON:-1} + +export MAX_JOBS=32 + +echo "1. install SGLang from source" +git clone -b v0.5.8 https://github.com/sgl-project/sglang.git +cd sglang +mv python/pyproject_other.toml python/pyproject.toml +pip install -e python[srt_npu] +cd .. + +echo "2. install torch & torch_npu & triton_ascend & other basic packages" +pip install torch==2.7.1 torch_npu==2.7.1.post2 torchvision==0.22.1 +pip install pybind11 click==8.2.1 mbridge "numpy<2.0.0" cachetools + + +echo "3. install sgl-kernel-npu form source, detailed readme in https://github.com/sgl-project/sgl-kernel-npu/blob/main/python/deep_ep/README.md" +git clone https://github.com/sgl-project/sgl-kernel-npu.git +cd sgl-kernel-npu +git checkout 46b73de +sed -i '101s/^/# /' build.sh +if [ "$NPU_DEVICE" = "A3" ]; then + bash build.sh +fi +if [ "$NPU_DEVICE" = "A2" ]; then + bash build.sh -a deepep2 +fi +pip install output/torch_memory_saver*.whl +pip install output/sgl_kernel_npu*.whl +pip install output/deep_ep*.whl +cd "$(pip show deep-ep | grep -E '^Location:' | awk '{print $2}')" && ln -s deep_ep/deep_ep_cpp*.so && cd - +python -c "import deep_ep; print(deep_ep.__path__)" +cd .. +# install sgl-kernel-npu from release whl +# if [ "$NPU_DEVICE" = "A3" ]; then +# wget https://github.com/sgl-project/sgl-kernel-npu/releases/download/2026.01.21/sgl-kernel-npu_2026.01.21_8.5.0_a3.zip +# fi +# if [ "$NPU_DEVICE" = "A2" ]; then +# wget https://github.com/sgl-project/sgl-kernel-npu/releases/download/2026.01.21/sgl-kernel-npu_2026.01.21_8.5.0_910b.zip +# fi +# unzip sgl-kernel-npu*.zip +# pip install output/torch_memory_saver*.whl +# pip install output/sgl_kernel_npu*.whl +# pip install output/deep_ep*.whl + +if [ $USE_MEGATRON -eq 1 ]; then + echo "4. install Megatron and MindSpeed" + git clone -b 2.3.0_core_r0.12.1 https://gitcode.com/Ascend/MindSpeed.git + pip install -e MindSpeed + pip install git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.12.1 +fi + +echo "5. May need to uninstall timm & triton" +pip uninstall -y timm triton +echo "Successfully installed all packages" diff --git a/verl/scripts/install_vllm_sglang_mcore.sh b/verl/scripts/install_vllm_sglang_mcore.sh new file mode 100644 index 0000000000000000000000000000000000000000..4ac686764744b29ba20b0e5798170816f54ed868 --- /dev/null +++ b/verl/scripts/install_vllm_sglang_mcore.sh @@ -0,0 +1,55 @@ +#!/bin/bash + +USE_MEGATRON=${USE_MEGATRON:-1} +USE_SGLANG=${USE_SGLANG:-1} + +export MAX_JOBS=32 + +echo "1. install inference frameworks and pytorch they need" +if [ $USE_SGLANG -eq 1 ]; then + pip install "sglang[all]==0.5.2" --no-cache-dir && pip install torch-memory-saver --no-cache-dir +fi +pip install --no-cache-dir "vllm==0.11.0" + +echo "2. install basic packages" +pip install "transformers[hf_xet]>=4.51.0" accelerate datasets peft hf-transfer \ + "numpy<2.0.0" "pyarrow>=15.0.0" pandas "tensordict>=0.8.0,<=0.10.0,!=0.9.0" torchdata \ + ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler \ + pytest py-spy pre-commit ruff tensorboard + +echo "pyext is lack of maintainace and cannot work with python 3.12." +echo "if you need it for prime code rewarding, please install using patched fork:" +echo "pip install git+https://github.com/ShaohonChen/PyExt.git@py311support" + +pip install "nvidia-ml-py>=12.560.30" "fastapi[standard]>=0.115.0" "optree>=0.13.0" "pydantic>=2.9" "grpcio>=1.62.1" + + +echo "3. install FlashAttention and FlashInfer" +# Install flash-attn-2.8.1 (cxx11abi=False) +wget -nv https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.1/flash_attn-2.8.1+cu12torch2.8cxx11abiFALSE-cp312-cp312-linux_x86_64.whl && \ + pip install --no-cache-dir flash_attn-2.8.1+cu12torch2.8cxx11abiFALSE-cp312-cp312-linux_x86_64.whl + +pip install --no-cache-dir flashinfer-python==0.3.1 + + +if [ $USE_MEGATRON -eq 1 ]; then + echo "4. install TransformerEngine and Megatron" + echo "Notice that TransformerEngine installation can take very long time, please be patient" + pip install "onnxscript==0.3.1" + NVTE_FRAMEWORK=pytorch pip3 install --no-deps git+https://github.com/NVIDIA/TransformerEngine.git@v2.6 + pip3 install --no-deps git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.13.1 +fi + + +echo "5. May need to fix opencv" +pip install opencv-python +pip install opencv-fixer && \ + python -c "from opencv_fixer import AutoFix; AutoFix()" + + +if [ $USE_MEGATRON -eq 1 ]; then + echo "6. Install cudnn python package (avoid being overridden)" + pip install nvidia-cudnn-cu12==9.10.2.21 +fi + +echo "Successfully installed all packages" diff --git a/verl/scripts/legacy_model_merger.py b/verl/scripts/legacy_model_merger.py new file mode 100644 index 0000000000000000000000000000000000000000..f1f20d7ea9a59ec1d7ff51bc2555a208966e8710 --- /dev/null +++ b/verl/scripts/legacy_model_merger.py @@ -0,0 +1,803 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +This script is used to merge huggingface model and test verl checkpoints from FSDP and Megatron backends. + +To merge FSDP checkpoints: +```sh +python scripts/legacy_model_merger.py merge \ + --backend fsdp \ + --local_dir checkpoints/verl_fsdp_gsm8k_examples/qwen2_5_0b5_fsdp_saveload/global_step_1/actor \ + --target_dir /path/to/merged_hf_model +``` + +To merge Megatron checkpoints: +```sh +python scripts/legacy_model_merger.py merge \ + --backend megatron \ + --tie-word-embedding \ + --local_dir checkpoints/verl_megatron_gsm8k_examples/qwen2_5_0b5_megatron_saveload/global_step_1/actor \ + --target_dir /path/to/merged_hf_model +``` + +For more details, please refer to documentation: +https://verl.readthedocs.io/en/latest/advance/checkpoint.html#convert-fsdp-and-megatron-checkpoints-to-huggingface-format-model +""" + +import argparse +import os +import re +import warnings +from abc import ABC, abstractmethod +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from pathlib import Path +from typing import Optional, Union + +import numpy as np +import torch +from accelerate import init_empty_weights +from safetensors.torch import load_file +from torch.distributed._tensor import Placement, Shard +from transformers import (AutoConfig, AutoModelForCausalLM, + AutoModelForTokenClassification, GenerationConfig, + PretrainedConfig) + +try: + # for torch 2.5+ + from torch.distributed.tensor import DTensor +except ImportError: + from torch.distributed._tensor import DTensor + +from tqdm import tqdm + +from verl.utils import hf_processor, hf_tokenizer +from verl.utils.transformers_compat import get_auto_model_for_vision2seq + +AutoModelForVision2Seq = get_auto_model_for_vision2seq() + + +@dataclass +class ModelMergerConfig: + operation: str # 'merge' or 'test' + backend: str + local_dir: str + hf_model_config_path: str + target_dir: Optional[str] = "tmp" + hf_upload_path: Optional[str] = None + private: bool = False + test_hf_dir: Optional[str] = None + tie_word_embedding: bool = False + is_value_model: bool = False + hf_model_path: Optional[str] = None + hf_upload: bool = field(init=False) + + def __post_init__(self): + self.hf_upload = self.operation == "merge" and bool(self.hf_upload_path) + if self.operation == "test": + self.target_dir = None + self.hf_upload_path = None + self.private = False + + +class BaseModelMerger(ABC): + def __init__(self, config: ModelMergerConfig): + self.config = config + self.hf_model_config_path = config.hf_model_config_path + + if config.hf_model_path: + print( + "Warning: --hf_model_path is deprecated and will be removed in a future version. Currently verl will save huggingface model configuration files into checkpoint directories. Therefore, there is no need to provide --hf_model_path. " + ) + self.hf_model_config_path = config.hf_model_path + + # Auto-detect huggingface subdirectory if it exists + huggingface_subdir = os.path.join(self.hf_model_config_path, "huggingface") + if os.path.isdir(huggingface_subdir): + self.hf_model_config_path = huggingface_subdir + + self.model_config = AutoConfig.from_pretrained(self.hf_model_config_path) + + def get_transformers_auto_model_class(self): + # Handle case where architectures might be None or empty + if self.model_config.architectures is None or len(self.model_config.architectures) == 0: + # Try to infer from model_type if architectures is missing + model_type = getattr(self.model_config, 'model_type', '').lower() + if 'vision' in model_type or 'vl' in model_type: + return AutoModelForVision2Seq + elif 'causal' in model_type or 'gpt' in model_type or 'llama' in model_type or 'qwen' in model_type: + return AutoModelForCausalLM + else: + raise NotImplementedError( + f"Cannot determine model class: architectures is None and model_type '{model_type}' is not recognized" + ) + + architecture = self.model_config.architectures[0] + if "ForTokenClassification" in architecture: + return AutoModelForTokenClassification + elif "ForCausalLM" in architecture: + return AutoModelForCausalLM + elif "ForConditionalGeneration" in architecture: + return AutoModelForVision2Seq + + raise NotImplementedError(f"Unknown architecture {self.model_config.architectures}") + + def patch_model_generation_config(self, model): + """ + The generation_config created from model config may be different to the pretrained model, + this may lead to error when generating: https://github.com/verl-project/verl/issues/1246 + + This function patch the generation_config created from model config to the pretrained model. + """ + if model.can_generate(): + try: + model.generation_config = GenerationConfig.from_pretrained(self.hf_model_config_path) + except OSError: + print( + f"Warning: Generation config file not found in {self.hf_model_config_path}, using a generation config created from the model config." + ) + return model + + def save_lora_adapter(self, state_dict: dict[str, torch.Tensor]): + """ + Save lora adapter to safetensors. + + Returns: + lora_path: str, the path to the lora adapter. None if no lora adapter found. + + Note: + This function change the 'state_dict' in place. + """ + lora_params_names = [name for name in state_dict.keys() if "lora_" in name] + + if len(lora_params_names) == 0: + return None + + import json + from typing import OrderedDict + + import peft + from safetensors.torch import save_file + + lora_params = OrderedDict() + target_modules = set() + lora_key = None + + for name in lora_params_names: + lora_key = name.replace(".default.weight", ".weight") + target_modules.add(lora_key.split(".")[-3]) + lora_params[lora_key] = state_dict.pop(name) + + lora_rank = min(lora_params[lora_key].shape[0], lora_params[lora_key].shape[1]) + peft_dict = { + "r": lora_rank, + "lora_alpha": 0, # lora_alpha is not set. An error should be raised to inform the user to set it manually. + "target_modules": list(target_modules), + } + peft_config = peft.LoraConfig(**peft_dict).to_dict() + peft_config["task_type"] = peft_config["task_type"].value if peft_config["task_type"] else None + peft_config["peft_type"] = peft_config["peft_type"].value if peft_config["peft_type"] else None + peft_config["target_modules"] = list(peft_config["target_modules"]) + + lora_path = os.path.join(self.config.target_dir, "lora_adapter") + os.makedirs(lora_path, exist_ok=True) + with open(os.path.join(lora_path, "adapter_config.json"), "w", encoding="utf-8") as f: + json.dump(peft_config, f, ensure_ascii=False, indent=4) + save_file(lora_params, os.path.join(lora_path, "adapter_model.safetensors")) + + for name in list(state_dict.keys()): + key = ( + name.replace("base_model.model.", "") + .replace(".base_layer.weight", ".weight") + .replace(".base_layer.bias", ".bias") + ) + state_dict[key] = state_dict.pop(name) + + return lora_path + + def save_hf_model_and_tokenizer(self, state_dict: dict[str, torch.Tensor]): + auto_model_class = self.get_transformers_auto_model_class() + with init_empty_weights(): + model = auto_model_class.from_config(self.model_config, torch_dtype=torch.bfloat16) + model.to_empty(device="cpu") + model = self.patch_model_generation_config(model) + + lora_path = self.save_lora_adapter(state_dict) + if lora_path: + print(f"Saving lora adapter to {lora_path}") + + print(f"Saving model to {self.config.target_dir}") + model.save_pretrained(self.config.target_dir, state_dict=state_dict) + del state_dict + del model + + processor = hf_processor(self.hf_model_config_path) + try: + tokenizer = hf_tokenizer(self.hf_model_config_path) + except Exception as e: + warnings.warn(f"Failed to create tokenizer: {e}. This may affect tokenizer saving", stacklevel=1) + tokenizer = None + if processor is not None: + print(f"Saving processor to {self.config.target_dir}") + processor.save_pretrained(self.config.target_dir) + if tokenizer is not None: + print(f"Saving tokenizer to {self.config.target_dir}") + tokenizer.save_pretrained(self.config.target_dir) + + def upload_to_huggingface(self): + from huggingface_hub import HfApi + + api = HfApi() + api.create_repo(repo_id=self.config.hf_upload_path, private=self.config.private, exist_ok=True) + api.upload_folder(folder_path=self.config.target_dir, repo_id=self.config.hf_upload_path, repo_type="model") + + @abstractmethod + def merge_and_save(self): + raise NotImplementedError("Subclasses should implement this method") + + +class FSDPModelMerger(BaseModelMerger): + def _get_world_size(self) -> int: + """Extracts the FSDP world_size from checkpoint filenames (e.g., 'model_world_size_8_rank_0.pt').""" + for filename in os.listdir(self.config.local_dir): + match = re.match(r"model_world_size_(\d+)_rank_0\.pt", filename) + if match: + return int(match.group(1)) + raise FileNotFoundError( + f"Could not determine world size. No file matching 'model_world_size_(\\d+)_rank_0.pt' found in {self.config.local_dir}" + ) + + def _load_rank_zero_state_dict(self, world_size: int) -> dict: + return torch.load( + Path(self.config.local_dir) / f"model_world_size_{world_size}_rank_0.pt", + map_location="cpu", + weights_only=False, + ) + + def _extract_device_mesh_info(self, state_dict: dict, world_size: int) -> tuple[np.ndarray, tuple[str, ...]]: + """ + Retrieves sharding information (device_mesh, mesh_dim_names) from a DTensor in the state_dict. + If no DTensor is found, infers a simple FSDP mesh based on world_size. + """ + pivot_key = sorted(list(state_dict.keys()))[0] + weight = state_dict[pivot_key] + + if isinstance(weight, DTensor): + # get sharding info + device_mesh = weight.device_mesh + mesh = device_mesh.mesh + mesh_dim_names = device_mesh.mesh_dim_names + else: + # for non-DTensor + mesh = np.array([world_size], dtype=np.int64) + mesh_dim_names = ("fsdp",) + + return mesh, mesh_dim_names + + def _calculate_shard_configuration( + self, mesh: np.ndarray, mesh_dim_names: tuple[str, ...] + ) -> tuple[int, tuple[int, ...]]: + """Calculates the total number of shards and the shape of the device mesh.""" + assert mesh_dim_names in (("fsdp",), ("ddp", "fsdp")), f"Unsupported mesh_dim_names {mesh_dim_names}" + + if "tp" in mesh_dim_names: + # TODO: "tp" is not supported yet due to the above assert + total_shards = mesh.shape[-1] * mesh.shape[-2] + mesh_shape = (mesh.shape[-2], mesh.shape[-1]) + else: + total_shards = mesh.shape[-1] + mesh_shape = (mesh.shape[-1],) + + return total_shards, mesh_shape + + def _merge_by_placement(self, tensors: list[torch.Tensor], placement: Placement) -> torch.Tensor: + """Merges a list of tensors based on their DTensor placement""" + if placement.is_replicate(): + return tensors[0] + elif placement.is_partial(): + raise NotImplementedError("Partial placement is not supported yet") + elif placement.is_shard(): + return torch.cat(tensors, dim=placement.dim).contiguous() + + raise NotImplementedError(f"Unsupported placement: {placement}") + + def _load_and_merge_state_dicts( + self, world_size: int, total_shards: int, mesh_shape: tuple[int, ...], mesh_dim_names: tuple[str, ...] + ) -> dict[str, torch.Tensor]: + model_state_dict_lst = [None] * total_shards + + def process_one_shard(rank: int, model_state_dict_lst: list): + model_path = Path(self.config.local_dir) / f"model_world_size_{world_size}_rank_{rank}.pt" + state_dict = torch.load(model_path, map_location="cpu", weights_only=False) + model_state_dict_lst[rank] = state_dict + return state_dict + + with ThreadPoolExecutor(max_workers=min(32, os.cpu_count())) as executor: + futures = [executor.submit(process_one_shard, rank, model_state_dict_lst) for rank in range(total_shards)] + for future in tqdm(futures, desc=f"Loading {total_shards} FSDP shards", total=total_shards): + future.result() + + # Merge state dicts from all shards + state_dict = {} + param_placements: dict[str, list] = {} + + for key in set(model_state_dict_lst[0].keys()): + state_dict[key] = [] + for model_state_shard in model_state_dict_lst: + # add tensor shard in order of rank to state_dict[key] + tensor = model_state_shard.pop(key) + if isinstance(tensor, DTensor): + state_dict[key].append(tensor._local_tensor.bfloat16()) + + placements = tuple(tensor.placements) + # replicated placement at dp dimension can be discarded + if mesh_dim_names[0] in ("dp", "ddp"): + placements = placements[1:] + + if key not in param_placements: + param_placements[key] = placements + else: + assert param_placements[key] == placements + else: + state_dict[key].append(tensor.bfloat16()) + + del model_state_dict_lst + + # Merge tensors + for key in sorted(state_dict): + if not isinstance(state_dict[key], list): + print(f"No need to merge key {key}") + continue + if key in param_placements: + # merge shards + placements: tuple[Shard] = param_placements[key] + if len(mesh_shape) == 1: + # 1-D list, FSDP without TP + assert len(placements) == 1 + shards = state_dict[key] + state_dict[key] = self._merge_by_placement(shards, placements[0]) + else: + # 2-D list, FSDP + TP + raise NotImplementedError("FSDP + TP is not supported yet") + else: + state_dict[key] = torch.cat(state_dict[key], dim=0) + + return state_dict + + def merge_and_save(self): + world_size = self._get_world_size() + rank_zero_state_dict = self._load_rank_zero_state_dict(world_size) + + mesh, mesh_dim_names = self._extract_device_mesh_info(rank_zero_state_dict, world_size) + print(f"Got device mesh {mesh}, mesh_dim_names {mesh_dim_names}") + + total_shards, mesh_shape = self._calculate_shard_configuration(mesh, mesh_dim_names) + print(f"Processing model shards with {total_shards} {mesh_shape} in total") + + merged_state_dict = self._load_and_merge_state_dicts(world_size, total_shards, mesh_shape, mesh_dim_names) + + if self.config.operation == "test": + if not self.config.test_hf_dir: + raise ValueError("test_hf_dir must be provided for test operation") + self._test_state_dict(merged_state_dict) + elif self.config.operation == "merge": + self.save_hf_model_and_tokenizer(merged_state_dict) + if self.config.hf_upload: + self.upload_to_huggingface() + else: + raise ValueError(f"Unknown operation: {self.config.operation}") + + def _test_state_dict(self, state_dict: dict[str, torch.Tensor]): + auto_model_class = self.get_transformers_auto_model_class() + + hf_model = auto_model_class.from_pretrained(self.config.test_hf_dir, torch_dtype=torch.bfloat16) + hf_state_dict = hf_model.state_dict() + del hf_model + + hf_model_keys = set(hf_state_dict.keys()) + collected_keys = set(state_dict.keys()) + + missing_keys = hf_model_keys - collected_keys + assert len(missing_keys) == 0, f"Missing keys in collected state dict: {list(sorted(missing_keys))}" + + extra_keys = collected_keys - hf_model_keys + assert len(extra_keys) == 0, f"Extra keys in collected state dict: {list(sorted(extra_keys))}" + + for key in hf_model_keys: + hf_shape = hf_state_dict[key].shape + collected_shape = state_dict[key].shape + assert hf_shape == collected_shape, ( + f"Shape mismatch for key '{key}': original {hf_shape} vs collected {collected_shape}" + ) + + hf_dtype = hf_state_dict[key].dtype + collected_dtype = state_dict[key].dtype + assert hf_dtype == collected_dtype, ( + f"Dtype mismatch for key '{key}': original {hf_dtype} vs collected {collected_dtype}" + ) + + torch.testing.assert_close(hf_state_dict[key], state_dict[key], atol=1e-6, rtol=1e-6) + + print("FSDP checks passed: The merged state_dict matches the hf model saved by FSDPCheckpointManager.") + + +class MegatronModelMerger(BaseModelMerger): + def __init__(self, config: ModelMergerConfig): + from verl.utils.megatron_utils import \ + get_hf_config_and_tokenizer_checkpoint_path + + config.hf_model_config_path = get_hf_config_and_tokenizer_checkpoint_path(config.local_dir) + super().__init__(config) + + self.params_mapping = { + # megatron core gpt model name, huggingface model name + # NOTICE: It's a little bit tricky, when 2 keys have the same prefix, we need to make sure the longer key within the containing relationship is processed first. + "embedding.word_embeddings": "model.embed_tokens", + # attn + "self_attention.linear_qkv.layer_norm_weight": "input_layernorm.weight", + "self_attention.linear_qkv.layer_norm_bias": "input_layernorm.bias", + "self_attention.linear_qkv": "self_attn.qkv_proj", + "self_attention.q_layernorm": "self_attn.q_norm", + "self_attention.k_layernorm": "self_attn.k_norm", + "self_attention.linear_proj": "self_attn.o_proj", + # mla + "self_attention.linear_q_proj": "self_attn.q_proj", + "self_attention.linear_q_down_proj": "self_attn.q_a_proj", + "self_attention.linear_q_up_proj.layer_norm_weight": "self_attn.q_a_layernorm.weight", + "self_attention.linear_q_up_proj": "self_attn.q_b_proj", + "self_attention.linear_kv_down_proj": "self_attn.kv_a_proj_with_mqa", + "self_attention.linear_kv_up_proj.layer_norm_weight": "self_attn.kv_a_layernorm.weight", + "self_attention.linear_kv_up_proj": "self_attn.kv_b_proj", + # mlp + "pre_mlp_layernorm": "post_attention_layernorm", + "mlp.linear_fc1.layer_norm_weight": "post_attention_layernorm.weight", + "mlp.linear_fc1.layer_norm_bias": "post_attention_layernorm.bias", + "mlp.linear_fc1": "mlp.gate_up_proj", + "mlp.linear_fc2": "mlp.down_proj", + # moe + "mlp.router.expert_bias": "mlp.gate.e_score_correction_bias", + "mlp.router": "mlp.gate", + "mlp.shared_experts.linear_fc1": "mlp.shared_experts.gate_up_proj", + "mlp.shared_experts.linear_fc2": "mlp.shared_experts.down_proj", + "linear_fc1": "gate_up_proj", + "linear_fc2": "down_proj", + # output + "final_layernorm": "norm", + "output_layer": "lm_head", + } + + def _get_tp_pp_rank_from_sharded_dir(self, sharded_dir: str) -> tuple[int, int]: + tp_rank = pp_rank = None + rank_list = sharded_dir.split("_")[2:] + if re.match(r"mp_rank_(\d\d)_(\d\d\d)", sharded_dir): + tp_rank = int(rank_list[0]) + pp_rank = int(rank_list[1]) + elif re.match(r"mp_rank_(\d\d)", sharded_dir): + tp_rank = int(rank_list[0]) + pp_rank = 0 + + assert tp_rank is not None and pp_rank is not None, f"Invalid sharded dir {sharded_dir}" + + return tp_rank, pp_rank + + def _check_megatron_checkpoint_path(self, model_path: str) -> tuple[list[str], int, int]: + """ + Validates the Megatron checkpoint structure (presence of 'model.pt' in sharded directories). + Determines TP and PP sizes from directory names. + """ + tp_size = 0 + pp_size = 0 + sharded_dirs = sorted(os.listdir(model_path)) + for sharded_dir in sharded_dirs: + assert "model.pt" in os.listdir(Path(model_path) / sharded_dir), f"model.pt not found in {sharded_dir}" + tp_rank, pp_rank = self._get_tp_pp_rank_from_sharded_dir(sharded_dir) + tp_size = max(tp_size, tp_rank + 1) + pp_size = max(pp_size, pp_rank + 1) + return sharded_dirs, tp_size, pp_size + + def _merge_across_tp( + self, + key: str, + tp_data: list[torch.Tensor], + config: PretrainedConfig, + tp_size: int, + is_value_model: bool = False, + ) -> Union[torch.Tensor, list[torch.Tensor]]: + if "linear_fc1.weight" in key: + # if the tensor is gate and proj + gate_lst = [] + up_lst = [] + for infer_param in tp_data: + gate, up = infer_param.chunk(2) + gate_lst.append(gate) + up_lst.append(up) + gate = torch.cat(gate_lst, dim=0) + up = torch.cat(up_lst, dim=0) + return [gate, up] + elif "self_attention.linear_qkv." in key and "layer_norm" not in key: + # if the tensor is qkv, for each param on tp, split into q, k, v + # concat q, k, v separately. + q_lst = [] + k_lst = [] + v_lst = [] + assert config.num_attention_heads % config.num_key_value_heads == 0 + num_q_per_kv = config.num_attention_heads // config.num_key_value_heads + assert tp_data[0].shape[0] % (num_q_per_kv + 2) == 0 + kv_size_per_tp = tp_data[0].shape[0] // (num_q_per_kv + 2) + split_size = [kv_size_per_tp * num_q_per_kv, kv_size_per_tp, kv_size_per_tp] + + for infer_param in tp_data: + num_query_groups_per_partition = config.num_key_value_heads // tp_size + for chunk in infer_param.chunk(num_query_groups_per_partition): + split_size = [ + kv_size_per_tp * num_q_per_kv // num_query_groups_per_partition, + kv_size_per_tp // num_query_groups_per_partition, + kv_size_per_tp // num_query_groups_per_partition, + ] + q, k, v = chunk.split(split_size) + q_lst.append(q) + k_lst.append(k) + v_lst.append(v) + + q = torch.cat(q_lst, dim=0) + k = torch.cat(k_lst, dim=0) + v = torch.cat(v_lst, dim=0) + return [q, k, v] + elif "layer_norm" in key or "layernorm" in key or "router" in key or ("output_layer" in key and is_value_model): + return tp_data[0] + else: + dim = 0 + if "linear_fc2.weight" in key or "self_attention.linear_proj" in key: + dim = 1 + return torch.cat(tp_data, dim=dim) + + def _load_state_dicts( + self, model_ckpt_path: str, sharded_dirs: list[str], tp_size: int, pp_size: int + ) -> list[list[dict]]: + model_state_dict_lst = [[None for _ in range(tp_size)] for _ in range(pp_size)] + + def _process_one_megatron_shard(sharded_dir: str): + model_file_path = Path(model_ckpt_path) / sharded_dir / "model.pt" + state_dict = torch.load(model_file_path, map_location="cpu", weights_only=False) + tp_rank, pp_rank = self._get_tp_pp_rank_from_sharded_dir(sharded_dir) + model_state_dict_lst[pp_rank][tp_rank] = state_dict + + with ThreadPoolExecutor(max_workers=min(32, os.cpu_count())) as executor: + futures = [executor.submit(_process_one_megatron_shard, sharded_dir) for sharded_dir in sharded_dirs] + for future in tqdm(futures, desc=f"Loading {len(sharded_dirs)} Megatron shards", total=len(sharded_dirs)): + future.result() + + return model_state_dict_lst + + def _check_megatron_state_key(self, key: str) -> bool: + """ + Checks if the key is a valid Megatron state key. + + Now the model merger only supports keys that start with "decoder/embedding/output_layer" in TransformerLayer. + Shall not use key starts with "model." + """ + if key.startswith("model."): + raise ValueError( + f"Invalid key {key} in Megatron state_dict. Expected keys to start with 'decoder/embedding/output_layer' in TransformerLayer." + ) + + skip_checking_keys = ["embedding.word_embeddings", "output_layer"] + for skip_key in skip_checking_keys: + if skip_key in key: + print(f"skip checking key {key}") + return + + # Exclude extra state keys + if not key.startswith("decoder"): + raise ValueError( + f"Invalid key {key} in Megatron state_dict. Expected keys to start with 'decoder' in TransformerLayer." + ) + + def _merge_state_dicts( + self, model_state_dict_lst: list[list[dict]], tp_size: int, pp_size: int + ) -> dict[str, torch.Tensor]: + state_dict = {} + vpp_size = len(model_state_dict_lst[0][0]) + layers_cum = 0 + + for vpp_rank in range(vpp_size): + for pp_rank in range(pp_size): + layers_handled = 0 + keys = model_state_dict_lst[pp_rank][0][vpp_rank].keys() + for key in keys: + if "extra_state" in key: + continue + if self.config.tie_word_embedding and ("output_layer" in key): + print("skip lm_head and reward_head loading because of tie_word_embeddings") + continue + + self._check_megatron_state_key(key) + hf_name = self._replace_name(key, self.params_mapping) + assert hf_name is not None, f"Failed to convert layer name [{key}] from megatron to huggingface." + if "model.layers." in hf_name: + local_layer_no = int(hf_name.split(".")[2]) + layers_handled = max(local_layer_no, layers_handled) + global_layer_no = local_layer_no + layers_cum + new_key_list = hf_name.split(".") + new_key_list[2] = str(global_layer_no) + hf_name = ".".join(new_key_list) + else: + warnings.warn(f"hf_name {hf_name} will not be fixed with layer number", stacklevel=2) + + tp_data = [model_state_dict_lst[pp_rank][tp_rank][vpp_rank][key] for tp_rank in range(tp_size)] + merged = self._merge_across_tp(key, tp_data, self.model_config, tp_size, self.config.is_value_model) + + if not isinstance(merged, list): + state_dict[hf_name] = merged + elif len(merged) == 3: + # split qkv + for n, d in zip(["q", "k", "v"], merged): + state_dict[hf_name.replace("qkv", n)] = d + elif len(merged) == 2: + # split gate up + state_dict[hf_name.replace("gate_up", "gate")] = merged[0] + state_dict[hf_name.replace("gate_up", "up")] = merged[1] + print( + f"converted {key} to {hf_name} with shape {merged.shape if isinstance(merged, torch.Tensor) else [t.shape for t in merged]}" + ) + + layers_cum += layers_handled + 1 # zero based + + return state_dict + + def merge_and_save(self): + from verl.utils.megatron_utils import get_model_checkpoint_path + + model_ckpt_path = get_model_checkpoint_path(self.config.local_dir) + sharded_dirs, tp_size, pp_size = self._check_megatron_checkpoint_path(model_ckpt_path) + print(f"sharded_dirs: {sharded_dirs}, tp_size: {tp_size}, pp_size: {pp_size}, mp_size: {len(sharded_dirs)}") + + model_state_dict_lst = self._load_state_dicts(model_ckpt_path, sharded_dirs, tp_size, pp_size) + merged_state_dict = self._merge_state_dicts(model_state_dict_lst, tp_size, pp_size) + del model_state_dict_lst + + if self.config.operation == "test": + if not self.config.test_hf_dir: + raise ValueError("test_hf_dir must be provided for test operation") + self._test_state_dict(merged_state_dict) + elif self.config.operation == "merge": + self.save_hf_model_and_tokenizer(merged_state_dict) + if self.config.hf_upload: + self.upload_to_huggingface() + else: + raise ValueError(f"Unknown operation: {self.config.operation}") + + def _test_state_dict(self, state_dict: dict[str, torch.Tensor]): + """ + Compares the merged Megatron state_dict against a reference safetensors model. + Applies necessary name mappings from Megatron to Hugging Face conventions using _replace_name. + """ + ref_state_dict = load_file(Path(self.config.test_hf_dir) / "model.safetensors") + + for name, loaded_weight in state_dict.items(): + # name = self._replace_name(original_name, self.params_mapping) + if not name or name.endswith(".bias") and name not in ref_state_dict: + continue + if "rotary_emb.inv_freq" in name: + continue + if self.config.tie_word_embedding and "lm_head.weight" in name: + continue + if name not in ref_state_dict: + raise RuntimeError(f"key: {name} not exist in state_dict") + param = ref_state_dict[name] + assert loaded_weight.dtype == param.dtype + torch.testing.assert_close(loaded_weight, param, atol=1e-2, rtol=5e-2) + + def _replace_name(self, megatron_name: str, name_mapping: dict[str, str]) -> str: + for m_name, v_name in name_mapping.items(): + if m_name not in megatron_name: + continue + + megatron_name = megatron_name.replace("decoder", "model") + param_name = megatron_name.replace(m_name, v_name) + return param_name + + return None # Return None if no mapping found + + +def main(): + parser = argparse.ArgumentParser(description="verl model merger") + subparsers = parser.add_subparsers(dest="operation", required=True, help="Specify 'merge' or 'test' operation.") + + base_op_parser = argparse.ArgumentParser(add_help=False) + base_op_parser.add_argument( + "--backend", type=str, required=True, choices=["fsdp", "megatron"], help="The backend of the model" + ) + base_op_parser.add_argument("--local_dir", type=str, required=True, help="Path to the saved model checkpoints") + base_op_parser.add_argument( + "--hf_model_path", + type=str, + default=None, + help="(Deprecated) Path to the original Hugging Face model for config.", + ) + base_op_parser.add_argument( + "--tie-word-embedding", + action="store_true", + help="Whether to tie word embedding weights (currently only Megatron supported)", + ) + base_op_parser.add_argument( + "--is-value-model", + action="store_true", + help="Whether the model is a value model (currently only Megatron supported)", + ) + + merge_parser = subparsers.add_parser("merge", parents=[base_op_parser], help="Merge model checkpoints and save.") + merge_parser.add_argument( + "--target_dir", default="tmp", type=str, help="Directory to save the merged huggingface model" + ) + merge_parser.add_argument( + "--hf_upload_path", default=None, type=str, help="Hugging Face repository ID to upload the model" + ) + merge_parser.add_argument( + "--private", action="store_true", help="Whether to upload the model to a private Hugging Face repository" + ) + + test_parser = subparsers.add_parser( + "test", parents=[base_op_parser], help="Test merged model against a reference Hugging Face model" + ) + test_parser.add_argument( + "--test_hf_dir", type=str, required=True, help="Path to the reference Hugging Face model directory for testing" + ) + + args = parser.parse_args() + + common_config_args = { + "operation": args.operation, + "backend": args.backend, + "tie_word_embedding": args.tie_word_embedding, + "is_value_model": args.is_value_model, + "local_dir": args.local_dir, + "hf_model_path": args.hf_model_path, + "hf_model_config_path": args.local_dir, + } + + if args.operation == "merge": + config = ModelMergerConfig( + **common_config_args, + target_dir=args.target_dir, + hf_upload_path=args.hf_upload_path, + private=args.private, + test_hf_dir=None, + ) + os.makedirs(config.target_dir, exist_ok=True) + elif args.operation == "test": + config = ModelMergerConfig( + **common_config_args, + test_hf_dir=args.test_hf_dir, + # the following args are not used by test operation + target_dir=None, + hf_upload_path=None, + private=False, + ) + else: + raise NotImplementedError(f"Unknown operation: {args.operation}") + + if config.backend == "fsdp": + merger = FSDPModelMerger(config) + elif config.backend == "megatron": + merger = MegatronModelMerger(config) + else: + raise NotImplementedError(f"Unknown backend: {config.backend}") + + merger.merge_and_save() + + +if __name__ == "__main__": + main() diff --git a/verl/scripts/megatron_merge_lora.py b/verl/scripts/megatron_merge_lora.py new file mode 100644 index 0000000000000000000000000000000000000000..619d162c24fc90d27b32a2ee48f5ce8b858a4729 --- /dev/null +++ b/verl/scripts/megatron_merge_lora.py @@ -0,0 +1,134 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""LoRA merger for the Megatron backend, built on the unified model engine. + +The legacy :mod:`verl.workers.megatron_workers` module was removed in favour of +:mod:`verl.workers.engine_workers`. This script ports the old ``CustomSaveWorker`` +onto the new ``ActorRolloutRefWorker`` / ``TrainingWorker`` stack while keeping +the same CLI contract: reuse the training Hydra config, supply an +``actor_rollout_ref.model.lora.adapter_path``, and the merged HuggingFace +weights are written next to the adapter checkpoint. +""" + +import os +from pprint import pprint + +import hydra +import ray +import torch +from omegaconf import OmegaConf + +from verl.single_controller.base.decorator import Dispatch, register +from verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup +from verl.utils.megatron_utils import get_hf_model_checkpoint_path +from verl.workers.engine_workers import ActorRolloutRefWorker + +os.environ["NCCL_DEBUG"] = "WARN" +os.environ["TOKENIZERS_PARALLELISM"] = "true" + + +class CustomSaveWorker(ActorRolloutRefWorker): + """Extends the unified :class:`ActorRolloutRefWorker` with a merge RPC. + + The actor ``TrainingWorker`` built by ``init_model`` already loads the base + HuggingFace weights via the Megatron bridge and applies the LoRA adapter + from ``config.model.lora.adapter_path`` (see + :func:`verl.utils.megatron_utils.make_megatron_module`). Here we just need + to dump the merged weights back to HuggingFace format. + """ + + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def save_merged_weights(self, hf_ckpt_path): + engine = self.actor.engine + + # Reload model parameters to the accelerator if they were offloaded. + if engine.is_param_offload_enabled: + self.actor.to(device="device", model=True, optimizer=False, grad=False) + + torch.distributed.barrier() + + print(f"[Rank {os.environ.get('RANK', '?')}] Saving merged weights to {hf_ckpt_path}...") + + # ``MegatronEngine`` exposes the bridge that was used to load the base + # weights; reuse it to export merged HF-format weights. + if engine.vanilla_bridge: + engine.bridge.save_weights(engine.module, hf_ckpt_path, distributed_filesystem=True, memory_efficient=True) + else: + engine.bridge.save_hf_weights(engine.module, hf_ckpt_path) + + return True + + +@hydra.main(config_path="../verl/trainer/config", config_name="ppo_megatron_trainer", version_base=None) +def main(config): + assert config.actor_rollout_ref.model.lora.adapter_path is not None, "adapter_path must be specified" + + if ( + config.actor_rollout_ref.actor.optim.lr_decay_steps is None + or config.actor_rollout_ref.actor.optim.lr_decay_steps < 1 + ): + # set to bypass OptimizerParamScheduler checks + config.actor_rollout_ref.actor.optim.lr_decay_steps = 100000 + + run_merge(config) + + +def run_merge(config) -> None: + if not ray.is_initialized(): + # this is for local ray cluster + default_runtime_env = {"env_vars": {"TOKENIZERS_PARALLELISM": "true", "NCCL_DEBUG": "WARN"}} + ray_init_kwargs = config.ray_kwargs.get("ray_init", {}) + runtime_env_kwargs = ray_init_kwargs.get("runtime_env", {}) + runtime_env = OmegaConf.merge(default_runtime_env, runtime_env_kwargs) + ray_init_kwargs = OmegaConf.create({**ray_init_kwargs, "runtime_env": runtime_env}) + print(f"ray init kwargs: {ray_init_kwargs}") + ray.init(**OmegaConf.to_container(ray_init_kwargs)) + + ray.get(main_task.remote(config)) + + +@ray.remote(num_cpus=1) +def main_task(config): + pprint(OmegaConf.to_container(config, resolve=True)) # resolve=True will eval symbol values + OmegaConf.resolve(config) + + ray_cls_with_init = RayClassWithInitArgs( + cls=ray.remote(CustomSaveWorker), config=config.actor_rollout_ref, role="actor" + ) + resource_pool = RayResourcePool(process_on_nodes=[config.trainer.n_gpus_per_node] * config.trainer.nnodes) + + worker = RayWorkerGroup( + resource_pool=resource_pool, + ray_cls_with_init=ray_cls_with_init, + device_name=config.trainer.device, + ) + worker.init_model() + + adapter_path = config.actor_rollout_ref.model.lora.adapter_path + hf_ckpt_path = get_hf_model_checkpoint_path(os.path.dirname(adapter_path)) + worker.save_merged_weights(hf_ckpt_path) + + +if __name__ == "__main__": + """ + Use the same config as your training script, besides **specifying the adapter_path**. + + For example, your training script starts with: + `python3 -m verl.trainer.main_ppo --config-name=ppo_megatron_trainer ...` + Now replace it with + `python3 ./scripts/megatron_merge_lora.py --config-name=ppo_megatron_trainer ...` + """ + main() diff --git a/verl/scripts/print_cfg.py b/verl/scripts/print_cfg.py new file mode 100644 index 0000000000000000000000000000000000000000..287756fb1b7dbaac84b5f7ec572ba7a172e347b3 --- /dev/null +++ b/verl/scripts/print_cfg.py @@ -0,0 +1,35 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +try: + import hydra +except ImportError as e: + raise ImportError("Please install hydra-core via 'pip install hydra-core' and retry.") from e + + +@hydra.main(config_path="../verl/trainer/config", config_name="ppo_trainer", version_base=None) +def main(config): + """Main entry point for PPO training with Hydra configuration management. + + Args: + config_dict: Hydra configuration dictionary containing training parameters. + """ + print(config) + from verl.utils.config import omega_conf_to_dataclass + + profiler_config = omega_conf_to_dataclass(config.critic.profiler) + print(profiler_config) + + +if __name__ == "__main__": + main() diff --git a/verl/scripts/rollout_viewer.py b/verl/scripts/rollout_viewer.py new file mode 100644 index 0000000000000000000000000000000000000000..eb0314edc56f7f3ad0ce719b5ee0ffb4f241d4d3 --- /dev/null +++ b/verl/scripts/rollout_viewer.py @@ -0,0 +1,565 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +import re +import traceback +from pathlib import Path +from typing import Annotated, Optional + +import aiofiles + +try: + import ujson as json +except ImportError: + import json +import typer +from rich.highlighter import ReprHighlighter +from rich.markdown import Markdown +from rich.table import Table +from rich.text import Text +from textual import on +from textual.app import App, ComposeResult +from textual.containers import Horizontal, Vertical, VerticalScroll +from textual.widgets import Input, ProgressBar, Select, SelectionList, Static + +INDEX_KEY = "__IDX" +FILE_SUFFIX = ".jsonl" + + +def check_textual_version(): + # check if textual version is equal to 0.52.1 + import textual + from packaging.version import Version + + if Version(textual.__version__) != Version("0.52.1"): + raise ImportError(f"Textual version {textual.__version__} is not supported, please pip install textual==0.52.1") + + +check_textual_version() + + +async def load_path(p: Path, data: dict, mask_strs: str, idx: int, pbar): + samples = [] + async with aiofiles.open(p, encoding="utf-8") as f: + async for line in f: + d = json.loads(line) + for k in d: + if isinstance(d[k], str): + if mask_strs: + d[k] = re.sub(rf"{mask_strs}", "*", d[k]) + else: + d[k] = json.dumps(d[k], ensure_ascii=False, indent=4) + + d[INDEX_KEY] = len(samples) + samples.append(d) + data[idx] = {"samples": samples} + + print(f"path {p} loaded") + pbar.advance(1) + + +async def load_dir(path: Path, data: dict[int, dict], pbar, mask_strs: str = ""): + paths = list(path.glob(f"*{FILE_SUFFIX}")) + paths = sorted(paths, key=lambda x: int(x.stem)) + + tasks = [load_path(p, data, mask_strs, i, pbar) for i, p in enumerate(paths)] + + await asyncio.gather(*tasks) + + +class Highlighter(ReprHighlighter): + highlights = ReprHighlighter.highlights + [ + r"(?P[][\<\>{}()\|（）【】\[\]=`])", + r"\<\|(?P[\w\W]*?)\|\>", + ] + + +def center_word_with_equals_exactly(word: str, total_length: int, char: str = "=") -> str: + if len(word) > total_length: + return word + + padding = total_length - len(word) + left_pad = (padding) // 2 + right_pad = (padding + 1) // 2 + return char * left_pad + " " + word + " " + char * right_pad + + +def highlight_keyword(content: str, keyword: Optional[str]): + if not keyword: + return Text(content) + text = Text() + parts = content.split(keyword) + for i, part in enumerate(parts): + text.append(part, style=None) + if i < len(parts) - 1: + # text.append(keyword, style=Style(color="#d154d1", bgcolor="yellow", bold=True)) + text.append(keyword, style="on #8f51b5") + return text + + +help_doc = """ +⌨️ keybinds： + +- `f/esc`: find/cancel +- `tab/←/→`: change focus +- `j/k`: page down/up +- `g/G`: scroll home/end +- `n/N`: next sample/step +- `p/P`: previous sample/step +- `s`: switch display mode + - plain text + - rich table + +""" + + +class JsonLineViewer(App): + BINDINGS = [ + ("left", "focus_previous", "Focus Previous"), + ("right", "focus_next", "Focus Next"), + ("s", "swith_render", "switch render"), + # control + ("n", "next_sample", "Next Sample"), + ("N", "next_step", "Next Step"), + ("p", "previous_sample", "Previous Sample"), + ("P", "previous_step", "Previous Step"), + # search + ("f", "toggle_search", "find"), + ("enter", "next_search", "find next"), + ("escape", "cancel_search", "cancel find"), + # scroll + ("j", "page_down", "page down"), + ("k", "page_up", "page up"), + ("g", "page_home", "page home"), + ("G", "page_end", "page end"), + ] + + CSS = """ + + Select:focus > SelectCurrent { + border: tall #8f51b5; + } + Select.-expanded > SelectCurrent { + border: tall #8f51b5; + } + #select-container { + width: 15%; + height: 100%; + align: center top; + } + #search-container { + height: 10%; + align: center top; + } + #search-box { + width: 50%; + } + #reqid-box { + width: 50%; + } + """ + + def __init__(self, step_num: int, data: dict[int, dict], pbar): + super().__init__() + self.step_num = step_num + + self.data = data + self.render_table = False + self.selected_step_index = 0 + self.selected_sample_index = 0 + self.pbar = pbar + + self.matches = [] + self.current_match_index = 0 + + self.highlighter = Highlighter() + + first_samples = data[list(data.keys())[0]]["samples"] + # Prepare the initial field filter list (all keys from the first sample) + self.filter_fields = [(f, f, True) for f in first_samples[0].keys()] + + # Internal set used for fast membership checks when we add new fields on the fly. + # We keep it here so that when new columns appear in later steps (e.g. `request_id`), + # they can be added to the UI automatically without restarting the viewer. + self._field_set: set[str] = set(first_samples[0].keys()) + self.sample_num = len(first_samples) + + def compose(self) -> ComposeResult: + with Horizontal(id="search-container"): + yield Input(placeholder="find something...", id="search-box") + yield Input(placeholder="request id...", id="reqid-box") + with Vertical(id="search-container2"): + yield self.pbar + yield Static("", id="search-status") + + with Horizontal(): + with Vertical(id="select-container"): + yield Static("\n") + yield Static( + renderable=Markdown( + help_doc, + ), + markup=False, + ) + yield Static("\n") + yield Select( + id="step-select", + value=0, + prompt="select step", + options=[("step: 1", 0)], + allow_blank=False, + ) + yield Select( + id="sample-select", + value=0, + prompt="select sample", + options=[("sample: 1", 0)], + allow_blank=False, + ) + yield Select( + id="sample-sort", + value=0, + prompt="排序", + options=[ + ("sort", 0), + ("score asc", 1), + ("score desc", 2), + ], + allow_blank=False, + ) + + yield SelectionList[int](("Select ALL", 1, True), id="fields-select-all") + with VerticalScroll(id="scroll-view2"): + yield SelectionList[str](*self.filter_fields, id="fields-select") + with VerticalScroll(id="scroll-view"): + yield Static(id="content", markup=False) + + async def on_mount(self) -> None: + self.step_select = self.query_one("#step-select", Select) + self.sample_select = self.query_one("#sample-select", Select) + self.sample_sort = self.query_one("#sample-sort", Select) + self.content_display = self.query_one("#content", Static) + self.search_box = self.query_one("#search-box", Input) + self.reqid_box = self.query_one("#reqid-box", Input) + self.scroll_view = self.query_one("#scroll-view", VerticalScroll) + self.search_status = self.query_one("#search-status", Static) + self.fields_select = self.query_one("#fields-select", SelectionList) + self.fields_select.border_title = "field filter" + + if self.data: + self.step_select.set_options([(f"step: {i + 1}", i) for i in range(self.step_num)]) + self.sample_select.set_options([(f"sample: {i + 1}", i) for i in range(self.sample_num)]) + self.step_select.focus() + await self.update_content() + + def update_result_options(self, offset: int = 0, sort_desc: Optional[bool] = None): + options = [] + if isinstance(self.selected_step_index, int) and self.selected_step_index < len(self.data): + if self.sample_num is None or sort_desc is not None: + samples = self.data[self.selected_step_index].get("samples", []) + if not samples: + self.selected_sample_index = offset + return + if sort_desc is not None: + samples = sorted( + samples, + key=lambda x: x.get("score", x.get("score_1", 0)), + reverse=sort_desc, + ) + + options = [(f"sample: {r[INDEX_KEY] + 1}", r[INDEX_KEY]) for r in samples] + self.sample_select.set_options(options) + self.sample_num = len(samples) + + if sort_desc is not None and options: + self.selected_sample_index = options[0][1] + else: + self.selected_sample_index = offset + + async def update_content(self, search_keyword: Optional[str] = None): + content = "" + try: + samples = self.data[self.selected_step_index].get("samples", []) + content_dict_full = samples[self.selected_sample_index] + + # Dynamically track any NEW keys that appear and add them to the field filter. + self._update_fields_select(content_dict_full.keys()) + + # Apply field selection filter (only show selected fields) + content_dict = {k: v for k, v in content_dict_full.items() if k in self.fields_select.selected} + if self.render_table: + content = Table("key", "value", show_lines=True) + for k in content_dict: + v = content_dict[k] + v = f"{v}" + content.add_row( + k, + self.highlighter(highlight_keyword(v, search_keyword)), + ) + else: + text = Text() + for k in content_dict: + v = content_dict[k] + s = center_word_with_equals_exactly(k, 64) + f"\n{v}\n" + text.append(highlight_keyword(s, search_keyword)) + content = self.highlighter(text) + except KeyError: + content = f"Loading data asynchronously, progress: {len(self.data)}/{self.step_num} step" + + except Exception: + content = self.highlighter(traceback.format_exc()) + + self.content_display.update(content) + + # --------------------------------------------------------------------- + # Request-ID jump logic + # --------------------------------------------------------------------- + + @on(Input.Submitted, "#reqid-box") + async def on_reqid_submitted(self, event: Input.Submitted) -> None: + """Jump to the sample that has a matching `request_id`.""" + + req_id_raw = event.value.strip() + # Remove hyphens so search is tolerant to different id formats + req_id = req_id_raw.replace("-", "") + if not req_id: + return + + found = False + for step_idx, step_data in self.data.items(): + for sample in step_data.get("samples", []): + sample_id = str(sample.get("request_id", "")) + if sample_id.replace("-", "") == req_id: + # Update selected indices + self.selected_step_index = step_idx + self.step_select.value = step_idx + + # Ensure sample list is updated and select sample + self.update_result_options(offset=sample[INDEX_KEY]) + self.selected_sample_index = sample[INDEX_KEY] + self.sample_select.value = sample[INDEX_KEY] + + await self._clear_search() + await self.update_content() + + found = True + break + if found: + break + + if not found: + self.search_status.update(Text(f"request_id '{req_id_raw}' not found", style="bold red")) + else: + # Keep the typed id in the input box so users see what was searched. + pass + + # --------------------------------------------------------------------- + # Helper: add new fields to SelectionList on-the-fly + # --------------------------------------------------------------------- + + def _update_fields_select(self, keys): + """Add any unseen *keys* to the field-selection widget so they can be toggled. + + The viewer is often launched with only the first step loaded. Later steps may + introduce new columns (e.g. `request_id`). This helper ensures those fields + become visible without requiring a restart. + """ + # Ensure we have the widget (only after on_mount) + if not hasattr(self, "fields_select"): + return + + for k in keys: + if k not in self._field_set: + self._field_set.add(k) + try: + # By default, new fields are selected so they appear immediately. + self.fields_select.add_option(k, k, selected=True) + except Exception: + # Fallback for older textual versions where signature is different. + self.fields_select.add_option((k, k, True)) + + @on(Select.Changed, "#step-select") + async def step_changed(self, event): + self.selected_step_index = event.value + self.update_result_options() + await self.update_content() + + @on(Select.Changed, "#sample-select") + async def sample_changed(self, event): + self.selected_sample_index = event.value + await self._clear_search() + await self.update_content() + + @on(Select.Changed, "#sample-sort") + async def sort_changed(self, event): + v = event.value + self.update_result_options(sort_desc=None if v == 0 else False if v == 1 else True) + await self.update_content() + + @on(SelectionList.SelectedChanged, "#fields-select") + async def fields_changed(self, event): + await self.update_content() + + @on(SelectionList.SelectedChanged, "#fields-select-all") + async def fields_all_changed(self, event): + s = self.query_one("#fields-select-all", SelectionList) + if s.selected: + self.fields_select.select_all() + else: + self.fields_select.deselect_all() + + def action_focus_previous(self): + self.screen.focus_previous() + + def action_focus_next(self): + self.screen.focus_next() + + async def action_next_step(self) -> None: + self.selected_step_index += 1 + if self.selected_step_index >= self.step_num: + self.selected_step_index = 0 + self.step_select.value = self.selected_step_index + self.update_result_options() + await self.update_content() + + async def action_next_sample(self) -> None: + self.selected_sample_index += 1 + if not self.sample_num or self.selected_sample_index >= self.sample_num: + self.selected_sample_index = 0 + self.sample_select.value = self.selected_sample_index + await self._clear_search() + await self.update_content() + + async def action_previous_step(self) -> None: + self.selected_step_index -= 1 + if self.selected_step_index < 0: + self.selected_step_index = self.step_num - 1 + self.step_select.value = self.selected_step_index + self.update_result_options() + await self.update_content() + + async def action_previous_sample(self) -> None: + self.selected_sample_index -= 1 + if self.selected_sample_index < 0: + self.selected_sample_index = self.sample_num - 1 + self.sample_select.value = self.selected_sample_index + await self._clear_search() + await self.update_content() + + async def action_swith_render(self): + self.render_table = not self.render_table + await self.update_content() + + def action_toggle_search(self) -> None: + self.search_box.focus() + + async def action_cancel_search(self) -> None: + self.search_box.value = "" + await self._clear_search() + await self.update_content() + + async def _clear_search(self): + self.matches = [] + self.search_status.update("") + self.current_match_index = 0 + + @on(Input.Submitted, "#search-box") + async def on_search_submitted(self, event: Input.Submitted) -> None: + self.matches = [] + self.current_match_index = 0 + if event.value: + await self.update_content(event.value) + renderable = self.content_display.render() + if isinstance(renderable, Table): + return + + assert isinstance(renderable, Text) + console = self.content_display._console + lines = renderable.wrap(console, self.scroll_view.container_size.width) + line_idx_recorded = set() + for line_idx, line in enumerate(lines): + if line_idx in line_idx_recorded: + continue + if event.value in line: + self.matches.append( + { + "line": line_idx, + "word": event.value, + } + ) + line_idx_recorded.add(line_idx) + self.scroll_view.focus() + await self.action_next_search() + + async def action_next_search(self) -> None: + if not self.matches or self.current_match_index >= len(self.matches): + return + + target_line = self.matches[self.current_match_index]["line"] + self.scroll_view.scroll_to(x=0, y=target_line * 1, animate=False) + self.current_match_index = (self.current_match_index + 1) % len(self.matches) + self.search_status.update( + Text( + f"Find ：{self.current_match_index + 1}/{len(self.matches)}", + style="bold on #8f51b5", + ) + ) + + def action_page_up(self): + self.scroll_view.scroll_page_up(animate=False) + + def action_page_down(self): + self.scroll_view.scroll_page_down(animate=False) + + def action_page_home(self): + self.scroll_view.scroll_home(animate=False) + + def action_page_end(self): + self.scroll_view.scroll_end(animate=False) + + +async def _run(path: Path, mask_str: str): + assert path.exists(), f"{path} not exist" + + paths = list(path.glob(f"*{FILE_SUFFIX}")) + paths = sorted(paths, key=lambda x: int(x.stem)) + + if not paths: + raise ValueError(f"no available reward dump files under f{path}") + + print(f"get jsonl file nums: {len(paths)}") + + pbar = ProgressBar(total=len(paths), name="data load progress") + data = {} + await load_path(paths[0], data, mask_str, 0, pbar) + app = JsonLineViewer(step_num=len(paths), data=data, pbar=pbar) + await asyncio.gather(load_dir(path, data, pbar, mask_str), app.run_async()) + + +app = typer.Typer() + + +@app.command(help="launch TUI APP") +def run( + rollout_data_dir: Path, + mask_str: Annotated[str, typer.Option(help="string that will be masked to *")] = r"<\|image_pad\|>|<\|imgpad\|>", +): + loop = asyncio.get_event_loop() + loop.run_until_complete(_run(rollout_data_dir, mask_str)) + + +if __name__ == "__main__": + app() diff --git a/verl/scripts/veomni/moe_merge.py b/verl/scripts/veomni/moe_merge.py new file mode 100644 index 0000000000000000000000000000000000000000..aa1c57d42e4de0a0665bc39466c8af498a0ef36a --- /dev/null +++ b/verl/scripts/veomni/moe_merge.py @@ -0,0 +1,121 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Merge individual MoE expert weights into stacked tensors for efficient loading. + +This script takes a HuggingFace checkpoint with individual expert weights +(e.g., model.layers.{i}.mlp.experts.{j}.gate_proj.weight) and merges them +into stacked tensors (e.g., model.layers.{i}.mlp.experts.gate_proj) for +faster loading and better memory efficiency in VeOmni. + +The merging process: +1. Loads individual expert weights from the HF checkpoint +2. Stacks them into single tensors for each projection type +3. Handles all three projection types: gate_proj, up_proj, down_proj +4. Supports both Qwen3-MoE (num_experts) and DeepSeek (n_routed_experts) formats +5. Handles models with initial dense layers (first_k_dense_replace) + +Usage: python moe_merge.py --raw_hf_path --merge_hf_path +""" + +import os +from argparse import ArgumentParser +from dataclasses import dataclass +from glob import glob +from typing import Generator + +import torch +from safetensors.torch import safe_open +from tqdm import tqdm +from transformers import AutoConfig +from veomni.models import build_tokenizer, save_model_weights + + +@dataclass +class StateDictIterator: + filepath: str + + def __iter__(self) -> Generator[tuple[str, "torch.Tensor"], None, None]: + if self.filepath.endswith(".safetensors"): + with safe_open(self.filepath, framework="pt", device="cpu") as f: + for key in f.keys(): + yield key, f.get_tensor(key) + + else: + state_dict = torch.load(self.filepath, map_location="cpu", weights_only=True, mmap=True) + for key in state_dict.keys(): + yield key, state_dict[key] + + +def main(raw_hf_path, merge_hf_path): + torch.set_default_dtype(torch.bfloat16) + os.makedirs(merge_hf_path, exist_ok=True) + + config = AutoConfig.from_pretrained(raw_hf_path) + tokenizer = build_tokenizer(raw_hf_path) + + safetensor_files = list(glob(os.path.join(raw_hf_path, "*.safetensors"))) + safetensor_files.sort() + state_dict_iterators = [StateDictIterator(shard_file) for shard_file in safetensor_files] + new_state_dict = {} + for state_dict_iterator in tqdm(state_dict_iterators, desc="Loading checkpoint shards"): + for name, tensor in state_dict_iterator: + new_state_dict[name] = tensor.cpu() + + print(new_state_dict.keys()) + + if hasattr(config, "num_experts"): + # qwen3moe + num_experts = config.num_experts + elif hasattr(config, "n_routed_experts"): + # deepseek + num_experts = config.n_routed_experts + else: + raise RuntimeError("could not find how many experts to assign") + num_hidden_layers = config.num_hidden_layers + + if hasattr(config, "first_k_dense_replace"): + # deepseek first k dense layer + moe_layer_start_idx = config.first_k_dense_replace + else: + # moe layer only in the model + moe_layer_start_idx = 0 + + for i in range(moe_layer_start_idx, num_hidden_layers): + gate_proj = [] + for j in range(num_experts): + gate_proj.append(new_state_dict.pop(f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight")) + + new_state_dict[f"model.layers.{i}.mlp.experts.gate_proj"] = torch.stack(gate_proj) + up_proj = [] + for j in range(num_experts): + up_proj.append(new_state_dict.pop(f"model.layers.{i}.mlp.experts.{j}.up_proj.weight")) + + new_state_dict[f"model.layers.{i}.mlp.experts.up_proj"] = torch.stack(up_proj) + down_proj = [] + for j in range(num_experts): + down_proj.append(new_state_dict.pop(f"model.layers.{i}.mlp.experts.{j}.down_proj.weight")) + + new_state_dict[f"model.layers.{i}.mlp.experts.down_proj"] = torch.stack(down_proj) + + model_assets = [config, tokenizer] + save_model_weights(merge_hf_path, new_state_dict, model_assets=model_assets) + + +if __name__ == "__main__": + parser = ArgumentParser() + parser.add_argument("--raw_hf_path", type=str, required=True) + parser.add_argument("--merge_hf_path", type=str, required=True) + args = parser.parse_args() + main(args.raw_hf_path, args.merge_hf_path) diff --git a/verl/scripts/veomni/moe_split.py b/verl/scripts/veomni/moe_split.py new file mode 100644 index 0000000000000000000000000000000000000000..f38a990466e87eb34aa68eaca71d8b2a38cb3ba4 --- /dev/null +++ b/verl/scripts/veomni/moe_split.py @@ -0,0 +1,96 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Reverse process of moe_merge.py - splits merged MoE expert weights back to individual experts. + +This script takes a HF checkpoint that has been processed by moe_merge.py (where expert weights +are stacked into single tensors) and splits them back to the original format with individual +expert weights. + +The process reverses the merging by: +1. Loading stacked tensors like model.layers.{i}.mlp.experts.gate_proj +2. Unstacking them back to individual experts model.layers.{i}.mlp.experts.{j}.gate_proj.weight +3. Handling all three projection types: gate_proj, up_proj, down_proj + +Usage: python moe_split.py --merge_hf_path --split_hf_path +""" + +import os +from argparse import ArgumentParser +from dataclasses import dataclass +from glob import glob +from typing import Generator + +import torch +from safetensors.torch import safe_open +from tqdm import tqdm +from transformers import AutoConfig +from veomni.models import build_tokenizer, save_model_weights + + +@dataclass +class StateDictIterator: + filepath: str + + def __iter__(self) -> Generator[tuple[str, "torch.Tensor"], None, None]: + if self.filepath.endswith(".safetensors"): + with safe_open(self.filepath, framework="pt", device="cpu") as f: + for key in f.keys(): + yield key, f.get_tensor(key) + + else: + state_dict = torch.load(self.filepath, map_location="cpu", weights_only=True, mmap=True) + for key in state_dict.keys(): + yield key, state_dict[key] + + +def main(merge_hf_path, split_hf_path): + torch.set_default_dtype(torch.bfloat16) + os.makedirs(split_hf_path, exist_ok=True) + + config = AutoConfig.from_pretrained(merge_hf_path) + tokenizer = build_tokenizer(merge_hf_path) + + safetensor_files = list(glob(os.path.join(merge_hf_path, "*.safetensors"))) + safetensor_files.sort() + state_dict_iterators = [StateDictIterator(shard_file) for shard_file in safetensor_files] + new_state_dict = {} + for state_dict_iterator in tqdm(state_dict_iterators, desc="Loading checkpoint shards"): + for name, tensor in state_dict_iterator: + new_state_dict[name] = tensor.cpu() + + num_experts = config.num_experts + num_hidden_layers = config.num_hidden_layers + for i in range(num_hidden_layers): + print(f"Converting layer {i}") + for proj_name in ["gate_proj", "up_proj", "down_proj"]: + stacked_key = f"model.layers.{i}.mlp.experts.{proj_name}" + if stacked_key in new_state_dict: + stacked_tensor = new_state_dict.pop(stacked_key) + for j in range(num_experts): + expert_key = f"model.layers.{i}.mlp.experts.{j}.{proj_name}.weight" + new_state_dict[expert_key] = stacked_tensor[j] + + model_assets = [config, tokenizer] + + print("Saving to safetensors") + save_model_weights(split_hf_path, new_state_dict, model_assets=model_assets) + + +if __name__ == "__main__": + parser = ArgumentParser() + parser.add_argument("--merge_hf_path", type=str, required=True) + parser.add_argument("--split_hf_path", type=str, required=True) + args = parser.parse_args() + main(args.merge_hf_path, args.split_hf_path) diff --git a/verl/setup.py b/verl/setup.py new file mode 100644 index 0000000000000000000000000000000000000000..7671e1b3f10dca8ef76ee47dbdf90a2858d8506a --- /dev/null +++ b/verl/setup.py @@ -0,0 +1,102 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# setup.py is the fallback installation script when pyproject.toml does not work +import os +from pathlib import Path + +from setuptools import find_packages, setup + +version_folder = os.path.dirname(os.path.join(os.path.abspath(__file__))) + +with open(os.path.join(version_folder, "verl/version/version")) as f: + __version__ = f.read().strip() + +install_requires = [ + "accelerate", + "codetiming", + "datasets", + "dill", + "hydra-core", + "numpy<2.0.0", + "pandas", + "peft", + "pyarrow>=19.0.0", + "pybind11", + "pylatexenc", + "ray[default]>=2.41.0", + "torchdata", + "tensordict>=0.8.0,<=0.10.0,!=0.9.0", + "transformers", + "wandb", + "packaging>=20.0", + "tensorboard", +] + +TEST_REQUIRES = ["pytest", "pre-commit", "py-spy", "pytest-asyncio", "pytest-rerunfailures"] +PRIME_REQUIRES = ["pyext"] +GEO_REQUIRES = ["mathruler", "torchvision", "qwen_vl_utils"] +GPU_REQUIRES = ["liger-kernel", "flash-attn"] +MATH_REQUIRES = ["math-verify"] # Add math-verify as an optional dependency +VLLM_REQUIRES = ["tensordict>=0.8.0,<=0.10.0,!=0.9.0", "vllm>=0.8.5,<=0.12.0"] +TRTLLM_REQUIRES = ["tensorrt-llm>=1.2.0rc6"] +SGLANG_REQUIRES = [ + "tensordict>=0.8.0,<=0.10.0,!=0.9.0", + "sglang[srt,openai]==0.5.8", + "torch==2.9.1", +] +TRL_REQUIRES = ["trl<=0.9.6"] +MCORE_REQUIRES = ["mbridge"] + +extras_require = { + "test": TEST_REQUIRES, + "prime": PRIME_REQUIRES, + "geo": GEO_REQUIRES, + "gpu": GPU_REQUIRES, + "math": MATH_REQUIRES, + "vllm": VLLM_REQUIRES, + "sglang": SGLANG_REQUIRES, + "trl": TRL_REQUIRES, + "mcore": MCORE_REQUIRES, + "trtllm": TRTLLM_REQUIRES, +} + + +this_directory = Path(__file__).parent +long_description = (this_directory / "README.md").read_text() + +setup( + name="verl", + version=__version__, + package_dir={"": "."}, + packages=find_packages(where="."), + url="https://github.com/verl-project/verl", + license="Apache 2.0", + author="Bytedance - Seed - MLSys", + author_email="zhangchi.usc1992@bytedance.com, gmsheng@connect.hku.hk", + description="verl: Volcano Engine Reinforcement Learning for LLM", + install_requires=install_requires, + extras_require=extras_require, + package_data={ + "": ["version/*"], + "verl": [ + "trainer/config/*.yaml", + "trainer/config/*/*.yaml", + "experimental/*/config/*.yaml", + ], + }, + include_package_data=True, + long_description=long_description, + long_description_content_type="text/markdown", +) diff --git a/verl/tests/README.md b/verl/tests/README.md new file mode 100644 index 0000000000000000000000000000000000000000..479f06933e4e536ee159b738794daa05364119bb --- /dev/null +++ b/verl/tests/README.md @@ -0,0 +1,30 @@ +# Tests layout + +Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance: +- `tests/trainer` for testing functionality related to `verl/trainer` +- `tests/models` for testing functionality related to `verl/models` +- ... + +There are a few folders with `special_` prefix, created for special purposes: +- `special_distributed`: unit tests that must run with multiple GPUs +- `special_e2e`: end-to-end tests with training/generation scripts +- `special_npu`: tests for NPUs +- `special_sanity`: a suite of quick sanity tests +- `special_standalone`: a set of test that are designed to run in dedicated environments + +Accelerators for tests +- By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`. +- For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment. + +# Workflow layout + +All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs: +1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml` +2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml` +3. End-to-end tests: `e2e_*.yml` +4. Unit tests + - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py` + - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix. + - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when + - new workflow yaml is added to `.github/workflows` + - new tests are added to workflow mentioned in 2. \ No newline at end of file diff --git a/verl/tests/__init__.py b/verl/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1ce90c5eb352d85c59105c0dc85b5f1dd576f095 --- /dev/null +++ b/verl/tests/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/verl/tests/checkpoint_engine/__init__.py b/verl/tests/checkpoint_engine/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1cd1e8433dffa0b3ba420be3e346f4f5cd062014 --- /dev/null +++ b/verl/tests/checkpoint_engine/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/verl/tests/checkpoint_engine/test_correctness_on_gpu.py b/verl/tests/checkpoint_engine/test_correctness_on_gpu.py new file mode 100644 index 0000000000000000000000000000000000000000..33218ebfb3d130b8d7d535f33460da854ee7f762 --- /dev/null +++ b/verl/tests/checkpoint_engine/test_correctness_on_gpu.py @@ -0,0 +1,200 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os + +import pytest +import ray +import torch + +from tests.checkpoint_engine.test_utils import create_rollout_worker_group, create_trainer_worker_group +from verl.checkpoint_engine import CheckpointEngineManager +from verl.single_controller.ray.base import ( + RayResourcePool, + split_resource_pool, +) +from verl.utils.device import get_device_name +from verl.utils.ray_utils import auto_await +from verl.workers.config import CheckpointEngineConfig, HFModelConfig, RolloutConfig + +_ngpus = torch.cuda.device_count() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("rebuild_group", [False, True]) +@pytest.mark.parametrize("num_trainer, num_rollout", [(2, _ngpus - 2)]) +@auto_await +async def test_nccl_checkpoint_engine( + rebuild_group, + num_trainer, + num_rollout, + num_nodes=1, + num_gpus_per_node=_ngpus, + bucket_size_mb=128, + check_allclose=True, + model_path="~/models/Qwen/Qwen3-8B-Base", +): + model_path = os.path.expanduser(model_path) + ray.init( + runtime_env={ + "env_vars": { + "UCX_TLS": "rc,tcp,cuda", + "UCX_MAX_RNDV_RAILS": "4", + "UCX_LOG_LEVEL": "INFO", + "VERL_LOGGING_LEVEL": "DEBUG", + } + } + ) + + # initialize config + checkpoint_engine_config = CheckpointEngineConfig( + backend="nccl", + update_weights_bucket_megabytes=bucket_size_mb, + engine_kwargs={"nccl": {"rebuild_group": rebuild_group}}, + ) + model_config = HFModelConfig(path=model_path, use_remove_padding=True) + rollout_config = RolloutConfig(name="vllm", checkpoint_engine=checkpoint_engine_config) + + # create trainer and rollout worker group + resource_pool = RayResourcePool(process_on_nodes=[num_gpus_per_node] * num_nodes, max_colocate_count=3) + trainer_pool, rollout_pool = split_resource_pool(resource_pool, [num_trainer, num_rollout]) + trainer = create_trainer_worker_group(trainer_pool, model_config, checkpoint_engine_config) + trainer.reset() + rollout, replicas = await create_rollout_worker_group(rollout_pool, model_config, rollout_config, check_allclose) + + # create checkpoint engine manager + checkpoint_manager = CheckpointEngineManager(config=checkpoint_engine_config, trainer=trainer, replicas=replicas) + for _ in range(3): + await checkpoint_manager.update_weights() + rollout.check_weights() + + ray.shutdown() + + +@pytest.mark.skip(reason="temporary skip since our ci environment is not ready") +@pytest.mark.asyncio +@pytest.mark.parametrize("device", ["cuda", "cpu"]) +@pytest.mark.parametrize("num_trainer, num_rollout", [(2, 6)]) +@auto_await +async def test_nixl_checkpoint_engine( + num_trainer, + num_rollout, + device, + num_nodes=1, + num_gpus_per_node=8, + bucket_size_mb=128, + check_allclose=True, + model_path="~/models/Qwen/Qwen3-8B-Base", +): + model_path = os.path.expanduser(model_path) + ray.init( + runtime_env={ + "env_vars": { + # TODO: it's pretty hard to set these environment variables right, please consult + # with your network admin. Maybe auto adjust UCX_* according to NCCL_IB_*? + "UCX_TLS": "rc,ud,cuda", + # "UCX_IB_GID_INDEX": "3", # NCCL_IB_GID_INDEX + # "UCX_IB_DEVICES": "mlx5_1:1,mlx5_2:1,mlx5_3:1", # NCCL_IB_HCA + "UCX_RC_TIMEOUT": "30s", # NCCL_IB_TIMEOUT + "UCX_RC_RETRY_COUNT": "7", # NCCL_IB_RETRY_COUNT + "UCX_KEEPALIVE_INTERVAL": "1s", + "UCX_KEEPALIVE_NUM_EPS": "10", + "UCX_MAX_RNDV_RAILS": "4", + "UCX_IB_ROCE_REACHABILITY_MODE": "all", + "UCX_LOG_LEVEL": "INFO", + "VERL_LOGGING_LEVEL": "DEBUG", + } + } + ) + + # initialize config + checkpoint_engine_config = CheckpointEngineConfig( + backend="nixl", update_weights_bucket_megabytes=bucket_size_mb, engine_kwargs={"nixl": {"device": device}} + ) + model_config = HFModelConfig(path=model_path, use_remove_padding=True) + rollout_config = RolloutConfig(name="vllm", checkpoint_engine=checkpoint_engine_config) + + # create trainer and rollout worker group + resource_pool = RayResourcePool(process_on_nodes=[num_gpus_per_node] * num_nodes, max_colocate_count=3) + trainer_pool, rollout_pool = split_resource_pool(resource_pool, [num_trainer, num_rollout]) + trainer = create_trainer_worker_group(trainer_pool, model_config, checkpoint_engine_config) + trainer.reset() + rollout, replicas = await create_rollout_worker_group(rollout_pool, model_config, rollout_config, check_allclose) + + # create checkpoint engine manager + checkpoint_manager = CheckpointEngineManager(config=checkpoint_engine_config, trainer=trainer, replicas=replicas) + for _ in range(3): + await checkpoint_manager.update_weights() + rollout.check_weights() + + ray.shutdown() + + +@pytest.mark.skip(reason="temporary skip since our ci environment is not ready") +@pytest.mark.asyncio +@pytest.mark.parametrize("rebuild_group", [False]) +@pytest.mark.parametrize("num_trainer, num_rollout", [(2, 6)]) +@auto_await +async def test_kimi_checkpoint_engine( + rebuild_group, + num_trainer, + num_rollout, + num_nodes=1, + num_gpus_per_node=8, + check_allclose=True, + model_path="~/models/Qwen/Qwen3-8B-Base", +): + model_path = os.path.expanduser(model_path) + ray.init( + runtime_env={ + "env_vars": { + "NCCL_IB_HCA": "mlx5", + "VERL_LOGGING_LEVEL": "DEBUG", + } + } + ) + + # initialize config + checkpoint_engine_config = CheckpointEngineConfig( + backend="kimi_ckpt_engine", engine_kwargs={"kimi_ckpt_engine": {"rebuild_group": rebuild_group}} + ) + model_config = HFModelConfig(path=model_path, use_remove_padding=True) + rollout_config = RolloutConfig(name="vllm", checkpoint_engine=checkpoint_engine_config) + + # create trainer and rollout worker group + resource_pool = RayResourcePool(process_on_nodes=[num_gpus_per_node] * num_nodes, max_colocate_count=3) + resource_pool.get_placement_groups(device_name=get_device_name()) + trainer_pool, rollout_pool = split_resource_pool(resource_pool, [num_trainer, num_rollout]) + trainer = create_trainer_worker_group(trainer_pool, model_config, checkpoint_engine_config) + trainer.reset() + rollout, replicas = await create_rollout_worker_group(rollout_pool, model_config, rollout_config, check_allclose) + + # create checkpoint engine manager + checkpoint_manager = CheckpointEngineManager(config=checkpoint_engine_config, trainer=trainer, replicas=replicas) + for _ in range(3): + await checkpoint_manager.update_weights() + rollout.check_weights() + + ray.shutdown() + + +if __name__ == "__main__": + test_nccl_checkpoint_engine( + rebuild_group=False, + num_trainer=2, + num_rollout=30, + num_nodes=4, + num_gpus_per_node=8, + check_allclose=False, + model_path=os.environ["HDFS_ROOT"] + "/model/Qwen3-30B-A3B-Base", + ) diff --git a/verl/tests/checkpoint_engine/test_correctness_on_npu.py b/verl/tests/checkpoint_engine/test_correctness_on_npu.py new file mode 100644 index 0000000000000000000000000000000000000000..4981a2870b9d8c43c70e98344dc71c96ef62d16a --- /dev/null +++ b/verl/tests/checkpoint_engine/test_correctness_on_npu.py @@ -0,0 +1,183 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os + +import pytest +import ray + +from tests.checkpoint_engine.test_utils import create_rollout_worker_group, create_trainer_worker_group +from verl.checkpoint_engine import CheckpointEngineManager +from verl.single_controller.ray.base import ( + RayResourcePool, + split_resource_pool, +) +from verl.utils.device import get_device_name +from verl.utils.ray_utils import auto_await +from verl.workers.config import CheckpointEngineConfig, HFModelConfig, RolloutConfig + + +@pytest.mark.asyncio +@pytest.mark.parametrize("rebuild_group", [False]) +@pytest.mark.parametrize("num_trainer, num_rollout", [(2, 6)]) +@auto_await +async def test_hccl_checkpoint_engine( + rebuild_group, + num_trainer, + num_rollout, + num_nodes=1, + num_gpus_per_node=8, + check_allclose=True, + model_path="~/models/Qwen/Qwen3-8B-Base", +): + model_path = os.path.expanduser(model_path) + ray.init( + runtime_env={ + "env_vars": { + "HCCL_CONNECT_TIMEOUT": "1500", + "HCCL_HOST_SOCKET_PORT_RANGE": "60000-60050", + "HCCL_NPU_SOCKET_PORT_RANGE": "61000-61050", + "VERL_LOGGING_LEVEL": "DEBUG", + } + } + ) + + # initialize config + checkpoint_engine_config = CheckpointEngineConfig( + backend="nccl", engine_kwargs={"nccl": {"rebuild_group": rebuild_group}} + ) + model_config = HFModelConfig(path=model_path, use_remove_padding=True) + rollout_config = RolloutConfig(name="vllm", checkpoint_engine=checkpoint_engine_config) + + # create trainer and rollout worker group + resource_pool = RayResourcePool(process_on_nodes=[num_gpus_per_node] * num_nodes, max_colocate_count=3) + resource_pool.get_placement_groups(device_name=get_device_name()) + trainer_pool, rollout_pool = split_resource_pool(resource_pool, [num_trainer, num_rollout]) + trainer = create_trainer_worker_group(trainer_pool, model_config, checkpoint_engine_config) + trainer.reset() + rollout, replicas = await create_rollout_worker_group(rollout_pool, model_config, rollout_config, check_allclose) + + # create checkpoint engine manager + checkpoint_manager = CheckpointEngineManager(config=checkpoint_engine_config, trainer=trainer, replicas=replicas) + for _ in range(3): + await checkpoint_manager.update_weights() + rollout.check_weights() + + ray.shutdown() + + +@pytest.mark.skip(reason="temporary skip since our ci environment is not ready") +@pytest.mark.asyncio +@pytest.mark.parametrize("rebuild_group", [False]) +@pytest.mark.parametrize("num_trainer, num_rollout", [(4, 28)]) +async def test_kimi_checkpoint_engine( + rebuild_group, + num_trainer, + num_rollout, + num_nodes=2, + num_gpus_per_node=16, + check_allclose=True, + model_path="~/models/Qwen/Qwen3-32B", +): + model_path = os.path.expanduser(model_path) + ray.init( + runtime_env={ + "env_vars": { + "HCCL_CONNECT_TIMEOUT": "1500", + "VERL_LOGGING_LEVEL": "DEBUG", + } + } + ) + + # initialize config + checkpoint_engine_config = CheckpointEngineConfig( + backend="kimi_ckpt_engine", engine_kwargs={"kimi_ckpt_engine": {"rebuild_group": rebuild_group}} + ) + model_config = HFModelConfig(path=model_path, use_remove_padding=True) + rollout_config = RolloutConfig(name="vllm", checkpoint_engine=checkpoint_engine_config) + + # create trainer and rollout worker group + resource_pool = RayResourcePool(process_on_nodes=[num_gpus_per_node] * num_nodes, max_colocate_count=3) + resource_pool.get_placement_groups(device_name=get_device_name()) + trainer_pool, rollout_pool = split_resource_pool(resource_pool, [num_trainer, num_rollout]) + trainer = create_trainer_worker_group(trainer_pool, model_config, checkpoint_engine_config) + trainer.reset() + rollout, replicas = await create_rollout_worker_group(rollout_pool, model_config, rollout_config, check_allclose) + + # create checkpoint engine manager + checkpoint_manager = CheckpointEngineManager(config=checkpoint_engine_config, trainer=trainer, replicas=replicas) + for _ in range(3): + await checkpoint_manager.update_weights() + rollout.check_weights() + + ray.shutdown() + + +@pytest.mark.skip(reason="temporary skip since our ci environment is not ready") +@pytest.mark.asyncio +@pytest.mark.parametrize("device", ["npu"]) +@pytest.mark.parametrize("num_trainer, num_rollout", [(2, 6)]) +async def test_mooncake_checkpoint_engine( + rebuild_group, + num_trainer, + num_rollout, + device, + num_nodes=1, + num_gpus_per_node=8, + check_allclose=True, + model_path="~/models/Qwen/Qwen3-8B-Base", +): + model_path = os.path.expanduser(model_path) + ray.init( + runtime_env={ + "env_vars": { + "ASCEND_USE_SHORT_CONNECTION": "1", + "VERL_LOGGING_LEVEL": "DEBUG", + } + } + ) + + # initialize config + checkpoint_engine_config = CheckpointEngineConfig( + backend="mooncake", engine_kwargs={"mooncake": {"device": device, "rebuild_group": rebuild_group}} + ) + model_config = HFModelConfig(path=model_path, use_remove_padding=True) + rollout_config = RolloutConfig(name="vllm", checkpoint_engine=checkpoint_engine_config) + + # create trainer and rollout worker group + resource_pool = RayResourcePool(process_on_nodes=[num_gpus_per_node] * num_nodes, max_colocate_count=3) + resource_pool.get_placement_groups(device_name=get_device_name()) + trainer_pool, rollout_pool = split_resource_pool(resource_pool, [num_trainer, num_rollout]) + trainer = create_trainer_worker_group(trainer_pool, model_config, checkpoint_engine_config) + trainer.reset() + rollout, replicas = await create_rollout_worker_group(rollout_pool, model_config, rollout_config, check_allclose) + + # create checkpoint engine manager + checkpoint_manager = CheckpointEngineManager(config=checkpoint_engine_config, trainer=trainer, replicas=replicas) + for _ in range(3): + await checkpoint_manager.update_weights() + rollout.check_weights() + + ray.shutdown() + + +if __name__ == "__main__": + test_hccl_checkpoint_engine( + rebuild_group=False, + num_trainer=2, + num_rollout=6, + num_nodes=1, + num_gpus_per_node=8, + check_allclose=False, + model_path=os.environ["HDFS_ROOT"] + "/model/Qwen3-30B-A3B-Base", + ) diff --git a/verl/tests/checkpoint_engine/test_special_server_adapter.py b/verl/tests/checkpoint_engine/test_special_server_adapter.py new file mode 100644 index 0000000000000000000000000000000000000000..44b7d64d92e95ed12dd48dc458772e26192aaef3 --- /dev/null +++ b/verl/tests/checkpoint_engine/test_special_server_adapter.py @@ -0,0 +1,241 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import asyncio +import os + +import pytest +import ray +from omegaconf import DictConfig +from transformers import PreTrainedTokenizer + +from tests.checkpoint_engine.test_utils import create_trainer_worker_group +from verl.checkpoint_engine import CheckpointEngineManager +from verl.experimental.fully_async_policy.fully_async_rollouter import FullyAsyncLLMServerClient +from verl.single_controller.ray import ( + RayResourcePool, +) +from verl.utils.config import omega_conf_to_dataclass +from verl.workers.config import CheckpointEngineConfig, HFModelConfig +from verl.workers.rollout.llm_server import LLMServerClient, LLMServerManager + + +@pytest.fixture +def init_config() -> DictConfig: + from hydra import compose, initialize_config_dir + + with initialize_config_dir(config_dir=os.path.abspath("verl/trainer/config")): + config = compose( + config_name="ppo_trainer", + overrides=[ + "+async_training.partial_rollout=True", + ], + ) + + config.actor_rollout_ref.model.path = os.path.expanduser("~/models/Qwen/Qwen3-VL-2B-Instruct") + config.actor_rollout_ref.rollout.name = os.environ["ROLLOUT_NAME"] + config.actor_rollout_ref.rollout.max_num_seqs = 256 + config.actor_rollout_ref.rollout.response_length = 4096 + config.actor_rollout_ref.rollout.checkpoint_engine.backend = "nccl" + config.actor_rollout_ref.rollout.nnodes = 1 + config.trainer.n_gpus_per_node = 4 + config.trainer.nnodes = 1 + + return config + + +async def _run_update_weights_with_global_steps_none( + server_manager: LLMServerClient, + checkpoint_manager: CheckpointEngineManager, + tokenizer: PreTrainedTokenizer, +): + await checkpoint_manager.update_weights(global_steps=None) + prompt = [{"role": "user", "content": "How to make a sandwich?"}] + prompt_ids = tokenizer.apply_chat_template(prompt, add_generation_prompt=True, tokenize=True) + output = await server_manager.generate( + request_id="test_0", + prompt_ids=prompt_ids, + sampling_params={ + "temperature": 1.0, + "logprobs": True, + }, + ) + assert output.stop_reason not in ("aborted", "abort"), ( + f"output.stop_reason is {output.stop_reason}, expected not abort" + ) + assert output.extra_fields["global_steps"] is None, ( + f"output.extra_fields['global_steps'] is {output.extra_fields['global_steps']}, expected None" + ) + print("========== [update_weights with global_steps=None] ==========") + print("[RESPONSE]", tokenizer.decode(output.token_ids, skip_special_tokens=True)) + + +async def _run_server_manager_without_resume( + initial_steps: int, + train_steps: int, + server_manager: LLMServerClient, + checkpoint_manager: CheckpointEngineManager, + prompts: list[list[dict]], + tokenizer: PreTrainedTokenizer, +): + for global_steps in range(initial_steps, initial_steps + train_steps): + tasks = [] + for i, prompt in enumerate(prompts): + prompt_ids = tokenizer.apply_chat_template(prompt, add_generation_prompt=True, tokenize=True) + tasks.append( + asyncio.create_task( + server_manager.generate( + request_id=f"test_{global_steps}_{i}", + prompt_ids=prompt_ids, + sampling_params={ + "temperature": 1.0, + "logprobs": True, + }, + ) + ) + ) + + # wait a while and update weights to interrupt the generation + await asyncio.sleep(2) + await checkpoint_manager.update_weights(global_steps=global_steps) + + outputs = await asyncio.gather(*tasks) + expected_steps = global_steps - 1 + for output in outputs: + global_steps = output.extra_fields["global_steps"] + assert output.stop_reason in ("aborted", "abort"), ( + f"output.stop_reason is {output.stop_reason}, expected in abort" + ) + assert global_steps == expected_steps, f"output.global_steps is {global_steps}, expected {expected_steps}" + print(f"========== [{initial_steps=}, {train_steps=}] ==========") + print("[RESPONSE]", tokenizer.decode(outputs[0].token_ids, skip_special_tokens=True)) + + +async def _run_server_manager_with_resume( + initial_steps: int, + train_steps: int, + server_manager: FullyAsyncLLMServerClient, + checkpoint_manager: CheckpointEngineManager, + prompts: list[list[dict]], + tokenizer: PreTrainedTokenizer, +): + # 1. rollout generate responses + tasks = [] + for i, prompt in enumerate(prompts): + prompt_ids = tokenizer.apply_chat_template(prompt, add_generation_prompt=True, tokenize=True) + tasks.append( + asyncio.create_task( + server_manager.generate( + request_id=f"test_{initial_steps}_{i}", + prompt_ids=prompt_ids, + sampling_params={ + "temperature": 1.0, + "logprobs": True, + }, + ) + ) + ) + + # 2. trainer update weights to rollout multiple times + for global_steps in range(initial_steps, initial_steps + train_steps): + # wait a while and update weights to interrupt the generation + await asyncio.sleep(2) + await checkpoint_manager.update_weights(global_steps=global_steps) + + # 3. wait for rollout generate responses finished + outputs = await asyncio.gather(*tasks) + expected_min_steps = initial_steps - 1 + for output in outputs: + min_global_steps = output.extra_fields["min_global_steps"] + max_global_steps = output.extra_fields["max_global_steps"] + assert min_global_steps == expected_min_steps, ( + f"output.min_global_steps is {min_global_steps}, expected {expected_min_steps}" + ) + assert max_global_steps > expected_min_steps, ( + f"output.max_global_steps is {max_global_steps}, expected > {expected_min_steps}" + ) + assert output.stop_reason not in ("aborted", "abort"), ( + f"output.stop_reason is {output.stop_reason}, expected not abort" + ) + print(f"========== [{initial_steps=}, {train_steps=}] ==========") + print("[RESPONSE]", tokenizer.decode(outputs[0].token_ids, skip_special_tokens=True)) + + +@pytest.mark.asyncio +async def test_server_adapter(init_config): + ray.init( + runtime_env={ + "env_vars": { + "TOKENIZERS_PARALLELISM": "true", + "NCCL_DEBUG": "WARN", + "VLLM_LOGGING_LEVEL": "INFO", + "VLLM_USE_V1": "1", + "VLLM_DISABLE_COMPILE_CACHE": "1", + } + } + ) + + # 1. create trainer worker group + model_config: HFModelConfig = omega_conf_to_dataclass(init_config.actor_rollout_ref.model) + checkpoint_engine_config: CheckpointEngineConfig = omega_conf_to_dataclass( + init_config.actor_rollout_ref.rollout.checkpoint_engine + ) + trainer_pool = RayResourcePool(process_on_nodes=[init_config.trainer.n_gpus_per_node], max_colocate_count=3) + trainer = create_trainer_worker_group(trainer_pool, model_config, checkpoint_engine_config) + trainer.reset() + + # 2. create standalone rollout with AgentLoopManager + llm_server_manager = await LLMServerManager.create(config=init_config) + + # 3. create checkpoint engine manager + checkpoint_manager = CheckpointEngineManager( + config=checkpoint_engine_config, trainer=trainer, replicas=llm_server_manager.get_replicas() + ) + + n = 4 + prompts = [ + [{"role": "user", "content": "Please write an article about the history of China, at least 1000 words."}], + [{"role": "user", "content": "Please write an article about the history of America, at least 1000 words."}], + [{"role": "user", "content": "Please write an article about the geography of China, at least 1000 words."}], + [{"role": "user", "content": "Please write an article about the geography of America, at least 1000 words."}], + ] * n + + # 4. test update_weights with global_steps=None + await _run_update_weights_with_global_steps_none( + server_manager=llm_server_manager.get_client(), + checkpoint_manager=checkpoint_manager, + tokenizer=model_config.tokenizer, + ) + + # 5. test LLMServerClient without partial rollout resume + await checkpoint_manager.update_weights(global_steps=0) + await _run_server_manager_without_resume( + initial_steps=1, + train_steps=3, + server_manager=llm_server_manager.get_client(), + checkpoint_manager=checkpoint_manager, + prompts=prompts, + tokenizer=model_config.tokenizer, + ) + + # 6. test FullyLLMServerClient with partial rollout resume + await _run_server_manager_with_resume( + initial_steps=4, + train_steps=3, + server_manager=llm_server_manager.get_client(client_cls=FullyAsyncLLMServerClient), + checkpoint_manager=checkpoint_manager, + prompts=prompts, + tokenizer=model_config.tokenizer, + ) + + ray.shutdown() diff --git a/verl/tests/checkpoint_engine/test_utils.py b/verl/tests/checkpoint_engine/test_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..cbb1e75cd3351cbd7e9b5a4e87cd45249cfaddae --- /dev/null +++ b/verl/tests/checkpoint_engine/test_utils.py @@ -0,0 +1,183 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import asyncio +from typing import Generator + +import ray +import torch +from transformers import AutoModelForCausalLM + +from verl.checkpoint_engine import CheckpointEngineRegistry, CheckpointEngineWorker +from verl.single_controller.base.decorator import Dispatch, register +from verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup +from verl.utils.device import get_device_name +from verl.utils.fs import copy_to_local +from verl.workers.config import CheckpointEngineConfig, FSDPEngineConfig, HFModelConfig, RolloutConfig +from verl.workers.engine_workers import TrainingWorker, TrainingWorkerConfig +from verl.workers.rollout import BaseRollout, RolloutReplica + + +class TrainingWorkerTest(TrainingWorker): + def __init__(self, config: TrainingWorkerConfig, checkpoint_engine_config: CheckpointEngineConfig) -> None: + super().__init__(config) + + backend = checkpoint_engine_config.backend + bucket_size = checkpoint_engine_config.update_weights_bucket_megabytes << 20 + engine_kwargs = checkpoint_engine_config.engine_kwargs.get(backend, {}) + if torch.distributed.get_rank() == 0: + engine_kwargs["is_master"] = True + self.checkpoint_engine = CheckpointEngineRegistry.new(backend, bucket_size=bucket_size, **engine_kwargs) + + @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False) + async def update_weights(self, global_steps: int = None, mode: str = "auto"): + per_tensor_param, _ = self.engine.get_per_tensor_param() + await self.checkpoint_engine.send_weights(per_tensor_param) + + @register(dispatch_mode=Dispatch.DP_COMPUTE, blocking=False) + def execute_checkpoint_engine(self, method: str, *args, **kwargs): + return getattr(self.checkpoint_engine, method)(*args, **kwargs) + + +class MockServerAdapter(BaseRollout): + def __init__(self, config: RolloutConfig, model_config: HFModelConfig, check_allclose: bool = True): + super().__init__(config, model_config, device_mesh=None) + self.check_allclose = check_allclose + self.model = None + self.received_weights: dict[str, torch.Tensor] = {} + + async def resume(self, tags: list[str]): + raise NotImplementedError() + + async def release(self): + raise NotImplementedError() + + async def update_weights( + self, + weights: Generator[tuple[str, torch.Tensor], None, None], + **kwargs, + ): + async for name, weight in weights: + weight = weight.clone() + if self.check_allclose: + self.received_weights[name] = weight.clone() + + def check_weights(self): + if not self.check_allclose: + return + + if self.model is None: + local_path = copy_to_local(self.model_config.path) + self.model = AutoModelForCausalLM.from_pretrained(local_path, torch_dtype=torch.bfloat16, device_map="cpu") + + for name, weight in self.model.state_dict().items(): + assert name in self.received_weights, f"weight {name} not received" + received = self.received_weights[name] + assert torch.allclose(weight.to(received.device), received), f"weight {name} not equal" + self.received_weights.clear() + print("Check passed, all weights are equal!") + + +class MockReplica(RolloutReplica): + async def init_hybrid(self, worker_group: RayWorkerGroup): + """Init hybrid rollout server, rollout engine and training engine(fsdp/megatron) fused in same process. + + Args: + worker_group: RayWorkerGroup, fused workers where training engine(fsdp/megatron) have been initialized. + """ + self.workers = worker_group.workers[ + self.world_size * self.replica_rank : self.world_size * (self.replica_rank + 1) + ] + + def get_ray_class_with_init_args(self) -> RayClassWithInitArgs: + """Get rollout worker actor class for colocated and standalone mode.""" + raise NotImplementedError + + async def launch_servers(self): + """Launch http server in each node.""" + raise NotImplementedError + + +class CheckpointEngineWorkerTest(CheckpointEngineWorker): + def __init__( + self, rollout_config: RolloutConfig, model_config: HFModelConfig, check_allclose: bool = True, *args, **kwargs + ) -> None: + server_adapter = MockServerAdapter(rollout_config, model_config, check_allclose) + super().__init__(rollout_config, model_config, server_adapter, *args, **kwargs) + + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def check_weights(self): + self.server_adapter.check_weights() + + +def create_trainer_worker_group( + resource_pool: RayResourcePool, model_config: HFModelConfig, checkpoint_engine_config: CheckpointEngineConfig +) -> RayWorkerGroup: + engine_config = FSDPEngineConfig(forward_only=True, fsdp_size=resource_pool.world_size, strategy="fsdp") + trainer_config = TrainingWorkerConfig( + model_type="language_model", + model_config=model_config, + engine_config=engine_config, + ) + + ray_cls_with_init = RayClassWithInitArgs( + cls=ray.remote(TrainingWorkerTest), + config=trainer_config, + checkpoint_engine_config=checkpoint_engine_config, + ) + ray_cls_with_init.update_options( + { + "runtime_env": { + "env_vars": { + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", + } + } + } + ) + wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init, device_name=get_device_name()) + return wg + + +async def create_rollout_worker_group( + resource_pool: RayResourcePool, + model_config: HFModelConfig, + rollout_config: RolloutConfig, + check_allclose: bool = True, +) -> tuple[RayWorkerGroup, list[MockReplica]]: + # create rollout worker group + ray_cls_with_init = RayClassWithInitArgs( + cls=ray.remote(CheckpointEngineWorkerTest), + model_config=model_config, + rollout_config=rollout_config, + check_allclose=check_allclose, + ) + wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init, device_name=get_device_name()) + + # create rollout replicas + rollout_world_size = ( + rollout_config.tensor_model_parallel_size + * rollout_config.data_parallel_size + * rollout_config.pipeline_model_parallel_size + ) + num_replicas = wg.world_size // rollout_world_size + replicas = [] + for replica_rank in range(num_replicas): + replica = MockReplica( + replica_rank=replica_rank, + config=rollout_config, + model_config=model_config, + ) + replicas.append(replica) + await asyncio.gather(*[replica.init_hybrid(wg) for replica in replicas]) + + return wg, replicas diff --git a/verl/tests/experimental/agent_loop/agent_utils.py b/verl/tests/experimental/agent_loop/agent_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..939e383e42a54005a02d56768c513f83a668b840 --- /dev/null +++ b/verl/tests/experimental/agent_loop/agent_utils.py @@ -0,0 +1,102 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import ray +from omegaconf import DictConfig + +from verl.checkpoint_engine import CheckpointEngineManager +from verl.experimental.agent_loop import AgentLoopManager +from verl.experimental.reward_loop import RewardLoopManager +from verl.single_controller.ray import RayClassWithInitArgs, RayWorkerGroup +from verl.single_controller.ray.base import create_colocated_worker_cls +from verl.trainer.ppo.ray_trainer import ResourcePoolManager, Role +from verl.utils import omega_conf_to_dataclass +from verl.utils.device import get_device_name +from verl.workers.engine_workers import ActorRolloutRefWorker +from verl.workers.rollout.llm_server import LLMServerManager + + +def init_agent_loop_manager(config: DictConfig) -> AgentLoopManager | RayWorkerGroup: + # =========================== 1. Create hybrid ActorRollout workers =========================== + # The unified model-engine ActorRolloutRefWorker supports both sync and async rollout modes. + actor_rollout_cls = ActorRolloutRefWorker + role_worker_mapping = { + Role.ActorRollout: ray.remote(actor_rollout_cls), + } + + global_pool_id = "global_pool" + resource_pool_spec = { + global_pool_id: [config.trainer.n_gpus_per_node] * config.trainer.nnodes, + } + mapping = { + Role.ActorRollout: global_pool_id, + } + if config.reward.reward_model.enable_resource_pool: + mapping[Role.RewardModel] = "reward_pool" + if config.reward.reward_model.n_gpus_per_node <= 0: + raise ValueError("config.reward.reward_model.n_gpus_per_node must be greater than 0") + if config.reward.reward_model.nnodes <= 0: + raise ValueError("config.reward.reward_model.nnodes must be greater than 0") + + reward_pool = [config.reward.reward_model.n_gpus_per_node] * config.reward.reward_model.nnodes + resource_pool_spec["reward_pool"] = reward_pool + resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping) + resource_pool_manager.create_resource_pool() + resource_pool_to_cls = {pool: {} for pool in resource_pool_manager.resource_pool_dict.values()} + + # create actor and rollout + resource_pool = resource_pool_manager.get_resource_pool(Role.ActorRollout) + actor_rollout_cls = RayClassWithInitArgs( + cls=role_worker_mapping[Role.ActorRollout], config=config.actor_rollout_ref, role="actor_rollout" + ) + resource_pool_to_cls[resource_pool]["actor_rollout"] = actor_rollout_cls + + device_name = get_device_name() + all_wg = {} + for resource_pool, class_dict in resource_pool_to_cls.items(): + worker_dict_cls = create_colocated_worker_cls(class_dict=class_dict) + wg_dict = RayWorkerGroup( + resource_pool=resource_pool, ray_cls_with_init=worker_dict_cls, device_name=device_name + ) + spawn_wg = wg_dict.spawn(prefix_set=class_dict.keys()) + all_wg.update(spawn_wg) + actor_rollout_wg = all_wg["actor_rollout"] + actor_rollout_wg.init_model() + + if config.actor_rollout_ref.rollout.mode == "sync": + raise ValueError("Agent loop tests require async rollout mode. Please set rollout.mode=async.") + + # =========================== 2. Create AgentLoopManager =========================== + rm_resource_pool = ( + resource_pool_manager.get_resource_pool(Role.RewardModel) if config.reward.reward_model.enable else None + ) + reward_loop_manager = RewardLoopManager( + config=config, + rm_resource_pool=rm_resource_pool, + ) + llm_server_manager = LLMServerManager.create(config=config, worker_group=actor_rollout_wg) + agent_loop_manager = AgentLoopManager.create( + config=config, + llm_client=llm_server_manager.get_client(), + reward_loop_worker_handles=reward_loop_manager.reward_loop_workers, + ) + checkpoint_manager = CheckpointEngineManager( + config=omega_conf_to_dataclass(config.actor_rollout_ref.rollout.checkpoint_engine), + trainer=actor_rollout_wg, + replicas=llm_server_manager.get_replicas(), + ) + checkpoint_manager.sleep_replicas() + checkpoint_manager.update_weights() + + return agent_loop_manager diff --git a/verl/tests/experimental/agent_loop/function_tool_examples.py b/verl/tests/experimental/agent_loop/function_tool_examples.py new file mode 100644 index 0000000000000000000000000000000000000000..2087d013580af3bbcfc0de407c8021feb8e0d07c --- /dev/null +++ b/verl/tests/experimental/agent_loop/function_tool_examples.py @@ -0,0 +1,86 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""``@function_tool`` examples loaded by the tests. + +Includes both sync and async functions to exercise the dispatcher's +``asyncio.iscoroutinefunction`` branch. +""" + +import asyncio + +from verl.tools.function_tool import function_tool + + +@function_tool +def get_weather(city: str) -> dict: + """Get the current weather for a city. + + Args: + city: The city to look up, e.g. "Tokyo" or "San Francisco". + """ + # Stubbed lookup table; in production this would hit a weather API. The + # values are deliberately unusual so a test can distinguish a real tool + # response from a number the model guessed. + table = { + "tokyo": {"temperature_c": 17.3, "condition": "drizzle"}, + "san francisco": {"temperature_c": 14.8, "condition": "fog"}, + "new york": {"temperature_c": 21.6, "condition": "sunny"}, + } + return table.get(city.lower(), {"temperature_c": -273.15, "condition": "unknown"}) + + +@function_tool +def calculator(expression: str) -> str: + """Evaluate an arithmetic expression and return the result. + + Supports +, -, *, /, **, parentheses, and unary minus. Use this for any + numerical computation instead of doing mental arithmetic. + + Args: + expression: A Python-style arithmetic expression, e.g. "(3+4)*5". + """ + import ast + import operator as op + + ops = {ast.Add: op.add, ast.Sub: op.sub, ast.Mult: op.mul, ast.Div: op.truediv, ast.Pow: op.pow, ast.USub: op.neg} + + def _eval(node): + if isinstance(node, ast.Constant) and isinstance(node.value, int | float): # noqa: UP038 + return node.value + if isinstance(node, ast.BinOp): + return ops[type(node.op)](_eval(node.left), _eval(node.right)) + if isinstance(node, ast.UnaryOp): + return ops[type(node.op)](_eval(node.operand)) + raise ValueError(f"unsupported node: {ast.dump(node)}") + + try: + return str(_eval(ast.parse(expression, mode="eval").body)) + except Exception as e: + return f"ERROR: {e}" + + +@function_tool +async def fetch_url(url: str) -> str: + """Fetch the contents of a URL (async). + + This is a stubbed example that demonstrates ``async def`` tools; the + real implementation would use ``aiohttp`` or ``httpx``. + + Args: + url: The URL to fetch. + """ + # Yield once so the tool actually behaves like an awaitable, then return + # a deterministic stub payload that tests can assert on. + await asyncio.sleep(0) + return f"" diff --git a/verl/tests/experimental/agent_loop/qwen_vl_tool_chat_template.jinja2 b/verl/tests/experimental/agent_loop/qwen_vl_tool_chat_template.jinja2 new file mode 100644 index 0000000000000000000000000000000000000000..9fea57ff86b54917ff806a28b3617bb79517c494 --- /dev/null +++ b/verl/tests/experimental/agent_loop/qwen_vl_tool_chat_template.jinja2 @@ -0,0 +1,150 @@ +{% set image_count = namespace(value=0) %} +{% set video_count = namespace(value=0) %} +{%- if tools %} +{{- '<|im_start|>system\n' }} +{%- if messages[0]['role'] == 'system' %} +{%- if messages[0]['content'] is string %} +{{- messages[0]['content'] }} +{%- else %} +{{- messages[0]['content'][0]['text'] }} +{%- endif %} +{%- else %} +{{- 'You are a helpful assistant.' }} +{%- endif %} +{{- "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within XML tags:\n" }} +{%- for tool in tools %} +{{- "\n" }} +{{- tool | tojson }} +{%- endfor %} +{{- "\n\n\nFor each function call, return a json object with function name and arguments within XML tags:\n\n{\"name\": , \"arguments\": }\n<|im_end|>\n" }} +{% for message in messages %} +{% if message['role'] != 'system' or loop.first == false %} +{%- if (message.role == "user") or (message.role == "system" and not loop.first) or (message.role == "assistant" and not message.tool_calls) %} +<|im_start|>{{ message['role'] }} +{% if message['content'] is string %} +{{ message['content'] }}<|im_end|> +{% else %} +{% for content in message['content'] %} +{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %} +{% set image_count.value = image_count.value + 1 %} +{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|> +{% elif content['type'] == 'video' or 'video' in content %} +{% set video_count.value = video_count.value + 1 %} +{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|> +{% elif 'text' in content %} +{{ content['text'] }} +{% endif %} +{% endfor %}<|im_end|> +{% endif %} +{%- elif message.role == "assistant" %} +{{- '<|im_start|>' + message.role }} +{%- if message.content %} +{{- '\n' + message.content }} +{%- endif %} +{%- for tool_call in message.tool_calls %} +{%- if tool_call.function is defined %} +{%- set tool_call = tool_call.function %} +{%- endif %} +{{- '\n\n{"name": "' }} +{{- tool_call.name }} +{{- '", "arguments": ' }} +{{- tool_call.arguments | tojson }} +{{- '}\n' }} +{%- endfor %} +{{- '<|im_end|>\n' }} +{%- elif message.role == "tool" %} +{%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != "tool") %} +{{- '<|im_start|>user' }} +{%- endif %} +{{- '\n\n' }} +{% if message['content'] is string %} +{{ message.content }} +{% else %} +{% for content in message['content'] %} +{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %} +{% set image_count.value = image_count.value + 1 %} +{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|> +{% elif content['type'] == 'video' or 'video' in content %} +{% set video_count.value = video_count.value + 1 %} +{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|> +{% elif content['type'] == 'text' or 'text' in content %} +{{ content['text'] }} +{% endif %} +{% endfor %} +{% endif %} +{{- '\n' }} +{%- if loop.last or (messages[loop.index0 + 1].role != "tool") %} +{{- '<|im_end|>\n' }} +{%- endif %} +{%- endif %} +{% endif %} +{% endfor %} +{%- else %} +{% for message in messages %} +{% if loop.first and message['role'] != 'system' %} +<|im_start|>system +You are a helpful assistant.<|im_end|> +{% endif %} +{%- if (message.role == "user") or (message.role == "system" and not loop.first) or (message.role == "assistant" and not message.tool_calls) %} +<|im_start|>{{ message['role'] }} +{% if message['content'] is string %} +{{ message['content'] }}<|im_end|> +{% else %} +{% for content in message['content'] %} +{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %} +{% set image_count.value = image_count.value + 1 %} +{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|> +{% elif content['type'] == 'video' or 'video' in content %} +{% set video_count.value = video_count.value + 1 %} +{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|> +{% elif 'text' in content %} +{{ content['text'] }} +{% endif %} +{% endfor %}<|im_end|> +{% endif %} +{%- elif message.role == "assistant" %} +{{- '<|im_start|>' + message.role }} +{%- if message.content %} +{{- '\n' + message.content }} +{%- endif %} +{%- for tool_call in message.tool_calls %} +{%- if tool_call.function is defined %} +{%- set tool_call = tool_call.function %} +{%- endif %} +{{- '\n\n{"name": "' }} +{{- tool_call.name }} +{{- '", "arguments": ' }} +{{- tool_call.arguments | tojson }} +{{- '}\n' }} +{%- endfor %} +{{- '<|im_end|>\n' }} +{%- elif message.role == "tool" %} +{%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != "tool") %} +{{- '<|im_start|>user' }} +{%- endif %} +{{- '\n\n' }} +{% if message['content'] is string %} +{{ message.content }} +{% else %} +{% for content in message['content'] %} +{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %} +{% set image_count.value = image_count.value + 1 %} +{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|> +{% elif content['type'] == 'video' or 'video' in content %} +{% set video_count.value = video_count.value + 1 %} +{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|> +{% elif content['type'] == 'text' or 'text' in content %} +{{ content['text'] }} +{% endif %} +{% endfor %} +{% endif %} +{{- '\n' }} +{%- if loop.last or (messages[loop.index0 + 1].role != "tool") %} +{{- '<|im_end|>\n' }} +{%- endif %} +{%- endif %} +{% endfor %} +{%- endif %} +{% if add_generation_prompt %} +<|im_start|>assistant +{% endif %} \ No newline at end of file diff --git a/verl/tests/experimental/agent_loop/test_agent_loop_extra_fields_schema_on_cpu.py b/verl/tests/experimental/agent_loop/test_agent_loop_extra_fields_schema_on_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..c28214877e3a19cf3f37d2a4cac24daec849903d --- /dev/null +++ b/verl/tests/experimental/agent_loop/test_agent_loop_extra_fields_schema_on_cpu.py @@ -0,0 +1,307 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import warnings +from typing import Any, Optional + +import numpy as np +import pytest +import torch +from omegaconf import OmegaConf + +from verl.experimental.agent_loop.agent_loop import ( + AgentLoopMetrics, + AgentLoopOutput, + AgentLoopWorker, + DictConfigWrap, + _InternalAgentLoopOutput, +) +from verl.experimental.agent_loop.single_turn_agent_loop import SingleTurnAgentLoop +from verl.utils.dataset.rl_dataset import RLHFDataset +from verl.workers.rollout.replica import TokenOutput + + +class _FakeServerManager: + async def generate( + self, + request_id: str, + *, + prompt_ids: list[int], + sampling_params: dict[str, Any], + image_data: Optional[list[Any]] = None, + video_data: Optional[list[Any]] = None, + audio_data: Optional[list[Any]] = None, + mm_processor_kwargs: Optional[dict[str, Any]] = None, + ) -> TokenOutput: + del request_id, sampling_params, image_data, video_data, audio_data, mm_processor_kwargs + # Return a short, deterministic "generation" for testing. + return TokenOutput(token_ids=prompt_ids[-1:] + [11, 12, 13], log_probs=[0.0, 0.0, 0.0, 0.0]) + + async def generate_for_partial( + self, + request_id: str, + *, + prompt_ids: list[int], + sampling_params: dict[str, Any], + image_data: Optional[list[Any]] = None, + video_data: Optional[list[Any]] = None, + audio_data: Optional[list[Any]] = None, + mm_processor_kwargs: Optional[dict[str, Any]] = None, + ) -> tuple[list[int], list[float], bool]: + del request_id, sampling_params, image_data, video_data, audio_data, mm_processor_kwargs + # Return a short partial generation and "not cancelled". + response_ids = prompt_ids[-1:] + [21, 22] + response_logprobs = [0.0] * len(response_ids) + return response_ids, response_logprobs, False + + +class _FakeTokenizer: + padding_side = "right" + + def apply_chat_template( + self, + messages: list[dict[str, Any]], + *, + tools: Optional[list[dict]] = None, + add_generation_prompt: bool = True, + tokenize: bool = True, + **kwargs, + ) -> list[int]: + del messages, tools, add_generation_prompt, tokenize, kwargs + # Minimal tokenization: return a small prompt. + return [101, 102] + + def pad( + self, + encoded_inputs: dict[str, list[int]], + *, + padding: str, + max_length: int, + return_tensors: str, + return_attention_mask: bool, + ) -> dict[str, torch.Tensor]: + del padding, return_tensors + input_ids = encoded_inputs["input_ids"] + if len(input_ids) > max_length: + if self.padding_side == "left": + input_ids = input_ids[-max_length:] + else: + input_ids = input_ids[:max_length] + + pad_len = max_length - len(input_ids) + if self.padding_side == "left": + padded_ids = [0] * pad_len + input_ids + attention_mask = [0] * pad_len + [1] * len(input_ids) + else: + padded_ids = input_ids + [0] * pad_len + attention_mask = [1] * len(input_ids) + [0] * pad_len + + output = {"input_ids": torch.tensor([padded_ids], dtype=torch.long)} + if return_attention_mask: + output["attention_mask"] = torch.tensor([attention_mask], dtype=torch.long) + return output + + def decode(self, ids: list[int] | torch.Tensor, skip_special_tokens: bool = True) -> str: + del ids, skip_special_tokens + return "" + + +def _pad_1d(ids: list[int], *, length: int, pad_id: int = 0) -> list[int]: + if len(ids) > length: + return ids[:length] + return ids + [pad_id] * (length - len(ids)) + + +def _to_internal( + *, + output_prompt_ids: list[int], + output_response_ids: list[int], + output_response_mask: list[int], + metrics: AgentLoopMetrics, + extra_fields: dict[str, Any], + num_turns: int, + prompt_len: int, + response_len: int, +) -> _InternalAgentLoopOutput: + prompt_ids = _pad_1d(output_prompt_ids, length=prompt_len, pad_id=0) + response_ids = _pad_1d(output_response_ids, length=response_len, pad_id=0) + response_mask = _pad_1d(output_response_mask, length=response_len, pad_id=0) + + seq_len = prompt_len + response_len + attention_mask = _pad_1d([1] * len(output_prompt_ids), length=prompt_len, pad_id=0) + _pad_1d( + [1] * len(output_response_ids), + length=response_len, + pad_id=0, + ) + input_ids = prompt_ids + response_ids + position_ids = list(range(seq_len)) + + def t(x: list[int]) -> torch.Tensor: + return torch.tensor([x], dtype=torch.long) + + return _InternalAgentLoopOutput( + prompt_ids=t(prompt_ids), + response_ids=t(response_ids), + response_mask=t(response_mask), + attention_mask=t(attention_mask), + input_ids=t(input_ids), + position_ids=t(position_ids), + response_logprobs=None, + routed_experts=None, + multi_modal_inputs=None, + multi_modal_data=None, + reward_score=None, + num_turns=num_turns, + metrics=metrics, + extra_fields=extra_fields, + ) + + +@pytest.mark.asyncio +async def test_agent_loop_extra_fields_schema_stable_for_training_concat_on_cpu(): + # Minimal config surface used by the agent loops. + config = OmegaConf.create( + { + "actor_rollout_ref": { + "rollout": {"prompt_length": 16, "response_length": 16, "multi_turn": {"tool_config_path": None}}, + "model": {}, + }, + "data": { + "tool_config_path": None, + "apply_chat_template_kwargs": {}, + }, + } + ) + + server_manager = _FakeServerManager() + tokenizer = _FakeTokenizer() + processor = None + + trainer_config = DictConfigWrap(config) + data_config = DictConfigWrap(config.data) + + single_turn = SingleTurnAgentLoop( + trainer_config=trainer_config, + server_manager=server_manager, + tokenizer=tokenizer, + processor=processor, + dataset_cls=RLHFDataset, + data_config=data_config, + ) + + raw_prompt = [{"role": "user", "content": "hi"}] + sampling_params: dict[str, Any] = {} + + out = await single_turn.run(sampling_params=sampling_params, raw_prompt=raw_prompt) + + # Agent loop outputs should always contain these fields with consistent types. + assert out.extra_fields["turn_scores"] == [] + assert out.extra_fields["tool_rewards"] == [] + + internal_a = _to_internal( + output_prompt_ids=out.prompt_ids, + output_response_ids=out.response_ids, + output_response_mask=out.response_mask, + metrics=out.metrics, + extra_fields=out.extra_fields, + num_turns=out.num_turns, + prompt_len=len(out.prompt_ids), + response_len=len(out.response_ids), + ) + + # Mimic two "worker chunks" and concatenate as in training. + dummy_worker = type( + "_DummyWorker", + (), + {"reward_loop_worker_handles": None, "distillation_enabled": False}, + )() + merged = AgentLoopWorker._postprocess( + dummy_worker, + inputs=[internal_a], + input_non_tensor_batch={ + "index": np.array([0], dtype=object), + "agent_name": np.array(["single_turn_agent"], dtype=object), + }, + ) + + # Stable schema: present regardless of which loop produced a sample. + stable_keys = ( + "turn_scores", + "tool_rewards", + "min_global_steps", + "max_global_steps", + "extras", + ) + for key in stable_keys: + assert key in merged.non_tensor_batch, f"missing key in merged batch: {key}" + assert merged.non_tensor_batch[key].shape == (1,), ( + f"invalid shape for {key}: {merged.non_tensor_batch[key].shape}" + ) + + # And the list-typed fields are actually lists (not missing / scalar). + assert merged.non_tensor_batch["turn_scores"][0] == [] + assert merged.non_tensor_batch["tool_rewards"][0] == [] + + +@pytest.mark.asyncio +async def test_agent_loop_postprocess_accepts_read_only_routed_experts_on_cpu(): + class _DummyWorker: + _compute_multi_modal_inputs = AgentLoopWorker._compute_multi_modal_inputs + _compute_position_ids = AgentLoopWorker._compute_position_ids + _get_mm_processor_kwargs = AgentLoopWorker._get_mm_processor_kwargs + _compute_score = AgentLoopWorker._compute_score + _compute_teacher_logprobs = AgentLoopWorker._compute_teacher_logprobs + distillation_enabled = False + + def __init__(self): + self.tokenizer = _FakeTokenizer() + self.rollout_config = OmegaConf.create({"prompt_length": 4, "response_length": 4}) + self.processor = None + self.mm_processor_kwargs = {} + self.reward_loop_worker_handles = None + + routed_experts = np.arange(8, dtype=np.int64).reshape(4, 2, 1) + routed_experts.setflags(write=False) + assert not routed_experts.flags.writeable + + output = AgentLoopOutput( + prompt_ids=[101, 102], + response_ids=[11, 12], + response_mask=[1, 1], + routed_experts=routed_experts, + metrics=AgentLoopMetrics(), + extra_fields={}, + ) + + with warnings.catch_warnings(): + warnings.filterwarnings( + "error", + message="The given NumPy array is not writable.*", + category=UserWarning, + ) + internal = await AgentLoopWorker._agent_loop_postprocess( + _DummyWorker(), + output, + validate=False, + raw_prompt=[{"role": "user", "content": "hi"}], + ) + + expected = torch.tensor(routed_experts.copy()).unsqueeze(0) + assert internal.routed_experts is not None + assert internal.routed_experts.shape == (1, 8, 2, 1) + torch.testing.assert_close(internal.routed_experts[:, 2:6], expected) + assert torch.count_nonzero(internal.routed_experts[:, :2]) == 0 + assert torch.count_nonzero(internal.routed_experts[:, 6:]) == 0 diff --git a/verl/tests/experimental/agent_loop/test_audio_server_contract_on_cpu.py b/verl/tests/experimental/agent_loop/test_audio_server_contract_on_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..0124a3e011323cd2d026724b2e9d0e07be7fe394 --- /dev/null +++ b/verl/tests/experimental/agent_loop/test_audio_server_contract_on_cpu.py @@ -0,0 +1,62 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import ast +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[3] +LLM_SERVER_SOURCE = REPO_ROOT / "verl/workers/rollout/llm_server.py" +FULLY_ASYNC_ROLLOUTER_SOURCE = REPO_ROOT / "verl/experimental/fully_async_policy/fully_async_rollouter.py" +VLLM_SERVER_SOURCE = REPO_ROOT / "verl/workers/rollout/vllm_rollout/vllm_async_server.py" + + +def _load_module_ast(path: Path) -> ast.Module: + return ast.parse(path.read_text(encoding="utf-8")) + + +def test_async_server_manager_generate_accepts_audio_and_mm_kwargs() -> None: + module = _load_module_ast(LLM_SERVER_SOURCE) + + generate_fn = None + for node in module.body: + if isinstance(node, ast.ClassDef) and node.name == "LLMServerClient": + for inner in node.body: + if isinstance(inner, ast.AsyncFunctionDef) and inner.name == "generate": + generate_fn = inner + break + + assert generate_fn is not None + arg_names = [arg.arg for arg in generate_fn.args.kwonlyargs] + assert "audio_data" in arg_names + assert "mm_processor_kwargs" in arg_names + + +def test_async_server_manager_generate_forwards_audio_and_mm_kwargs() -> None: + source = LLM_SERVER_SOURCE.read_text(encoding="utf-8") + assert 'multimodal_kwargs["audio_data"] = audio_data' in source + assert 'multimodal_kwargs["mm_processor_kwargs"] = mm_processor_kwargs' in source + + +def test_fully_async_server_manager_generate_forwards_audio_and_mm_kwargs() -> None: + source = FULLY_ASYNC_ROLLOUTER_SOURCE.read_text(encoding="utf-8") + assert "audio_data=audio_data" in source + assert "mm_processor_kwargs=mm_processor_kwargs" in source + + +def test_vllm_generate_includes_audio_and_mm_processor_kwargs() -> None: + source = VLLM_SERVER_SOURCE.read_text(encoding="utf-8") + assert 'multi_modal_data["audio"] = audio_data' in source + assert 'prompt_kwargs["mm_processor_kwargs"] = mm_processor_kwargs' in source diff --git a/verl/tests/experimental/agent_loop/test_basic_agent_loop.py b/verl/tests/experimental/agent_loop/test_basic_agent_loop.py new file mode 100644 index 0000000000000000000000000000000000000000..44e27556e520e3fa9756ceb21177de0663640a19 --- /dev/null +++ b/verl/tests/experimental/agent_loop/test_basic_agent_loop.py @@ -0,0 +1,557 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import json +import os +from typing import Any + +import numpy as np +import pytest +import ray +from omegaconf import DictConfig +from transformers.utils import get_json_schema + +from tests.experimental.agent_loop.agent_utils import init_agent_loop_manager +from verl.experimental.agent_loop import get_trajectory_info +from verl.protocol import DataProto +from verl.tools.base_tool import BaseTool, OpenAIFunctionToolSchema +from verl.tools.schemas import ToolResponse +from verl.utils import hf_tokenizer +from verl.workers.rollout.llm_server import GlobalRequestLoadBalancer + + +@pytest.fixture +def init_config() -> DictConfig: + from hydra import compose, initialize_config_dir + + with initialize_config_dir(config_dir=os.path.abspath("verl/trainer/config")): + config = compose( + config_name="ppo_trainer", + overrides=[ + "actor_rollout_ref.actor.use_dynamic_bsz=true", + # test sleep/wake_up with fsdp offload + "actor_rollout_ref.actor.fsdp_config.param_offload=True", + "actor_rollout_ref.actor.fsdp_config.optimizer_offload=True", + "reward.reward_manager.name=dapo", + "+reward.reward_kwargs.overlong_buffer_cfg.enable=False", + "+reward.reward_kwargs.overlong_buffer_cfg.len=3072", + "+reward.reward_kwargs.max_resp_len=4096", + ], + ) + + model_path = os.path.expanduser("~/models/Qwen/Qwen2.5-1.5B-Instruct") + config.actor_rollout_ref.model.path = model_path + config.actor_rollout_ref.rollout.name = os.environ["ROLLOUT_NAME"] + config.actor_rollout_ref.rollout.mode = "async" + config.actor_rollout_ref.rollout.enforce_eager = True + config.actor_rollout_ref.rollout.prompt_length = 4096 + config.actor_rollout_ref.rollout.response_length = 4096 + config.actor_rollout_ref.rollout.n = 4 + config.actor_rollout_ref.rollout.agent.num_workers = 2 + config.actor_rollout_ref.rollout.skip_tokenizer_init = True + + return config + + +def test_single_turn(init_config): + ray.init( + runtime_env={ + "env_vars": { + "TOKENIZERS_PARALLELISM": "true", + "NCCL_DEBUG": "WARN", + "VLLM_LOGGING_LEVEL": "INFO", + "VLLM_USE_V1": "1", + } + } + ) + + agent_loop_manager = init_agent_loop_manager(init_config) + + raw_prompts = [ + [ + { + "role": "user", + "content": "Let's play a role playing game. Your name is Alice, your favorite color is blue.", + } + ], + [{"role": "user", "content": "Let's play a role playing game. Your name is Bob, your favorite color is red."}], + ] + batch = DataProto( + non_tensor_batch={ + "raw_prompt": np.array(raw_prompts), + "agent_name": np.array(["single_turn_agent"] * len(raw_prompts)), + "data_source": np.array(["openai/gsm8k"] * len(raw_prompts)), + "reward_model": np.array([{"style": "rule", "ground_truth": "1.0"}] * len(raw_prompts)), + }, + ) + n = init_config.actor_rollout_ref.rollout.n + batch = batch.repeat(n) + result = agent_loop_manager.generate_sequences(prompts=batch) + assert len(result) == len(raw_prompts) * n + + # check result + seq_len = result.batch["prompts"].size(1) + result.batch["responses"].size(1) + assert result.batch["input_ids"].size(1) == seq_len + assert result.batch["attention_mask"].size(1) == seq_len + assert result.batch["position_ids"].size(1) == seq_len + + if init_config.actor_rollout_ref.rollout.calculate_log_probs: + assert result.batch["rollout_log_probs"].size(1) == result.batch["responses"].size(1) + + # check compute score + assert result.batch["rm_scores"].shape == result.batch["responses"].shape + reward_tensor = result.batch["rm_scores"] + reward_extra_keys = result.meta_info.get("reward_extra_keys", []) + reward_extra_info = {key: result.non_tensor_batch[key] for key in reward_extra_keys} + assert reward_tensor.shape == result.batch["responses"].shape + assert "acc" in reward_extra_info, f"reward_extra_info {reward_extra_info} should contain 'acc'" + assert reward_extra_info["acc"].shape == (len(result),), f"invalid acc: {reward_extra_info['acc']}" + + # check turns + num_turns = result.non_tensor_batch["__num_turns__"] + assert np.all(num_turns == 2) + + print("Test passed!") + ray.shutdown() + + +class WeatherTool(BaseTool): + def get_current_temperature(self, location: str, unit: str = "celsius"): + """Get current temperature at a location. + + Args: + location: The location to get the temperature for, in the format "City, State, Country". + unit: The unit to return the temperature in. Defaults to "celsius". (choices: ["celsius", "fahrenheit"]) + + Returns: + the temperature, the location, and the unit in a dict + """ + print(f"[DEBUG] get_current_temperature: {location}, {unit}") + return { + "temperature": 26.1, + "location": location, + "unit": unit, + } + + def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema: + schema = get_json_schema(self.get_current_temperature) + return OpenAIFunctionToolSchema(**schema) + + async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[ToolResponse, float, dict]: + try: + result = self.get_current_temperature(**parameters) + return ToolResponse(text=json.dumps(result)), 0, {} + except Exception as e: + return ToolResponse(text=str(e)), 0, {} + + +class WeatherToolWithData(BaseTool): + def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema: + schema = get_json_schema(self.get_temperature_date) + return OpenAIFunctionToolSchema(**schema) + + def get_temperature_date(self, location: str, date: str, unit: str = "celsius"): + """Get temperature at a location and date. + + Args: + location: The location to get the temperature for, in the format "City, State, Country". + date: The date to get the temperature for, in the format "Year-Month-Day". + unit: The unit to return the temperature in. Defaults to "celsius". (choices: ["celsius", "fahrenheit"]) + + Returns: + the temperature, the location, the date and the unit in a dict + """ + print(f"[DEBUG] get_temperature_date: {location}, {date}, {unit}") + return { + "temperature": 25.9, + "location": location, + "date": date, + "unit": unit, + } + + async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[ToolResponse, float, dict]: + try: + result = self.get_temperature_date(**parameters) + return ToolResponse(text=json.dumps(result)), 0, {} + except Exception as e: + return ToolResponse(text=str(e)), 0, {} + + +def test_tool_agent(init_config): + ray.init( + runtime_env={ + "env_vars": { + "TOKENIZERS_PARALLELISM": "true", + "NCCL_DEBUG": "WARN", + "VLLM_LOGGING_LEVEL": "INFO", + "VLLM_USE_V1": "1", + } + }, + ignore_reinit_error=True, + ) + + # =========================== 1. Init rollout manager =========================== + tool_config = { + "tools": [ + { + "class_name": "tests.experimental.agent_loop.test_basic_agent_loop.WeatherTool", + "config": {"type": "native"}, + }, + { + "class_name": "tests.experimental.agent_loop.test_basic_agent_loop.WeatherToolWithData", + "config": {"type": "native"}, + }, + ] + } + tool_config_path = "/tmp/tool_config.json" + with open(tool_config_path, "w") as f: + json.dump(tool_config, f) + + n = 2 + init_config.actor_rollout_ref.rollout.n = n + init_config.actor_rollout_ref.rollout.multi_turn.tool_config_path = tool_config_path + init_config.actor_rollout_ref.rollout.multi_turn.max_parallel_calls = 2 + init_config.actor_rollout_ref.rollout.calculate_log_probs = True + agent_loop_manager = init_agent_loop_manager(init_config) + + # =========================== 2. Generate sequences =========================== + raw_prompts = [ + [ + {"role": "user", "content": "How are you?"}, + ], + [ + {"role": "user", "content": "What's the temperature in Los Angeles now?"}, + ], + [ + {"role": "user", "content": "What's the temperature in New York now?"}, + ], + [ + { + "role": "system", + "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant.\n\n" + "Current Date: 2024-09-30", + }, + {"role": "user", "content": "What's the temperature in San Francisco now? How about tomorrow?"}, + ], + ] + batch = DataProto( + non_tensor_batch={ + "raw_prompt": np.array([np.array(prompt) for prompt in raw_prompts], dtype=object), + "agent_name": np.array(["tool_agent"] * len(raw_prompts)), + "data_source": np.array(["openai/gsm8k"] * len(raw_prompts)), + "reward_model": np.array([{"style": "rule", "ground_truth": "1.0"}] * len(raw_prompts)), + }, + ) + batch = batch.repeat(n) + result = agent_loop_manager.generate_sequences(prompts=batch) + assert len(result) == len(raw_prompts) * n + + # Check turns + num_turns = result.non_tensor_batch["__num_turns__"] + print(f"num_turns: {num_turns}") + for i in range(len(num_turns)): + if i // n == 0: + # [user, assistant] + assert num_turns[i] == 2 + else: + # [user, assistant, tool, assistant] + assert num_turns[i] == 4 + + # Check response_mask + tokenizer = hf_tokenizer(init_config.actor_rollout_ref.model.path) + responses = result.batch["responses"] + response_mask = result.batch["response_mask"] + attention_mask = result.batch["attention_mask"] + assert result.batch["rm_scores"].size(1) == responses.size(1) + assert responses.size() == response_mask.size(), f"{responses.size()} != {response_mask.size()}" + assert result.batch["rollout_log_probs"].size(1) == result.batch["responses"].size(1) + + response_length = response_mask.size(1) + for i in range(len(responses)): + # response with tool response + valid_tokens = responses[i][attention_mask[i][-response_length:].bool()] + response_with_obs = tokenizer.decode(valid_tokens) + + # response without tool response + valid_tokens = responses[i][response_mask[i].bool()] + response_without_obs = tokenizer.decode(valid_tokens) + + assert "" not in response_without_obs, ( + f"found in response: {response_without_obs}" + ) + assert "" not in response_without_obs, ( + f"found in response: {response_without_obs}" + ) + print("=========================") + print(response_with_obs) + print("---") + print(response_without_obs) + + print("Test passed!") + ray.shutdown() + + +def test_function_tool_agent(init_config): + """End-to-end coverage for ``rollout.multi_turn.function_tool_path``.""" + ray.init( + runtime_env={ + "env_vars": { + "TOKENIZERS_PARALLELISM": "true", + "NCCL_DEBUG": "WARN", + "VLLM_LOGGING_LEVEL": "INFO", + "VLLM_USE_V1": "1", + } + }, + ignore_reinit_error=True, + ) + + function_tool_path = os.path.join(os.path.dirname(__file__), "function_tool_examples.py") + + n = 2 + init_config.actor_rollout_ref.rollout.n = n + init_config.actor_rollout_ref.rollout.multi_turn.function_tool_path = function_tool_path + init_config.actor_rollout_ref.rollout.multi_turn.max_parallel_calls = 2 + agent_loop_manager = init_agent_loop_manager(init_config) + + raw_prompts = [ + [{"role": "user", "content": "Hi! Please reply with a one-word greeting."}], + [{"role": "user", "content": "What is the current temperature in Tokyo, in celsius?"}], + ] + batch = DataProto( + non_tensor_batch={ + "raw_prompt": np.array([np.array(prompt) for prompt in raw_prompts], dtype=object), + "agent_name": np.array(["tool_agent"] * len(raw_prompts)), + "data_source": np.array(["openai/gsm8k"] * len(raw_prompts)), + "reward_model": np.array([{"style": "rule", "ground_truth": "1.0"}] * len(raw_prompts)), + }, + ) + batch = batch.repeat(n) + result = agent_loop_manager.generate_sequences(prompts=batch) + assert len(result) == len(raw_prompts) * n + + num_turns = result.non_tensor_batch["__num_turns__"] + print(f"num_turns: {num_turns}") + greeting_idx = list(range(0, n)) + weather_idx = list(range(n, 2 * n)) + assert all(num_turns[i] == 2 for i in greeting_idx), ( + f"greeting prompt should not trigger a tool: {num_turns[greeting_idx]}" + ) + assert any(num_turns[i] == 4 for i in weather_idx), ( + f"expected at least one weather prompt to trigger a tool call (==4 turns); got {num_turns[weather_idx]}" + ) + + tokenizer = hf_tokenizer(init_config.actor_rollout_ref.model.path) + responses = result.batch["responses"] + response_mask = result.batch["response_mask"] + attention_mask = result.batch["attention_mask"] + response_length = response_mask.size(1) + saw_stub_value = False + for i in weather_idx: + if num_turns[i] != 4: + continue + valid_with_obs = responses[i][attention_mask[i][-response_length:].bool()] + full_response = tokenizer.decode(valid_with_obs) + if "17.3" in full_response: + saw_stub_value = True + break + assert saw_stub_value, ( + "expected the stub temperature '17.3' to appear in at least one " + "weather rollout that triggered a tool call; this is the only " + "value that proves the tool was actually invoked rather than " + "hallucinated by the model." + ) + + # Tool responses must not leak into the masked-for-loss response stream + # (same invariant ``test_tool_agent`` checks for native tools). + for i in range(len(responses)): + valid_tokens = responses[i][response_mask[i].bool()] + response_without_obs = tokenizer.decode(valid_tokens) + assert "" not in response_without_obs, ( + f"found in response: {response_without_obs}" + ) + assert "" not in response_without_obs, ( + f"found in response: {response_without_obs}" + ) + + print("Test passed!") + ray.shutdown() + + +@pytest.mark.asyncio +async def test_get_trajectory_info(): + """Tests the get_trajectory_info method.""" + # Initialize the class to set up class-level attributes + step = 10 + index = [1, 1, 3, 3] + expected_info = [ + {"step": step, "sample_index": 1, "rollout_n": 0, "validate": False}, + {"step": step, "sample_index": 1, "rollout_n": 1, "validate": False}, + {"step": step, "sample_index": 3, "rollout_n": 0, "validate": False}, + {"step": step, "sample_index": 3, "rollout_n": 1, "validate": False}, + ] + + trajectory_info = await get_trajectory_info(step, index, validate=False) + + assert trajectory_info == expected_info + + +# ────────────────────────────────────────────────────────────────────── +# GlobalRequestLoadBalancer unit tests (lightweight, no GPU required) +# ────────────────────────────────────────────────────────────────────── + + +@pytest.fixture(scope="module") +def ray_for_lb(): + ray.init(ignore_reinit_error=True) + yield + ray.shutdown() + + +class TestLoadBalancerRouting: + """Least-loaded selection.""" + + def test_distributes_across_servers(self, ray_for_lb): + lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None, "s2": None}) + servers = [ray.get(lb.acquire_server.remote(request_id=f"r{i}"))[0] for i in range(3)] + assert sorted(servers) == ["s0", "s1", "s2"] + + def test_new_requests_route_to_least_loaded(self, ray_for_lb): + lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None, "s2": None}) + # Load s0 with 3 inflight requests + ray.get(lb.acquire_server.remote(request_id="a"))[0] # -> s0 + ray.get(lb.acquire_server.remote(request_id="a"))[0] # sticky -> s0 + ray.get(lb.acquire_server.remote(request_id="a"))[0] # sticky -> s0 + # Load s1 with 1 inflight request + ray.get(lb.acquire_server.remote(request_id="b"))[0] # -> s1 + # s2 has 0 inflight, so next new request must go to s2 + s_new = ray.get(lb.acquire_server.remote(request_id="d"))[0] + assert s_new == "s2" + + def test_release_rebalances(self, ray_for_lb): + lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None}) + s0 = ray.get(lb.acquire_server.remote(request_id="r0"))[0] + s1 = ray.get(lb.acquire_server.remote(request_id="r1"))[0] + assert s0 != s1 + ray.get(lb.release_server.remote(server_id=s0)) + ray.get(lb.release_server.remote(server_id=s1)) + s2 = ray.get(lb.acquire_server.remote(request_id="r2"))[0] + s3 = ray.get(lb.acquire_server.remote(request_id="r3"))[0] + assert s2 != s3 + + def test_release_invalid_server_silently_ignored(self, ray_for_lb): + """Releasing a nonexistent server is silently ignored (hybrid-safe).""" + lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None}) + # Should not raise + ray.get(lb.release_server.remote(server_id="nonexistent")) + + def test_release_without_inflight_silently_ignored(self, ray_for_lb): + """Releasing a server with no inflight requests is silently ignored (hybrid-safe).""" + lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None}) + # Should not raise even though s1 has 0 inflight + ray.get(lb.release_server.remote(server_id="s1")) + + +class TestLoadBalancerStickySession: + """Request-level sticky session.""" + + def test_same_request_id_same_server(self, ray_for_lb): + lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None, "s2": None, "s3": None}) + s0 = ray.get(lb.acquire_server.remote(request_id="conv-abc"))[0] + ray.get(lb.release_server.remote(server_id=s0)) + s1 = ray.get(lb.acquire_server.remote(request_id="conv-abc"))[0] + assert s0 == s1 + + +class TestLoadBalancerHybrid: + """Dynamic server add/remove for hybrid scaling.""" + + def test_add_server(self, ray_for_lb): + lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None}) + ray.get(lb.add_servers.remote(servers={"s2": None})) + status = ray.get(lb.get_status.remote()) + assert "s2" in status["servers"] + assert status["servers"]["s2"] == 0 + + def test_remove_server_purges_handle(self, ray_for_lb): + lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None}) + ray.get(lb.remove_servers.remote(server_ids=["s1"])) + # remove_server now purges from both _inflight_requests and _servers + status = ray.get(lb.get_status.remote()) + assert "s1" not in status["servers"] + assert "s1" not in status["registered_handles"] + # New requests should only go to s0 + s = ray.get(lb.acquire_server.remote(request_id="r1"))[0] + assert s == "s0" + + def test_removed_server_invalidates_sticky_session(self, ray_for_lb): + """When a sticky session points to a removed server, cache is invalidated.""" + lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None}) + # Occupy s0 so that the sticky request is assigned to s1 + ray.get(lb.acquire_server.remote(request_id="occupy-s0"))[0] # -> s0 + # Pin request to s1 (least-loaded now) + s1 = ray.get(lb.acquire_server.remote(request_id="sticky-req"))[0] + assert s1 == "s1" + ray.get(lb.release_server.remote(server_id=s1)) + # Remove s1 + ray.get(lb.remove_servers.remote(server_ids=["s1"])) + # Sticky session should be invalidated and reroute to s0 + s_new = ray.get(lb.acquire_server.remote(request_id="sticky-req"))[0] + assert s_new == "s0" + + def test_remove_server_also_purges_registry(self, ray_for_lb): + """remove_servers atomically purges from both LB pool and handle registry.""" + lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None}) + ray.get(lb.remove_servers.remote(server_ids=["s1"])) + status = ray.get(lb.get_status.remote()) + # Both _inflight_requests and _servers are cleaned up (no separate cleanup step needed) + assert "s1" not in status["servers"] + assert "s1" not in status["registered_handles"] + + def test_get_all_servers_excludes_removed(self, ray_for_lb): + lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None, "s2": None}) + ray.get(lb.remove_servers.remote(server_ids=["s1"])) + all_servers = ray.get(lb.get_all_servers.remote()) + assert "s0" in all_servers + assert "s2" in all_servers + assert "s1" not in all_servers + + def test_no_available_servers_raises(self, ray_for_lb): + lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None}) + ray.get(lb.remove_servers.remote(server_ids=["s0", "s1"])) + with pytest.raises(ray.exceptions.RayTaskError, match="No available servers"): + ray.get(lb.acquire_server.remote(request_id="r1")) + + def test_add_server_readds_previously_removed(self, ray_for_lb): + """Re-adding a previously removed server makes it routable again.""" + lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None}) + ray.get(lb.remove_servers.remote(server_ids=["s1"])) + # s1 is removed, only s0 is available + assert ray.get(lb.acquire_server.remote(request_id="r1"))[0] == "s0" + # Re-add s1 + ray.get(lb.add_servers.remote(servers={"s1": None})) + # Now both s0 and s1 should be available + s = ray.get(lb.acquire_server.remote(request_id="r2"))[0] + assert s in ("s0", "s1") + + def test_get_inflight_count(self, ray_for_lb): + lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None}) + assert ray.get(lb.get_inflight_count.remote(server_id="s0")) == 0 + ray.get(lb.acquire_server.remote(request_id="r1"))[0] # -> s0 (least loaded) + assert ray.get(lb.get_inflight_count.remote(server_id="s0")) == 1 + + def test_get_status_reports_active_correctly(self, ray_for_lb): + lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None, "s2": None}) + ray.get(lb.remove_servers.remote(server_ids=["s1"])) + status = ray.get(lb.get_status.remote()) + assert status["active_servers"] == 2 # s0 and s2 + assert status["total_inflight"] == 0 diff --git a/verl/tests/experimental/agent_loop/test_call_tool_on_cpu.py b/verl/tests/experimental/agent_loop/test_call_tool_on_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..48ccf51398edca01e4e15db576b3dc0732123b67 --- /dev/null +++ b/verl/tests/experimental/agent_loop/test_call_tool_on_cpu.py @@ -0,0 +1,192 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for ToolAgentLoop._call_tool error handling (no GPU required). + +Tests that malformed tool calls return specific, actionable error messages +instead of generic exception strings. +""" + +import unittest +from dataclasses import dataclass, field +from typing import Any +from unittest.mock import MagicMock + +from verl.tools.schemas import ToolResponse + + +@dataclass +class FakeFunctionCall: + """Minimal FunctionCall for testing.""" + + name: str + arguments: str + + +@dataclass +class FakeAgentData: + """Minimal AgentData for testing.""" + + tools_kwargs: dict = field(default_factory=dict) + + +class FakeTool: + """A fake tool that succeeds.""" + + def __init__(self, name: str): + self.name = name + + async def create(self, create_kwargs=None): + return "instance_1", ToolResponse() + + async def execute(self, instance_id, parameters, **kwargs): + return ToolResponse(text=f"OK: {parameters}"), 1.0, {} + + async def release(self, instance_id): + pass + + +class FakeFailingTool(FakeTool): + """A fake tool that raises during execute.""" + + async def execute(self, instance_id, parameters, **kwargs): + raise RuntimeError("database connection failed") + + +class FakeLongResponseTool(FakeTool): + """A fake tool that returns a long response.""" + + def __init__(self, name: str, text: str): + super().__init__(name) + self.text = text + + async def execute(self, instance_id, parameters, **kwargs): + return ToolResponse(text=self.text), 1.0, {} + + +def _make_tool_agent_loop( + tools: dict[str, Any], + max_tool_response_length: int = 10000, + tool_response_truncate_side: str = "left", +): + """Create a minimal ToolAgentLoop instance with only the fields _call_tool needs.""" + from verl.experimental.agent_loop.tool_agent_loop import ToolAgentLoop + + mock = MagicMock(spec=ToolAgentLoop) + mock.tools = tools + mock.max_tool_response_length = max_tool_response_length + mock.tool_response_truncate_side = tool_response_truncate_side + # Bind the real _call_tool method to our mock + mock._call_tool = ToolAgentLoop._call_tool.__get__(mock, ToolAgentLoop) + return mock + + +class TestCallToolErrorHandling(unittest.IsolatedAsyncioTestCase): + """Test ToolAgentLoop._call_tool error handling for malformed tool calls.""" + + def setUp(self): + self.tools = { + "calculator": FakeTool("calculator"), + "search": FakeTool("search"), + } + self.loop = _make_tool_agent_loop(self.tools) + self.agent_data = FakeAgentData() + + async def test_valid_tool_call(self): + """Valid tool call should succeed.""" + tool_call = FakeFunctionCall(name="calculator", arguments='{"a": 3, "b": 5}') + response, reward, _ = await self.loop._call_tool(tool_call, {}, self.agent_data) + assert reward == 1.0 + assert "OK" in response.text + + async def test_unknown_function_name(self): + """Unknown function name should list available tools.""" + tool_call = FakeFunctionCall(name="calculater", arguments='{"a": 3}') + response, reward, _ = await self.loop._call_tool(tool_call, {}, self.agent_data) + assert reward == 0.0 + assert "Unknown function" in response.text + assert "calculater" in response.text + assert "calculator" in response.text + assert "search" in response.text + + async def test_invalid_json_arguments(self): + """Invalid JSON arguments should report parse error.""" + tool_call = FakeFunctionCall(name="calculator", arguments="{a: 3}") + response, reward, _ = await self.loop._call_tool(tool_call, {}, self.agent_data) + assert reward == 0.0 + assert "Invalid JSON" in response.text + assert "calculator" in response.text + + async def test_empty_arguments(self): + """Empty string arguments should report parse error.""" + tool_call = FakeFunctionCall(name="calculator", arguments="") + response, reward, _ = await self.loop._call_tool(tool_call, {}, self.agent_data) + assert reward == 0.0 + assert "Invalid JSON" in response.text + + async def test_none_arguments(self): + """None arguments should report error.""" + tool_call = FakeFunctionCall(name="calculator", arguments=None) + response, reward, _ = await self.loop._call_tool(tool_call, {}, self.agent_data) + assert reward == 0.0 + assert "Invalid JSON" in response.text + + async def test_tool_execution_error(self): + """Tool execution failure should include tool name in error.""" + tools = {"failing_tool": FakeFailingTool("failing_tool")} + loop = _make_tool_agent_loop(tools) + tool_call = FakeFunctionCall(name="failing_tool", arguments='{"query": "test"}') + response, reward, _ = await loop._call_tool(tool_call, {}, self.agent_data) + assert reward == 0.0 + assert "failing_tool" in response.text + assert "database connection failed" in response.text + + async def test_left_truncation_keeps_response_tail(self): + """Left truncation should drop the left side and preserve the response tail.""" + tool_response = ( + "Search results for capital of France:\n" + "1. Lyon is a major city with a long Roman history.\n" + "2. Marseille is a large port city in southern France.\n" + "3. The final retrieved snippet says the capital is Paris.\n" + "Final answer: Paris" + ) + tools = {"search": FakeLongResponseTool("search", tool_response)} + loop = _make_tool_agent_loop(tools, max_tool_response_length=19, tool_response_truncate_side="left") + tool_call = FakeFunctionCall(name="search", arguments="{}") + response, reward, _ = await loop._call_tool(tool_call, {}, self.agent_data) + assert reward == 1.0 + assert response.text.startswith("(truncated)...") + assert response.text.endswith("Final answer: Paris") + + async def test_right_truncation_keeps_response_head(self): + """Right truncation should drop the right side and preserve the response head.""" + tool_response = ( + "Search results for capital of France:\n" + "1. Lyon is a major city with a long Roman history.\n" + "2. Marseille is a large port city in southern France.\n" + "3. The final retrieved snippet says the capital is Paris.\n" + "Final answer: Paris" + ) + tools = {"search": FakeLongResponseTool("search", tool_response)} + loop = _make_tool_agent_loop(tools, max_tool_response_length=19, tool_response_truncate_side="right") + tool_call = FakeFunctionCall(name="search", arguments="{}") + response, reward, _ = await loop._call_tool(tool_call, {}, self.agent_data) + assert reward == 1.0 + assert response.text.startswith("Search results") + assert response.text.endswith("...(truncated)") + assert "Final answer: Paris" not in response.text + + +if __name__ == "__main__": + unittest.main() diff --git a/verl/tests/experimental/agent_loop/test_gpt_oss_tool_parser.py b/verl/tests/experimental/agent_loop/test_gpt_oss_tool_parser.py new file mode 100644 index 0000000000000000000000000000000000000000..a58c977a1b0d4eac2cbd542aab1fd0b8b691f1df --- /dev/null +++ b/verl/tests/experimental/agent_loop/test_gpt_oss_tool_parser.py @@ -0,0 +1,34 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import pytest +from transformers import AutoTokenizer + +from verl.experimental.agent_loop.tool_parser import GptOssToolParser + + +@pytest.mark.asyncio +@pytest.mark.skip(reason="local test only") +async def test_gpt_oss_tool_parser(): + example_text = """ +<|start|>assistant<|channel|>commentary to=functions.get_current_weather \ +<|constrain|>json<|message|>{"location": "Tokyo"}<|call|> +<|start|>functions.get_current_weather to=assistant<|channel|>commentary<|message|>\ +{ "temperature": 20, "sunny": true }<|end|>""" + tokenizer = AutoTokenizer.from_pretrained("openai/gpt-oss-20b") + response_ids = tokenizer.encode(example_text) + tool_parser = GptOssToolParser(tokenizer) + _, function_calls = await tool_parser.extract_tool_calls(response_ids) + assert len(function_calls) == 1 + assert function_calls[0].name == "get_current_weather" + assert function_calls[0].arguments == '{"location": "Tokyo"}' diff --git a/verl/tests/experimental/agent_loop/test_multi_modal.py b/verl/tests/experimental/agent_loop/test_multi_modal.py new file mode 100644 index 0000000000000000000000000000000000000000..6130a85ab4b44d29e42a4d8a6e38a200091ec272 --- /dev/null +++ b/verl/tests/experimental/agent_loop/test_multi_modal.py @@ -0,0 +1,441 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import json +import os +from typing import Any + +import numpy as np +import pytest +import ray +from omegaconf import DictConfig +from PIL import Image +from transformers.utils import get_json_schema + +from tests.experimental.agent_loop.agent_utils import init_agent_loop_manager +from verl.protocol import DataProto +from verl.tools.base_tool import BaseTool, OpenAIFunctionToolSchema +from verl.tools.schemas import ToolResponse +from verl.utils import hf_tokenizer + + +def parse_multi_modal_type(messages: list[dict]) -> str: + message = messages[-1] + if isinstance(message["content"], str): + return "text" + + for content in message["content"]: + if content["type"] == "image": + return "image" + elif content["type"] == "video": + return "video" + + return "text" + + +@pytest.fixture +def init_config() -> DictConfig: + from hydra import compose, initialize_config_dir + + with initialize_config_dir(config_dir=os.path.abspath("verl/trainer/config")): + config = compose( + config_name="ppo_trainer", + overrides=[ + "actor_rollout_ref.actor.use_dynamic_bsz=true", + # test sleep/wake_up with fsdp offload + "actor_rollout_ref.actor.fsdp_config.param_offload=True", + "actor_rollout_ref.actor.fsdp_config.optimizer_offload=True", + ], + ) + + model_path = os.path.expanduser("~/models/Qwen/Qwen2.5-VL-3B-Instruct") + config.actor_rollout_ref.model.path = model_path + config.actor_rollout_ref.rollout.name = os.environ["ROLLOUT_NAME"] + config.actor_rollout_ref.rollout.mode = "async" + config.actor_rollout_ref.rollout.enforce_eager = True + config.actor_rollout_ref.rollout.prompt_length = 10240 + config.actor_rollout_ref.rollout.response_length = 4096 + config.actor_rollout_ref.rollout.n = 4 + config.actor_rollout_ref.rollout.agent.num_workers = 2 + config.actor_rollout_ref.rollout.skip_tokenizer_init = True + + return config + + +class ImageGeneratorTool(BaseTool): + def generate_image(self, description: str, size: str = "256x256"): + """Generate a simple image based on description. + + Args: + description: The description of the image to generate. + size: The size of the image. Defaults to "256x256". (choices: ["256x256", "512x512"]) + + Returns: + A generated image + """ + print(f"[DEBUG] generate_image: {description}, {size}") + # Create a simple colored image for testing + width, height = map(int, size.split("x")) + + # Create different colors based on description + if "red" in description.lower(): + color = (255, 0, 0) + elif "blue" in description.lower(): + color = (0, 0, 255) + elif "green" in description.lower(): + color = (0, 255, 0) + else: + color = (128, 128, 128) # gray + + # Create image + image = Image.new("RGB", (width, height), color) + + # Add some pattern to make it more interesting + for i in range(0, width, 50): + for j in range(0, height, 50): + # Add white squares in a grid pattern + for x in range(i, min(i + 20, width)): + for y in range(j, min(j + 20, height)): + image.putpixel((x, y), (255, 255, 255)) + + return image + + def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema: + schema = get_json_schema(self.generate_image) + return OpenAIFunctionToolSchema(**schema) + + async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[ToolResponse, float, dict]: + try: + image = self.generate_image(**parameters) + # Return the PIL Image directly - the framework should handle the conversion + return ToolResponse(image=[image]), 0, {} + except Exception as e: + return ToolResponse(text=str(e)), 0, {} + + +@pytest.mark.flaky(reruns=3) +def test_multimodal_tool_agent(init_config): + """Test agent loop with multimodal tool that returns images using Qwen VL model.""" + ray.shutdown() + ray.init( + runtime_env={ + "env_vars": { + "TOKENIZERS_PARALLELISM": "true", + "NCCL_DEBUG": "WARN", + "VLLM_LOGGING_LEVEL": "INFO", + "VLLM_USE_V1": "1", + } + }, + ignore_reinit_error=True, + ) + + # Add custom chat template to enable tool calling support (same as recipe/deepeyes) + template_path = os.path.join(os.path.dirname(__file__), "qwen_vl_tool_chat_template.jinja2") + with open(template_path, encoding="utf-8") as f: + custom_chat_template = f.read() + + init_config.actor_rollout_ref.model.custom_chat_template = custom_chat_template + + # =========================== 1. Init rollout manager with image tool =========================== + tool_config = { + "tools": [ + { + "class_name": "tests.experimental.agent_loop.test_multi_modal.ImageGeneratorTool", + "config": {"type": "native"}, + }, + ] + } + tool_config_path = "/tmp/multimodal_tool_config.json" + with open(tool_config_path, "w") as f: + json.dump(tool_config, f) + + n = 2 + init_config.actor_rollout_ref.rollout.n = n + init_config.actor_rollout_ref.rollout.multi_turn.tool_config_path = tool_config_path + init_config.actor_rollout_ref.rollout.multi_turn.max_parallel_calls = 1 + init_config.actor_rollout_ref.rollout.multi_turn.max_user_turns = 1 + agent_loop_manager = init_agent_loop_manager(init_config) + + # =========================== 2. Generate sequences with multimodal prompts =========================== + raw_prompts = [ + [ + {"role": "user", "content": "How are you?"}, + ], + [ + { + "role": "user", + "content": [ + { + "type": "video", + "video": os.path.expanduser("~/models/hf_data/test-videos/space_woaudio.mp4"), + "min_pixels": 4 * 32 * 32, + "max_pixels": 256 * 32 * 32, + "total_pixels": 4096 * 32 * 32, + }, + { + "type": "text", + "text": "Describe this video. Then you must call the " + "image generator tool to generate a green image for me.", + }, + ], + }, + ], + [ + {"role": "user", "content": "Please generate a red image for me."}, + ], + [ + {"role": "user", "content": "Can you create a blue picture with size 512x512?"}, + ], + [ + { + "role": "system", + "content": ( + "You are Qwen VL, created by Alibaba Cloud. You are a helpful " + "assistant that can generate and analyze images." + ), + }, + {"role": "user", "content": "Generate a green landscape image and describe what you see in it."}, + ], + ] + + batch = DataProto( + non_tensor_batch={ + "raw_prompt": np.array([np.array(prompt) for prompt in raw_prompts], dtype=object), + "agent_name": np.array(["tool_agent"] * len(raw_prompts)), + "data_source": np.array(["openai/gsm8k"] * len(raw_prompts)), + "reward_model": np.array([{"style": "rule", "ground_truth": "1.0"}] * len(raw_prompts)), + }, + ) + batch = batch.repeat(n) + result = agent_loop_manager.generate_sequences(prompts=batch) + assert len(result) == len(raw_prompts) * n + + # Check turns + num_turns = result.non_tensor_batch["__num_turns__"] + multi_modal_inputs = result.non_tensor_batch["multi_modal_inputs"] + print(f"num_turns: {num_turns}") + for i in range(len(num_turns)): + multi_modal_type = parse_multi_modal_type(raw_prompts[i // n]) + if multi_modal_type == "video": + assert "pixel_values_videos" in multi_modal_inputs[i], f"Sample {i} should have pixel_values_videos" + assert "video_grid_thw" in multi_modal_inputs[i], f"Sample {i} should have video_grid_thw" + + if i // n == 0: + # First prompt: "How are you?" - should have 2 turns [user, assistant] + assert num_turns[i] == 2, f"Expected 2 turns but got {num_turns[i]} for sample {i}" + elif i // n == 1: + # TODO: prompt with video not generate tool call as expected + assert num_turns[i] == 2 or num_turns[i] == 4, ( + f"Expected 2 or 4 turns but got {num_turns[i]} for sample {i}" + ) + else: + # Tool-calling prompts should have 4 turns [user, assistant, tool, assistant] + assert num_turns[i] == 4, f"Expected 4 turns but got {num_turns[i]} for sample {i}" + assert "pixel_values" in multi_modal_inputs[i], f"Sample {i} should have pixel_values" + assert "image_grid_thw" in multi_modal_inputs[i], f"Sample {i} should have image_grid_thw" + + # Check that images were properly returned in the tool responses + tokenizer = hf_tokenizer(init_config.actor_rollout_ref.model.path) + responses = result.batch["responses"] + response_mask = result.batch["response_mask"] + attention_mask = result.batch["attention_mask"] + assert responses.size() == response_mask.size(), f"{responses.size()} != {response_mask.size()}" + response_length = response_mask.size(1) + + image_found_count = 0 + for i in range(len(responses)): + # response with tool response (including images) + valid_tokens = responses[i][attention_mask[i][-response_length:].bool()] + response_with_obs = tokenizer.decode(valid_tokens) + + # response without tool response + valid_tokens = responses[i][response_mask[i].bool()] + response_without_obs = tokenizer.decode(valid_tokens) + + # Check that tool responses were properly masked out from training + assert "" not in response_without_obs, ( + f"found in response: {response_without_obs}" + ) + assert "" not in response_without_obs, ( + f"found in response: {response_without_obs}" + ) + + # Check that images were included in the full response + if "" in response_with_obs or "image" in response_with_obs.lower(): + image_found_count += 1 + + print("=========================") + print("Response with tool observations:") + print(response_with_obs) + print("---") + print("Response without tool observations:") + print(response_without_obs) + + # Verify that tool-calling responses contained image-related content + print(f"Found {image_found_count} responses with image content out of {len(responses)}") + # We should have at least some image content from the tool-calling prompts + # Note: First prompt might not use tools, so we don't expect 100% image content + expected_tool_calls = sum(1 for i in range(len(num_turns)) if num_turns[i] == 4) + assert image_found_count >= 0, ( + f"No image-related content found, but expected at least some from {expected_tool_calls} tool calls" + ) + + print("Multimodal tool test passed!") + ray.shutdown() + + +def test_multimodal_single_turn_agent(init_config): + """Test single turn agent loop with multimodal inputs using Qwen VL model.""" + ray.shutdown() + ray.init( + runtime_env={ + "env_vars": { + "TOKENIZERS_PARALLELISM": "true", + "NCCL_DEBUG": "WARN", + "VLLM_LOGGING_LEVEL": "INFO", + "VLLM_USE_V1": "1", + } + }, + ignore_reinit_error=True, + ) + + # =========================== 1. Init rollout manager =========================== + n = 2 + init_config.actor_rollout_ref.rollout.n = n + init_config.actor_rollout_ref.rollout.multi_turn.max_parallel_calls = 1 + init_config.actor_rollout_ref.rollout.multi_turn.max_user_turns = 1 + agent_loop_manager = init_agent_loop_manager(init_config) + + # =========================== 2. Generate sequences with multimodal prompts =========================== + # Create a simple test image + test_image = Image.new("RGB", (256, 256), (100, 150, 200)) + test_image2 = Image.new("RGB", (512, 512), (100, 150, 200)) + + raw_prompts = [ + # text + [ + {"role": "user", "content": "Hello, how are you?"}, + ], + # image + [ + { + "role": "user", + "content": [ + {"type": "image", "image": test_image}, + {"type": "text", "text": "What color is this image?"}, + ], + }, + ], + # system + image + [ + { + "role": "system", + "content": "You are Qwen VL, created by Alibaba Cloud. You are a helpful assistant.", + }, + { + "role": "user", + "content": [ + {"type": "image", "image": test_image2}, + {"type": "text", "text": "Describe this image in detail."}, + ], + }, + ], + # video + [ + { + "role": "user", + "content": [ + { + "type": "video", + "video": os.path.expanduser("~/models/hf_data/test-videos/space_woaudio.mp4"), + "min_pixels": 4 * 32 * 32, + "max_pixels": 256 * 32 * 32, + "total_pixels": 4096 * 32 * 32, + }, + {"type": "text", "text": "Describe this video."}, + ], + }, + ], + ] + + batch = DataProto( + non_tensor_batch={ + "raw_prompt": np.array([np.array(prompt) for prompt in raw_prompts], dtype=object), + "agent_name": np.array(["single_turn_agent"] * len(raw_prompts)), + "data_source": np.array(["openai/gsm8k"] * len(raw_prompts)), + "reward_model": np.array([{"style": "rule", "ground_truth": "1.0"}] * len(raw_prompts)), + }, + ) + + batch = batch.repeat(n) + result = agent_loop_manager.generate_sequences(prompts=batch) + assert len(result) == len(raw_prompts) * n + + # Check turns - all should be single turn (2: user + assistant) + num_turns = result.non_tensor_batch["__num_turns__"] + print(f"num_turns: {num_turns}") + for i in range(len(num_turns)): + assert num_turns[i] == 2, f"Expected 2 turns but got {num_turns[i]} for sample {i}" + + # Verify responses + tokenizer = hf_tokenizer(init_config.actor_rollout_ref.model.path) + prompts = result.batch["prompts"] + responses = result.batch["responses"] + response_mask = result.batch["response_mask"] + input_ids = result.batch["input_ids"] + position_ids = result.batch["position_ids"] + multi_modal_inputs = result.non_tensor_batch["multi_modal_inputs"] + assert responses.size() == response_mask.size(), f"{responses.size()} != {response_mask.size()}" + assert position_ids.size() == (input_ids.size(0), 4, input_ids.size(1)) # (batch_size, 4, seq_len) + + # Check for image pads in prompts + image_pad_count = 0 + for i in range(len(prompts)): + prompt_ids = prompts[i][prompts[i] != tokenizer.pad_token_id].tolist() + prompt_text = tokenizer.decode(prompt_ids) + + # Check if this sample should have image pads (samples with index 1 and 2 in each repeat have images) + sample_idx = i // n + has_image_pad = "<|image_pad|>" in prompt_text or "<|vision_start|>" in prompt_text + + print("=========================") + print(f"Sample {i} (original prompt index: {sample_idx}):") + print(f"Prompt length: {len(prompt_ids)} tokens") + print(f"Has image_pad: {has_image_pad}") + + # Check multi-modal type + multi_modal_type = parse_multi_modal_type(raw_prompts[sample_idx]) + + if multi_modal_type == "text": + assert len(multi_modal_inputs[i]) == 0, f"Sample {i} should not have multi-modal inputs" + elif multi_modal_type == "image": + assert "pixel_values" in multi_modal_inputs[i], f"Sample {i} should have pixel_values" + assert "image_grid_thw" in multi_modal_inputs[i], f"Sample {i} should have image_grid_thw" + else: + assert "pixel_values_videos" in multi_modal_inputs[i], f"Sample {i} should have pixel_values_videos" + assert "video_grid_thw" in multi_modal_inputs[i], f"Sample {i} should have video_grid_thw" + + # Show first 200 chars of prompt + print(f"Prompt text (first 200 chars): {prompt_text[:200]}...") + + for i in range(len(responses)): + valid_tokens = responses[i][response_mask[i].bool()] + response_text = tokenizer.decode(valid_tokens) + print(f"Sample {i} response: {response_text[:100]}...") + + # Verify that we found image pads in multimodal samples + expected_multimodal_samples = 2 * n # 2 prompts with images, repeated n times + print(f"\nFound {image_pad_count} samples with image_pad out of {expected_multimodal_samples} expected") + + print("Single turn multimodal test passed!") + ray.shutdown() diff --git a/verl/tests/experimental/agent_loop/test_standalone_rollout.py b/verl/tests/experimental/agent_loop/test_standalone_rollout.py new file mode 100644 index 0000000000000000000000000000000000000000..a46138c702cf2aaa52142d74741ed5defa2762eb --- /dev/null +++ b/verl/tests/experimental/agent_loop/test_standalone_rollout.py @@ -0,0 +1,151 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os + +import pytest +import ray +from omegaconf import DictConfig +from openai import AsyncOpenAI, OpenAI + +from tests.experimental.agent_loop.agent_utils import init_agent_loop_manager +from verl.checkpoint_engine import CheckpointEngineManager +from verl.utils import omega_conf_to_dataclass +from verl.workers.rollout.llm_server import LLMServerManager + + +@pytest.fixture +def init_config() -> DictConfig: + from hydra import compose, initialize_config_dir + + with initialize_config_dir(config_dir=os.path.abspath("verl/trainer/config")): + config = compose(config_name="ppo_trainer") + + config.trainer.n_gpus_per_node = 4 + config.trainer.nnodes = 2 + config.actor_rollout_ref.actor.use_dynamic_bsz = True + config.actor_rollout_ref.model.path = os.path.expanduser("~/models/Qwen/Qwen2.5-1.5B-Instruct") + config.actor_rollout_ref.rollout.name = os.environ["ROLLOUT_NAME"] + config.actor_rollout_ref.rollout.mode = "async" + config.actor_rollout_ref.rollout.skip_tokenizer_init = False + + return config + + +@pytest.mark.asyncio +@pytest.mark.parametrize("tp_size", [2, 4]) +async def test_standalone_rollout(init_config, tp_size): + """Test standalone rollout single node and multi nodes.""" + ray.init( + runtime_env={ + "env_vars": { + "TOKENIZERS_PARALLELISM": "true", + "NCCL_DEBUG": "WARN", + "VLLM_LOGGING_LEVEL": "INFO", + "VLLM_USE_V1": "1", + "NCCL_P2P_DISABLE": "1", # disable p2p in L20 + } + } + ) + + init_config.actor_rollout_ref.rollout.nnodes = init_config.trainer.nnodes + init_config.actor_rollout_ref.rollout.tensor_model_parallel_size = tp_size + num_replicas = (init_config.trainer.n_gpus_per_node * init_config.trainer.nnodes) // tp_size + + # create standalone rollout server + llm_server_manager = await LLMServerManager.create( + config=init_config, + worker_group=None, + rollout_resource_pool=None, + ) + + server_addresses = llm_server_manager.get_addresses() + assert len(server_addresses) == num_replicas + + os.environ.pop("HTTPS_PROXY", None) + os.environ.pop("HTTP_PROXY", None) + os.environ.pop("NO_PROXY", None) + + client = AsyncOpenAI( + api_key="123-abc", + base_url=f"http://{server_addresses[0]}/v1", + ) + + completion = await client.chat.completions.create( + model=init_config.actor_rollout_ref.model.path, + messages=[{"role": "user", "content": "What can you do?"}], + ) + print(completion.choices[0].message.content) + + ray.shutdown() + + +@pytest.mark.skip(reason="local test only") +def test_hybrid_rollout_with_ep(init_config): + """Test hybrid rollout with expert parallelism, DP=2, TP=4, EP=8.""" + ray.init( + runtime_env={ + "env_vars": { + "TOKENIZERS_PARALLELISM": "true", + "NCCL_DEBUG": "WARN", + "VLLM_LOGGING_LEVEL": "INFO", + "VLLM_USE_V1": "1", + } + } + ) + + model_path = os.path.expanduser("~/models/Qwen/Qwen3-30B-A3B-Instruct-2507") + init_config.actor_rollout_ref.model.path = model_path + + # parallelism config + init_config.actor_rollout_ref.rollout.tensor_model_parallel_size = 2 + init_config.actor_rollout_ref.rollout.data_parallel_size = 4 + init_config.actor_rollout_ref.rollout.expert_parallel_size = 8 + + # 1. init hybrid worker: FSDP+rollout + # - build FSDP model and optimizer + # - offload FSDP model and optimizer, build rollout + # - sleep rollout and load FSDP model and optimizer + agent_loop_manager = init_agent_loop_manager(init_config) + checkpoint_manager = CheckpointEngineManager( + config=omega_conf_to_dataclass(init_config.actor_rollout_ref.rollout.checkpoint_engine), + trainer=agent_loop_manager.worker_group, + replicas=agent_loop_manager.rollout_replicas, + ) + checkpoint_manager.sleep_replicas() + checkpoint_manager.update_weights() + + # 3. test async openai call + server_address = agent_loop_manager.server_addresses[0] + client = OpenAI( + api_key="123-abc", + base_url=f"http://{server_address}/v1", + ) + + smapling_params = { + "temperature": 1.0, + "top_p": 1.0, + "max_tokens": 512, + } + + response = client.chat.completions.create( + model=model_path, + messages=[{"role": "user", "content": "What can you do?"}], + **smapling_params, + ) + + completion = response.choices[0].message.content + print(f"response: {completion}") + + print("Test passed!") + ray.shutdown() diff --git a/verl/tests/experimental/fully_async_policy/__init__.py b/verl/tests/experimental/fully_async_policy/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1cd1e8433dffa0b3ba420be3e346f4f5cd062014 --- /dev/null +++ b/verl/tests/experimental/fully_async_policy/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/verl/tests/experimental/fully_async_policy/test_async_genrm_config_on_cpu.py b/verl/tests/experimental/fully_async_policy/test_async_genrm_config_on_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..a081f17af4ae97c172ecd7d8c90992ac21d15dcf --- /dev/null +++ b/verl/tests/experimental/fully_async_policy/test_async_genrm_config_on_cpu.py @@ -0,0 +1,106 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for async GenRM config validation logic (no GPU required). + +Tests the config parsing and assertion logic added to FullyAsyncRollouter +to support GenRM/DisRM in fully async training mode. +""" + +import unittest + +import pytest +from omegaconf import OmegaConf + +from verl.trainer.ppo.utils import need_reward_model + + +def _make_config(reward_model_enable=False, enable_resource_pool=False): + """Create a minimal OmegaConf config for testing reward model settings.""" + return OmegaConf.create( + { + "reward": { + "reward_model": { + "enable": reward_model_enable, + "enable_resource_pool": enable_resource_pool, + "n_gpus_per_node": 2, + "nnodes": 1, + "model_path": "dummy/model", + "rollout": { + "name": "vllm", + "tensor_model_parallel_size": 1, + "gpu_memory_utilization": 0.5, + "skip_tokenizer_init": False, + }, + }, + "custom_reward_function": { + "path": None, + "name": None, + }, + }, + } + ) + + +class TestNeedRewardModel(unittest.TestCase): + """Test that need_reward_model correctly reads config.""" + + def test_rm_disabled(self): + config = _make_config(reward_model_enable=False) + assert need_reward_model(config) is False + + def test_rm_enabled(self): + config = _make_config(reward_model_enable=True) + assert need_reward_model(config) is True + + +class TestAsyncRollouterRMAssert(unittest.TestCase): + """Test the assertion logic that enforces standalone mode for async RM. + + This replicates the validation logic from FullyAsyncRollouter.__init__ + without instantiating the full class (which requires Ray, worker groups, etc.). + """ + + @staticmethod + def _validate_async_rm_config(config): + """Replicate the RM validation logic from FullyAsyncRollouter.__init__.""" + use_rm = need_reward_model(config) + if use_rm: + assert config.reward.reward_model.enable_resource_pool, ( + "GenRM/DisRM in fully async mode requires standalone mode (enable_resource_pool=True). " + "Colocate mode is not supported because async rollout never pauses." + ) + return use_rm + + def test_rm_disabled_passes(self): + """use_rm=False should pass regardless of enable_resource_pool.""" + config = _make_config(reward_model_enable=False, enable_resource_pool=False) + use_rm = self._validate_async_rm_config(config) + assert use_rm is False + + def test_rm_enabled_standalone_passes(self): + """use_rm=True + enable_resource_pool=True (standalone) should pass.""" + config = _make_config(reward_model_enable=True, enable_resource_pool=True) + use_rm = self._validate_async_rm_config(config) + assert use_rm is True + + def test_rm_enabled_colocate_fails(self): + """use_rm=True + enable_resource_pool=False (colocate) should assert.""" + config = _make_config(reward_model_enable=True, enable_resource_pool=False) + with pytest.raises(AssertionError, match="standalone mode"): + self._validate_async_rm_config(config) + + +if __name__ == "__main__": + unittest.main() diff --git a/verl/tests/experimental/reward_loop/reward_fn.py b/verl/tests/experimental/reward_loop/reward_fn.py new file mode 100644 index 0000000000000000000000000000000000000000..6a751c512d047e6240d01bd3129f4d7d654d3c39 --- /dev/null +++ b/verl/tests/experimental/reward_loop/reward_fn.py @@ -0,0 +1,100 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import os + +import aiohttp +from openai.types.chat import ChatCompletion +from transformers import PreTrainedTokenizer + +GRM_PROMPT_TEMPLATE = """ +You are given a problem and a proposed solution. + +Problem: +{problem} + +Solution: +{solution} + +Please evaluate how well the solution addresses the problem. +Give a score from 1 to 10, where: +- 1 means the solution is completely irrelevant or incorrect. +- 5 means the solution is partially correct but incomplete or not well reasoned. +- 10 means the solution is fully correct, well-reasoned, and directly solves the problem. + +Only output the score as a single number (integer). +""".strip() + + +async def chat_complete(router_address: str, chat_complete_request: dict): + url = f"http://{router_address}/v1/chat/completions" + try: + timeout = aiohttp.ClientTimeout(total=None) + session = aiohttp.ClientSession(timeout=timeout) + async with session.post(url, json=chat_complete_request) as resp: + output = await resp.text() + output = json.loads(output) + return ChatCompletion(**output) + except Exception as e: + raise e + finally: + await session.close() + + +async def compute_score_gsm8k( + data_source: str, + solution_str: str, + ground_truth: str, + extra_info: dict, + reward_router_address: str, + reward_model_tokenizer: PreTrainedTokenizer, +): + """Compute the reward score.""" + + grm_prompt = GRM_PROMPT_TEMPLATE.format(problem=extra_info["question"], solution=solution_str) + messages = [{"role": "user", "content": grm_prompt}] + sampling_params = {"temperature": 0.7, "top_p": 0.8, "max_tokens": 4096} + model_name = os.environ.get("GENRM_MODEL_NAME", os.path.expanduser("~/models/Qwen/Qwen2.5-1.5B-Instruct")) + chat_complete_request = { + "messages": messages, + "model": model_name, + **sampling_params, + } + result = await chat_complete( + router_address=reward_router_address, + chat_complete_request=chat_complete_request, + ) + grm_response = result.choices[0].message.content + try: + score = int(grm_response.split("\n\n")[-1].strip()) + except Exception: + score = 0 + return {"score": score, "acc": score == 10, "genrm_response": grm_response} + + +def compute_score_math_verify( + data_source: str, + solution_str: str, + ground_truth: str, + extra_info: dict, + **kwargs, +): + """Compute the reward score.""" + from verl.utils.reward_score.math_verify import compute_score + + return compute_score( + model_output=solution_str, + ground_truth=ground_truth, + ) diff --git a/verl/tests/experimental/reward_loop/test_agent_reward_loop_colocate.py b/verl/tests/experimental/reward_loop/test_agent_reward_loop_colocate.py new file mode 100644 index 0000000000000000000000000000000000000000..33b8fbf9c3337c12e0bd197f0bca4b9310ec8b76 --- /dev/null +++ b/verl/tests/experimental/reward_loop/test_agent_reward_loop_colocate.py @@ -0,0 +1,173 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os + +import ray +from hydra import compose, initialize_config_dir +from torchdata.stateful_dataloader import StatefulDataLoader +from transformers import AutoTokenizer + +from verl.checkpoint_engine import CheckpointEngineManager +from verl.experimental.agent_loop import AgentLoopManager +from verl.experimental.reward_loop import RewardLoopManager +from verl.protocol import DataProto +from verl.single_controller.ray import RayClassWithInitArgs, RayWorkerGroup +from verl.trainer.main_ppo import create_rl_sampler +from verl.trainer.ppo.ray_trainer import ResourcePoolManager +from verl.utils import omega_conf_to_dataclass +from verl.utils.dataset.rl_dataset import RLHFDataset, collate_fn +from verl.utils.device import get_device_name +from verl.workers.engine_workers import ActorRolloutRefWorker +from verl.workers.rollout.llm_server import LLMServerManager + + +def test_agent_reward_loop_standalone(): + ray.init( + runtime_env={ + "env_vars": { + "TOKENIZERS_PARALLELISM": "true", + "NCCL_DEBUG": "WARN", + "VLLM_LOGGING_LEVEL": "INFO", + "VLLM_USE_V1": "1", + } + } + ) + with initialize_config_dir(config_dir=os.path.abspath("verl/trainer/config")): + config = compose(config_name="ppo_trainer") + + rollout_model_path = os.path.expanduser("~/models/Qwen/Qwen2.5-0.5B-Instruct") + reward_model_path = os.path.expanduser("~/models/Qwen/Qwen2.5-1.5B-Instruct") + + # actor_rollout_ref config + config.data.return_raw_chat = True + config.data.max_prompt_length = 1024 + config.data.max_response_length = 4096 + config.actor_rollout_ref.model.path = rollout_model_path + config.actor_rollout_ref.actor.use_dynamic_bsz = True + config.actor_rollout_ref.rollout.name = os.getenv("ROLLOUT_NAME", "vllm") + config.actor_rollout_ref.rollout.mode = "async" + config.actor_rollout_ref.rollout.tensor_model_parallel_size = 2 + config.actor_rollout_ref.rollout.gpu_memory_utilization = 0.8 + config.actor_rollout_ref.rollout.enforce_eager = True + config.actor_rollout_ref.rollout.prompt_length = 1024 + config.actor_rollout_ref.rollout.response_length = 4096 + config.actor_rollout_ref.rollout.skip_tokenizer_init = True + config.trainer.n_gpus_per_node = 8 + config.trainer.nnodes = 1 + + config.reward.reward_manager.name = "dapo" + config.reward.reward_model.enable = True + config.reward.reward_model.enable_resource_pool = False + config.reward.reward_model.n_gpus_per_node = 8 + config.reward.reward_model.model_path = reward_model_path + config.reward.reward_model.rollout.name = os.getenv("ROLLOUT_NAME", "vllm") + config.reward.reward_model.rollout.gpu_memory_utilization = 0.8 + config.reward.reward_model.rollout.tensor_model_parallel_size = 2 + config.reward.reward_model.rollout.skip_tokenizer_init = False + config.reward.reward_model.rollout.prompt_length = 5120 + config.reward.reward_model.rollout.response_length = 4096 + config.reward.custom_reward_function.path = "tests/experimental/reward_loop/reward_fn.py" + config.reward.custom_reward_function.name = "compute_score_gsm8k" + + # 1. init reward model manager + # The unified model-engine ActorRolloutRefWorker supports both sync and async rollout modes. + actor_rollout_cls = ActorRolloutRefWorker + global_pool_id = "global_pool" + resource_pool_spec = { + global_pool_id: [config.trainer.n_gpus_per_node] * config.trainer.nnodes, + } + resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=None) + resource_pool_manager.create_resource_pool() + resource_pool = resource_pool_manager.resource_pool_dict[global_pool_id] + actor_rollout_cls = RayClassWithInitArgs( + cls=ray.remote(actor_rollout_cls), config=config.actor_rollout_ref, role="actor_rollout" + ) + actor_rollout_wg = RayWorkerGroup( + resource_pool=resource_pool, ray_cls_with_init=actor_rollout_cls, device_name=get_device_name() + ) + actor_rollout_wg.init_model() + + llm_server_manager = LLMServerManager.create(config=config, worker_group=actor_rollout_wg) + agent_loop_manager = AgentLoopManager.create( + config=config, + llm_client=llm_server_manager.get_client(), + ) + # sleep rollout replicas + checkpoint_manager = CheckpointEngineManager( + config=omega_conf_to_dataclass(config.actor_rollout_ref.rollout.checkpoint_engine), + trainer=actor_rollout_wg, + replicas=llm_server_manager.get_replicas(), + ) + checkpoint_manager.sleep_replicas() + reward_loop_manager = RewardLoopManager(config, rm_resource_pool=resource_pool) + + # 2. init test data + local_folder = os.path.expanduser("~/data/gsm8k/") + + data_files = [os.path.join(local_folder, "train.parquet")] + tokenizer = AutoTokenizer.from_pretrained(rollout_model_path) + + dataset = RLHFDataset( + data_files=data_files, + tokenizer=tokenizer, + config=config.data, + processor=None, + ) + + batch_size = 64 + sampler = create_rl_sampler(config.data, dataset) + dataloader = StatefulDataLoader( + dataset=dataset, + batch_size=batch_size, + num_workers=config.data.dataloader_num_workers, + drop_last=True, + collate_fn=collate_fn, + sampler=sampler, + ) + + # 3. generate responses + batch_dict = next(iter(dataloader)) + batch = DataProto.from_single_dict(batch_dict) + + def _get_gen_batch(batch: DataProto) -> DataProto: + reward_keys = set({"data_source", "reward_model", "extra_info", "uid"}) & batch.non_tensor_batch.keys() + + # pop those keys for generation + batch_keys_to_pop = [] + non_tensor_batch_keys_to_pop = set(batch.non_tensor_batch.keys()) - reward_keys + gen_batch = batch.pop( + batch_keys=batch_keys_to_pop, + non_tensor_batch_keys=list(non_tensor_batch_keys_to_pop), + ) + + # For agent loop, we need reward model keys to compute score. + gen_batch.non_tensor_batch.update(batch.non_tensor_batch) + + return gen_batch + + # wake up rollout replicas via update_weight + checkpoint_manager.update_weights() + gen_batch = _get_gen_batch(batch) + gen_batch = agent_loop_manager.generate_sequences(gen_batch) + checkpoint_manager.sleep_replicas() + + batch = batch.union(gen_batch) + rm_outputs = reward_loop_manager.compute_rm_score(batch) + + for output in rm_outputs[:5]: + print(output.non_tensor_batch) + + print("done") + + ray.shutdown() diff --git a/verl/tests/experimental/reward_loop/test_agent_reward_loop_standalone.py b/verl/tests/experimental/reward_loop/test_agent_reward_loop_standalone.py new file mode 100644 index 0000000000000000000000000000000000000000..026ddc168029d5f863dd792b1bc506372e11426b --- /dev/null +++ b/verl/tests/experimental/reward_loop/test_agent_reward_loop_standalone.py @@ -0,0 +1,122 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os + +import ray +from hydra import compose, initialize_config_dir +from torchdata.stateful_dataloader import StatefulDataLoader + +from verl.experimental.agent_loop import AgentLoopManager +from verl.experimental.reward_loop import RewardLoopManager +from verl.protocol import DataProto +from verl.trainer.main_ppo import create_rl_sampler +from verl.utils import hf_tokenizer +from verl.utils.dataset.rl_dataset import RLHFDataset, collate_fn +from verl.workers.rollout.llm_server import LLMServerManager + + +def test_agent_reward_loop_standalone(): + ray.init( + runtime_env={ + "env_vars": { + "TOKENIZERS_PARALLELISM": "true", + "NCCL_DEBUG": "WARN", + "VLLM_LOGGING_LEVEL": "INFO", + "VLLM_USE_V1": "1", + } + } + ) + with initialize_config_dir(config_dir=os.path.abspath("verl/trainer/config")): + config = compose(config_name="ppo_trainer") + + rollout_model_path = os.path.expanduser("~/models/Qwen/Qwen2.5-0.5B-Instruct") + reward_model_path = os.path.expanduser("~/models/Qwen/Qwen2.5-1.5B-Instruct") + + # actor_rollout_ref config + config.data.return_raw_chat = True + config.data.max_prompt_length = 1024 + config.data.max_response_length = 4096 + config.actor_rollout_ref.model.path = rollout_model_path + config.actor_rollout_ref.actor.use_dynamic_bsz = True + config.actor_rollout_ref.rollout.name = os.getenv("ROLLOUT_NAME", "vllm") + config.actor_rollout_ref.rollout.mode = "async" + config.actor_rollout_ref.rollout.tensor_model_parallel_size = 2 + config.actor_rollout_ref.rollout.gpu_memory_utilization = 0.9 + config.actor_rollout_ref.rollout.enforce_eager = True + config.actor_rollout_ref.rollout.prompt_length = 1024 + config.actor_rollout_ref.rollout.response_length = 4096 + config.actor_rollout_ref.rollout.skip_tokenizer_init = True + config.actor_rollout_ref.rollout.nnodes = 1 + config.trainer.n_gpus_per_node = 4 + config.trainer.nnodes = 1 + + config.reward.reward_manager.name = "dapo" + config.reward.reward_model.enable = True + config.reward.reward_model.enable_resource_pool = True + config.reward.reward_model.n_gpus_per_node = 4 + config.reward.reward_model.nnodes = 1 + config.reward.reward_model.model_path = reward_model_path + config.reward.reward_model.rollout.name = os.getenv("ROLLOUT_NAME", "vllm") + config.reward.reward_model.rollout.gpu_memory_utilization = 0.9 + config.reward.reward_model.rollout.tensor_model_parallel_size = 2 + config.reward.reward_model.rollout.skip_tokenizer_init = False + config.reward.reward_model.rollout.prompt_length = 5120 + config.reward.reward_model.rollout.response_length = 4096 + config.reward.custom_reward_function.path = "tests/experimental/reward_loop/reward_fn.py" + config.reward.custom_reward_function.name = "compute_score_gsm8k" + + # 1. init reward model manager + reward_loop_manager = RewardLoopManager(config) + llm_server_manager = LLMServerManager.create(config=config) + agent_loop_manager = AgentLoopManager.create( + config=config, + llm_client=llm_server_manager.get_client(), + reward_loop_worker_handles=reward_loop_manager.reward_loop_workers, + ) + + # 2. init test data + local_folder = os.path.expanduser("~/data/gsm8k/") + data_files = [os.path.join(local_folder, "train.parquet")] + tokenizer = hf_tokenizer(rollout_model_path) + + dataset = RLHFDataset( + data_files=data_files, + tokenizer=tokenizer, + config=config.data, + processor=None, + ) + + batch_size = 64 + sampler = create_rl_sampler(config.data, dataset) + dataloader = StatefulDataLoader( + dataset=dataset, + batch_size=batch_size, + num_workers=config.data.dataloader_num_workers, + drop_last=True, + collate_fn=collate_fn, + sampler=sampler, + ) + + # 3. generate responses + batch_dict = next(iter(dataloader)) + batch = DataProto.from_single_dict(batch_dict) + + # standalone reward model should wake up for agent_reward_loop + gen_batch = agent_loop_manager.generate_sequences(prompts=batch) + + rm_scores = gen_batch.batch["rm_scores"] + sample_scores = rm_scores.sum(dim=1) + print(sample_scores) + + ray.shutdown() diff --git a/verl/tests/experimental/reward_loop/test_async_token_bucket_on_cpu.py b/verl/tests/experimental/reward_loop/test_async_token_bucket_on_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..70906fb51bd3848aa9e925261f2f5c4f71718e17 --- /dev/null +++ b/verl/tests/experimental/reward_loop/test_async_token_bucket_on_cpu.py @@ -0,0 +1,267 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +import time + +import pytest + +from verl.experimental.reward_loop.reward_manager.limited import AsyncTokenBucket + + +class TestAsyncTokenBucket: + """Unit tests for AsyncTokenBucket rate limiter.""" + + @pytest.mark.asyncio + async def test_basic_acquire(self): + """Test basic token acquisition.""" + bucket = AsyncTokenBucket(rate_limit=10.0, max_tokens=10.0) + + # Should be able to acquire tokens immediately when bucket is full + start = time.time() + await bucket.acquire(5.0) + elapsed = time.time() - start + + assert elapsed < 0.1, "Initial acquire should be immediate" + assert bucket.tokens == pytest.approx(5.0, abs=0.1) + + @pytest.mark.asyncio + async def test_refill_mechanism(self): + """Test that tokens refill over time.""" + bucket = AsyncTokenBucket(rate_limit=10.0, max_tokens=10.0) + + # Consume all tokens + await bucket.acquire(10.0) + assert bucket.tokens == pytest.approx(0.0, abs=0.1) + + # Wait for refill (should get ~5 tokens in 0.5 seconds at 10 tokens/sec) + await asyncio.sleep(0.5) + + # Try to acquire 4 tokens (should succeed without waiting) + start = time.time() + await bucket.acquire(4.0) + elapsed = time.time() - start + + assert elapsed < 0.1, "Acquire should be quick after refill" + + @pytest.mark.asyncio + async def test_waiting_for_tokens(self): + """Test that acquire waits when insufficient tokens available.""" + bucket = AsyncTokenBucket(rate_limit=10.0, max_tokens=10.0) + + # Consume all tokens + await bucket.acquire(10.0) + + # Try to acquire more tokens (should wait ~0.5 seconds for 5 tokens) + start = time.time() + await bucket.acquire(5.0) + elapsed = time.time() - start + + # Should wait approximately 0.5 seconds (5 tokens / 10 tokens per second) + assert 0.4 < elapsed < 0.7, f"Expected ~0.5s wait, got {elapsed:.3f}s" + + @pytest.mark.asyncio + async def test_max_tokens_cap(self): + """Test that tokens don't exceed max_tokens capacity.""" + bucket = AsyncTokenBucket(rate_limit=10.0, max_tokens=5.0) + + # Wait for potential overflow + await asyncio.sleep(1.0) + + # Tokens should be capped at max_tokens + await bucket.acquire(1.0) + + # After 1 second at 10 tokens/sec, should have max_tokens (5.0) + # After acquiring 1, should have 4.0 remaining + assert bucket.tokens <= 5.0, "Tokens should not exceed max_tokens" + + @pytest.mark.asyncio + async def test_fractional_tokens(self): + """Test acquiring fractional tokens.""" + bucket = AsyncTokenBucket(rate_limit=100.0, max_tokens=100.0) + + # Acquire fractional amounts + await bucket.acquire(0.5) + await bucket.acquire(1.5) + await bucket.acquire(2.3) + + assert bucket.tokens == pytest.approx(100.0 - 0.5 - 1.5 - 2.3, abs=0.1) + + @pytest.mark.asyncio + async def test_concurrent_acquires(self): + """Test multiple concurrent acquire operations.""" + bucket = AsyncTokenBucket(rate_limit=10.0, max_tokens=10.0) + + async def acquire_task(num_tokens: float, task_id: int): + await bucket.acquire(num_tokens) + return task_id + + # Launch 5 concurrent tasks, each acquiring 3 tokens (15 total) + # Bucket only has 10, so some will need to wait + start = time.time() + tasks = [acquire_task(3.0, i) for i in range(5)] + results = await asyncio.gather(*tasks) + elapsed = time.time() - start + + # Should take at least 0.5 seconds to refill 5 tokens + # (15 needed - 10 available) / 10 tokens per second = 0.5 seconds + assert elapsed >= 0.4, f"Expected >=0.4s for concurrent acquires, got {elapsed:.3f}s" + assert len(results) == 5, "All tasks should complete" + + @pytest.mark.asyncio + async def test_high_rate_limit(self): + """Test with high rate limit (simulating high-throughput scenarios).""" + bucket = AsyncTokenBucket(rate_limit=1000.0, max_tokens=1000.0) + + # Rapidly acquire tokens + start = time.time() + for _ in range(100): + await bucket.acquire(10.0) # 1000 tokens total + elapsed = time.time() - start + + # Should complete in approximately 1 second + assert elapsed < 1.5, f"High rate limit test took too long: {elapsed:.3f}s" + + @pytest.mark.asyncio + async def test_zero_initial_state(self): + """Test that bucket starts with full tokens.""" + bucket = AsyncTokenBucket(rate_limit=10.0, max_tokens=10.0) + + assert bucket.tokens == 10.0, "Bucket should start full" + assert bucket.last_update is None, "last_update should be None initially" + + # After first acquire, last_update should be set + await bucket.acquire(1.0) + assert bucket.last_update is not None, "last_update should be set after acquire" + + @pytest.mark.asyncio + async def test_rate_limit_accuracy(self): + """Test rate limit accuracy over time.""" + rate = 50.0 # 50 tokens per second + bucket = AsyncTokenBucket(rate_limit=rate, max_tokens=rate) + + # Consume all tokens and measure refill time for 25 tokens + await bucket.acquire(50.0) + + start = time.time() + await bucket.acquire(25.0) + elapsed = time.time() - start + + expected_time = 25.0 / rate # 0.5 seconds + # Allow 20% margin for timing inaccuracy + assert abs(elapsed - expected_time) < expected_time * 0.2, f"Expected ~{expected_time:.3f}s, got {elapsed:.3f}s" + + @pytest.mark.asyncio + async def test_sequential_acquires(self): + """Test sequential acquire operations.""" + bucket = AsyncTokenBucket(rate_limit=20.0, max_tokens=20.0) + + # Sequential acquires without waiting + await bucket.acquire(5.0) + await bucket.acquire(5.0) + await bucket.acquire(5.0) + await bucket.acquire(5.0) + + # Bucket should be empty + assert bucket.tokens == pytest.approx(0.0, abs=0.1) + + # Next acquire should wait + start = time.time() + await bucket.acquire(10.0) + elapsed = time.time() - start + + assert elapsed >= 0.4, "Should wait for token refill" + + @pytest.mark.asyncio + async def test_default_max_tokens(self): + """Test that max_tokens defaults to rate_limit.""" + bucket = AsyncTokenBucket(rate_limit=15.0) + + assert bucket.max_tokens == 15.0, "max_tokens should default to rate_limit" + assert bucket.tokens == 15.0, "Initial tokens should equal max_tokens" + + @pytest.mark.asyncio + async def test_single_token_acquire(self): + """Test default acquire of 1 token.""" + bucket = AsyncTokenBucket(rate_limit=10.0, max_tokens=10.0) + + await bucket.acquire() # Default num_tokens=1.0 + + assert bucket.tokens == pytest.approx(9.0, abs=0.1) + + @pytest.mark.asyncio + async def test_large_token_acquire(self): + """Test acquiring more tokens than bucket capacity.""" + bucket = AsyncTokenBucket(rate_limit=10.0, max_tokens=10.0) + + # Try to acquire 50 tokens (5x capacity) + start = time.time() + await bucket.acquire(50.0) + elapsed = time.time() - start + + # Should wait for: (50 - 10) / 10 = 4 seconds + assert 3.5 < elapsed < 5.0, f"Expected ~4s wait for large acquire, got {elapsed:.3f}s" + + @pytest.mark.asyncio + async def test_thread_safety_with_lock(self): + """Test that lock prevents race conditions.""" + bucket = AsyncTokenBucket(rate_limit=100.0, max_tokens=100.0) + results = [] + + async def acquire_and_record(): + await bucket.acquire(10.0) + results.append(1) + + # Launch many concurrent tasks + tasks = [acquire_and_record() for _ in range(10)] + await asyncio.gather(*tasks) + + # All tasks should complete + assert len(results) == 10, "All tasks should complete successfully" + + # Bucket should have consumed exactly 100 tokens + assert bucket.tokens == pytest.approx(0.0, abs=0.5) + + @pytest.mark.asyncio + async def test_multiple_wait_cycles(self): + """Test multiple wait cycles in the acquire loop.""" + bucket = AsyncTokenBucket(rate_limit=10.0, max_tokens=10.0) + + # Consume all tokens + await bucket.acquire(10.0) + + # Acquire tokens that require multiple refill cycles + start = time.time() + await bucket.acquire(15.0) + elapsed = time.time() - start + + # Should wait for 15 tokens / 10 tokens per second = 1.5 seconds + assert 1.3 < elapsed < 1.8, f"Expected ~1.5s for multiple refill cycles, got {elapsed:.3f}s" + + @pytest.mark.asyncio + async def test_rapid_small_acquires(self): + """Test many rapid small acquisitions.""" + bucket = AsyncTokenBucket(rate_limit=100.0, max_tokens=100.0) + + start = time.time() + for _ in range(50): + await bucket.acquire(2.0) # 100 tokens total + elapsed = time.time() - start + + # Should complete quickly since we're within capacity + assert elapsed < 0.5, f"Rapid small acquires took too long: {elapsed:.3f}s" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/verl/tests/experimental/reward_loop/test_math_verify.py b/verl/tests/experimental/reward_loop/test_math_verify.py new file mode 100644 index 0000000000000000000000000000000000000000..f1d64c7a59e9c018eb90080f39e012f470a658ee --- /dev/null +++ b/verl/tests/experimental/reward_loop/test_math_verify.py @@ -0,0 +1,100 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os + +import ray +from hydra import compose, initialize_config_dir +from torchdata.stateful_dataloader import StatefulDataLoader +from transformers import AutoTokenizer + +from tests.experimental.agent_loop.agent_utils import init_agent_loop_manager +from verl.protocol import DataProto +from verl.trainer.main_ppo import create_rl_sampler +from verl.utils.dataset.rl_dataset import RLHFDataset, collate_fn + + +def test_agent_reward_loop_standalone(): + ray.init( + runtime_env={ + "env_vars": { + "TOKENIZERS_PARALLELISM": "true", + "NCCL_DEBUG": "WARN", + "VLLM_LOGGING_LEVEL": "INFO", + "VLLM_USE_V1": "1", + } + } + ) + with initialize_config_dir(config_dir=os.path.abspath("verl/trainer/config")): + config = compose(config_name="ppo_trainer") + + rollout_model_path = os.path.expanduser("~/models/Qwen/Qwen2.5-1.5B-Instruct") + + # actor_rollout_ref config + config.data.return_raw_chat = True + config.data.max_prompt_length = 1024 + config.data.max_response_length = 4096 + config.actor_rollout_ref.model.path = rollout_model_path + config.actor_rollout_ref.actor.use_dynamic_bsz = True + config.actor_rollout_ref.rollout.name = os.getenv("ROLLOUT_NAME", "vllm") + config.actor_rollout_ref.rollout.mode = "async" + config.actor_rollout_ref.rollout.tensor_model_parallel_size = 2 + config.actor_rollout_ref.rollout.gpu_memory_utilization = 0.9 + config.actor_rollout_ref.rollout.enforce_eager = True + config.actor_rollout_ref.rollout.prompt_length = 2048 + config.actor_rollout_ref.rollout.response_length = 4096 + config.actor_rollout_ref.rollout.skip_tokenizer_init = True + config.trainer.n_gpus_per_node = 8 + config.trainer.nnodes = 1 + + config.reward.reward_manager.name = "remote" + config.reward.num_workers = 2 + config.reward.custom_reward_function.path = "tests/experimental/reward_loop/reward_fn.py" + config.reward.custom_reward_function.name = "compute_score_math_verify" + + # 1. init reward model manager + agent_loop_manager = init_agent_loop_manager(config) + + # 2. init test data + local_folder = os.path.expanduser("~/data/math/") + data_files = [os.path.join(local_folder, "train.parquet")] + tokenizer = AutoTokenizer.from_pretrained(rollout_model_path) + + dataset = RLHFDataset( + data_files=data_files, + tokenizer=tokenizer, + config=config.data, + processor=None, + ) + + batch_size = 64 + sampler = create_rl_sampler(config.data, dataset) + dataloader = StatefulDataLoader( + dataset=dataset, + batch_size=batch_size, + num_workers=config.data.dataloader_num_workers, + drop_last=True, + collate_fn=collate_fn, + sampler=sampler, + ) + + # 3. generate responses + batch_dict = next(iter(dataloader)) + batch = DataProto.from_single_dict(batch_dict) + gen_batch = agent_loop_manager.generate_sequences(prompts=batch) + + rm_scores = gen_batch.batch["rm_scores"] + accuracy = rm_scores.sum(dim=-1).mean() + print(accuracy) + + ray.shutdown() diff --git a/verl/tests/experimental/reward_loop/test_rate_limited_reward_manager_on_cpu.py b/verl/tests/experimental/reward_loop/test_rate_limited_reward_manager_on_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..a6f2666423ace5df8d50ada4316755330e2d93c4 --- /dev/null +++ b/verl/tests/experimental/reward_loop/test_rate_limited_reward_manager_on_cpu.py @@ -0,0 +1,505 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +import os.path +import time + +import numpy as np +import pytest +import torch +from omegaconf import DictConfig +from transformers import AutoTokenizer + +from verl import DataProto +from verl.experimental.reward_loop.reward_manager.limited import RateLimitedRewardManager + + +# Mock API reward functions for testing +class MockAPICounter: + """Shared counter to track API calls across tests.""" + + def __init__(self): + self.call_count = 0 + self.call_times = [] + self.lock = asyncio.Lock() + + async def record_call(self): + async with self.lock: + self.call_count += 1 + self.call_times.append(time.time()) + + def reset(self): + self.call_count = 0 + self.call_times.clear() + + def get_rate_per_second(self, window_start: float = None): + """Calculate API call rate over a time window.""" + if window_start is None: + if not self.call_times: + return 0.0 + window_start = self.call_times[0] + + if not self.call_times: + return 0.0 + + window_end = self.call_times[-1] + duration = window_end - window_start + + if duration <= 0: + return 0.0 + + calls_in_window = sum(1 for t in self.call_times if t >= window_start) + return calls_in_window / duration + + +# Global counter instance +api_counter = MockAPICounter() + + +def mock_sync_reward_function( + data_source: str, solution_str: str, ground_truth: str, extra_info: dict, **kwargs +) -> float: + """Synchronous mock reward function that simulates API call.""" + # Simulate API processing time + time.sleep(0.01) + + # Simple scoring logic + score = 1.0 if solution_str.strip() == ground_truth.strip() else 0.0 + return score + + +async def mock_async_reward_function( + data_source: str, solution_str: str, ground_truth: str, extra_info: dict, **kwargs +) -> float: + """Asynchronous mock reward function that simulates API call.""" + # Record API call for rate tracking + await api_counter.record_call() + + # Simulate async API call (e.g., HTTP request) + await asyncio.sleep(0.01) + + # Simple scoring logic + score = 1.0 if solution_str.strip() == ground_truth.strip() else 0.0 + return score + + +async def mock_slow_api_function( + data_source: str, solution_str: str, ground_truth: str, extra_info: dict, **kwargs +) -> float: + """Slow mock API function for timeout testing.""" + await asyncio.sleep(2.0) # Simulate slow API + return 0.5 + + +async def mock_failing_api_function( + data_source: str, solution_str: str, ground_truth: str, extra_info: dict, **kwargs +) -> float: + """Mock API function that raises an exception.""" + await api_counter.record_call() + raise ValueError("Simulated API error") + + +async def mock_dict_result_function( + data_source: str, solution_str: str, ground_truth: str, extra_info: dict, **kwargs +) -> dict: + """Mock API function that returns dict result.""" + await api_counter.record_call() + await asyncio.sleep(0.01) + + correct = solution_str.strip() == ground_truth.strip() + return {"score": 1.0 if correct else 0.0, "correct": correct, "reasoning": "Mock reasoning"} + + +def create_test_data_proto(tokenizer, response_text: str, ground_truth: str, data_source: str = "test"): + """Helper to create DataProto for testing.""" + response_ids = tokenizer.encode(response_text, add_special_tokens=False) + response_tensor = torch.tensor([response_ids], dtype=torch.long) + attention_mask = torch.ones_like(response_tensor) + + data = DataProto.from_dict( + { + "responses": response_tensor, + "attention_mask": attention_mask, + } + ) + + # Wrap non-tensor values in numpy arrays to match batch dimension + data.non_tensor_batch = { + "data_source": np.array([data_source], dtype=object), + "reward_model": np.array([{"ground_truth": ground_truth}], dtype=object), + } + + return data + + +class TestRateLimitedRewardManager: + """Integration tests for RateLimitedRewardManager with mock API functions.""" + + @pytest.fixture(autouse=True) + def setup_and_teardown(self): + """Reset global state before each test.""" + api_counter.reset() + # Reset class state + RateLimitedRewardManager._class_initialized = False + RateLimitedRewardManager._semaphore = None + RateLimitedRewardManager._rpm_limiter = None + RateLimitedRewardManager._tpm_limiter = None + yield + # Cleanup + api_counter.reset() + + @pytest.fixture + def tokenizer(self): + """Load a simple tokenizer for testing.""" + return AutoTokenizer.from_pretrained(os.path.expanduser("~/models/Qwen/Qwen2.5-0.5B-Instruct")) + + @pytest.mark.asyncio + async def test_basic_reward_computation(self, tokenizer): + """Test basic reward computation without rate limiting.""" + config = DictConfig({"reward": {"max_concurrent": 10, "timeout": 10.0}}) + + RateLimitedRewardManager.init_class(config, tokenizer) + manager = RateLimitedRewardManager(config=config, tokenizer=tokenizer, compute_score=mock_async_reward_function) + + # Create test data + data = create_test_data_proto(tokenizer, "correct answer", "correct answer") + + # Compute reward + result = await manager.run_single(data) + + assert "reward_score" in result + assert result["reward_score"] == 1.0 + assert api_counter.call_count == 1 + + @pytest.mark.asyncio + async def test_rpm_rate_limiting(self, tokenizer): + """Test request per minute (RPM) rate limiting.""" + # Set RPM limit to 60 (1 request per second) + config = DictConfig( + { + "reward": { + "max_concurrent": 10, + "max_rpm": 60, # 1 request per second + "timeout": 10.0, + } + } + ) + + RateLimitedRewardManager.init_class(config, tokenizer) + manager = RateLimitedRewardManager(config=config, tokenizer=tokenizer, compute_score=mock_async_reward_function) + + # Create test data + data = create_test_data_proto(tokenizer, "answer", "answer") + + # Make 3 requests - should be rate limited + start_time = time.time() + + results = [] + for _ in range(3): + result = await manager.run_single(data) + results.append(result) + + elapsed = time.time() - start_time + + # Should take at least ~2 seconds for 3 requests at 1 req/sec + assert elapsed >= 1.8, f"RPM limiting failed: {elapsed:.3f}s for 3 requests" + assert all(r["reward_score"] == 1.0 for r in results) + assert api_counter.call_count == 3 + + @pytest.mark.asyncio + async def test_tpm_rate_limiting(self, tokenizer): + """Test tokens per minute (TPM) rate limiting.""" + # Set TPM limit to 6000 (100 tokens per second) + # With 2000 tokens per request, that's 0.05 req/sec or 20 seconds per request + config = DictConfig( + { + "reward": { + "max_concurrent": 10, + "max_tpm": 6000, # 100 tokens per second + "estimated_tokens_per_request": 2000, # Each request = 2000 tokens + "timeout": 30.0, + } + } + ) + + RateLimitedRewardManager.init_class(config, tokenizer) + manager = RateLimitedRewardManager(config=config, tokenizer=tokenizer, compute_score=mock_async_reward_function) + + data = create_test_data_proto(tokenizer, "answer", "answer") + + # Make 2 requests + start_time = time.time() + + result1 = await manager.run_single(data) + result2 = await manager.run_single(data) + + elapsed = time.time() - start_time + + # First request: consumes 2000 tokens (immediate) + # Second request: needs 2000 tokens, waits for refill + # Wait time: 2000 tokens / 100 tokens per second = 20 seconds + assert elapsed >= 18.0, f"TPM limiting failed: {elapsed:.3f}s for 2 requests" + assert result1["reward_score"] == 1.0 + assert result2["reward_score"] == 1.0 + + @pytest.mark.asyncio + async def test_concurrency_limiting(self, tokenizer): + """Test concurrent request limiting.""" + config = DictConfig( + { + "reward": { + "max_concurrent": 2, # Only 2 concurrent requests + "timeout": 10.0, + } + } + ) + + RateLimitedRewardManager.init_class(config, tokenizer) + manager = RateLimitedRewardManager(config=config, tokenizer=tokenizer, compute_score=mock_async_reward_function) + + data = create_test_data_proto(tokenizer, "answer", "answer") + + # Launch 5 concurrent requests + start_time = time.time() + + tasks = [manager.run_single(data) for _ in range(5)] + results = await asyncio.gather(*tasks) + + elapsed = time.time() - start_time + + # All should succeed + assert len(results) == 5 + assert all(r["reward_score"] == 1.0 for r in results) + + # With concurrency=2 and 0.01s per request, should take at least 0.03s + # (3 batches: 2+2+1) + assert elapsed >= 0.02, f"Concurrency limiting may not be working: {elapsed:.3f}s" + + @pytest.mark.asyncio + async def test_timeout_handling(self, tokenizer): + """Test timeout handling for slow API.""" + config = DictConfig( + { + "reward": { + "max_concurrent": 10, + "timeout": 0.5, # 500ms timeout + } + } + ) + + RateLimitedRewardManager.init_class(config, tokenizer) + manager = RateLimitedRewardManager(config=config, tokenizer=tokenizer, compute_score=mock_slow_api_function) + + data = create_test_data_proto(tokenizer, "answer", "answer") + + # Should timeout and return 0.0 + result = await manager.run_single(data) + + assert result["reward_score"] == 0.0 + assert result["reward_extra_info"].get("timeout") is True + assert result["reward_extra_info"].get("acc") == 0.0 + + @pytest.mark.asyncio + async def test_error_handling(self, tokenizer): + """Test error handling for failing API.""" + config = DictConfig({"reward": {"max_concurrent": 10, "timeout": 10.0}}) + + RateLimitedRewardManager.init_class(config, tokenizer) + manager = RateLimitedRewardManager(config=config, tokenizer=tokenizer, compute_score=mock_failing_api_function) + + data = create_test_data_proto(tokenizer, "answer", "answer") + + # Should catch exception and return 0.0 + result = await manager.run_single(data) + + assert result["reward_score"] == 0.0 + assert "error" in result["reward_extra_info"] + assert "Simulated API error" in result["reward_extra_info"]["error"] + assert result["reward_extra_info"].get("acc") == 0.0 + assert api_counter.call_count == 1 + + @pytest.mark.asyncio + async def test_dict_result_format(self, tokenizer): + """Test handling of dict return format from reward function.""" + config = DictConfig({"reward": {"max_concurrent": 10, "timeout": 10.0}}) + + RateLimitedRewardManager.init_class(config, tokenizer) + manager = RateLimitedRewardManager(config=config, tokenizer=tokenizer, compute_score=mock_dict_result_function) + + data = create_test_data_proto(tokenizer, "correct", "correct") + + result = await manager.run_single(data) + + assert result["reward_score"] == 1.0 + assert result["reward_extra_info"]["score"] == 1.0 + assert result["reward_extra_info"]["correct"] is True + assert result["reward_extra_info"]["reasoning"] == "Mock reasoning" + + @pytest.mark.asyncio + async def test_sync_reward_function(self, tokenizer): + """Test that synchronous reward functions work correctly.""" + config = DictConfig({"reward": {"max_concurrent": 10, "timeout": 10.0}}) + + RateLimitedRewardManager.init_class(config, tokenizer) + manager = RateLimitedRewardManager(config=config, tokenizer=tokenizer, compute_score=mock_sync_reward_function) + + data = create_test_data_proto(tokenizer, "answer", "answer") + + result = await manager.run_single(data) + + assert result["reward_score"] == 1.0 + assert manager.is_async_reward_score is False + + @pytest.mark.asyncio + async def test_combined_rate_limits(self, tokenizer): + """Test all three rate limiting layers together.""" + config = DictConfig( + { + "reward": { + "max_concurrent": 2, + "max_rpm": 120, # 2 requests per second + "max_tpm": 12000, # 200 tokens per second + "estimated_tokens_per_request": 100, # 0.5 seconds per request + "timeout": 10.0, + } + } + ) + + RateLimitedRewardManager.init_class(config, tokenizer) + manager = RateLimitedRewardManager(config=config, tokenizer=tokenizer, compute_score=mock_async_reward_function) + + data = create_test_data_proto(tokenizer, "answer", "answer") + + # Make 6 requests to exceed burst capacity (RPM bucket starts with 2 tokens) + start_time = time.time() + + tasks = [manager.run_single(data) for _ in range(6)] + results = await asyncio.gather(*tasks) + + elapsed = time.time() - start_time + + # Bucket starts with 2 RPM tokens and 200 TPM tokens + # First 2 requests: use burst capacity (2 RPM tokens, 200 TPM tokens) + # Next 4 requests: need 4 RPM tokens (wait 2 seconds) and 400 TPM tokens (wait 2 seconds) + # Limiting factor: RPM at 2 seconds + assert elapsed >= 1.8, f"Combined rate limiting: {elapsed:.3f}s" + assert all(r["reward_score"] == 1.0 for r in results) + assert api_counter.call_count == 6 + + @pytest.mark.asyncio + async def test_correct_vs_incorrect_answers(self, tokenizer): + """Test scoring of correct vs incorrect answers.""" + config = DictConfig({"reward": {"max_concurrent": 10, "timeout": 10.0}}) + + RateLimitedRewardManager.init_class(config, tokenizer) + manager = RateLimitedRewardManager(config=config, tokenizer=tokenizer, compute_score=mock_async_reward_function) + + # Test correct answer + data_correct = create_test_data_proto(tokenizer, "right answer", "right answer") + result_correct = await manager.run_single(data_correct) + + # Test incorrect answer + data_incorrect = create_test_data_proto(tokenizer, "wrong answer", "right answer") + result_incorrect = await manager.run_single(data_incorrect) + + assert result_correct["reward_score"] == 1.0 + assert result_incorrect["reward_score"] == 0.0 + + @pytest.mark.asyncio + async def test_high_throughput(self, tokenizer): + """Test high throughput with many concurrent requests.""" + config = DictConfig( + { + "reward": { + "max_concurrent": 20, + "max_rpm": 6000, # 100 requests per second + "timeout": 10.0, + } + } + ) + + RateLimitedRewardManager.init_class(config, tokenizer) + manager = RateLimitedRewardManager(config=config, tokenizer=tokenizer, compute_score=mock_async_reward_function) + + data = create_test_data_proto(tokenizer, "answer", "answer") + + # Launch 200 concurrent requests (more than burst capacity of 100) + start_time = time.time() + + tasks = [manager.run_single(data) for _ in range(200)] + results = await asyncio.gather(*tasks) + + elapsed = time.time() - start_time + + assert len(results) == 200 + assert all(r["reward_score"] == 1.0 for r in results) + + # Bucket starts with 100 tokens (burst capacity) + # First 100 requests: use burst capacity instantly + # Next 100 requests: need to wait for refill at 100 tokens/sec = 1 second minimum + # Total time should be at least 1 second + assert elapsed >= 0.9, f"Should take at least 0.9s for rate limiting, took {elapsed:.3f}s" + + # Calculate actual rate over the time window + actual_rate = api_counter.call_count / elapsed + + # Average rate should not significantly exceed 100 req/sec + # Allow some burst overhead due to initial capacity + assert actual_rate <= 200, f"Rate limiting failed: {actual_rate:.1f} req/sec (max 200)" + + @pytest.mark.asyncio + async def test_class_initialization_once(self, tokenizer): + """Test that class initialization only happens once.""" + config = DictConfig({"reward": {"max_concurrent": 5, "timeout": 10.0}}) + + # Initialize multiple times + RateLimitedRewardManager.init_class(config, tokenizer) + first_semaphore = RateLimitedRewardManager._semaphore + + RateLimitedRewardManager.init_class(config, tokenizer) + second_semaphore = RateLimitedRewardManager._semaphore + + # Should be the same object + assert first_semaphore is second_semaphore + + @pytest.mark.asyncio + async def test_extra_info_handling(self, tokenizer): + """Test that extra_info is properly passed to reward function.""" + received_extra_info = {} + + async def mock_reward_with_extra_info( + data_source: str, solution_str: str, ground_truth: str, extra_info: dict, **kwargs + ): + received_extra_info.update(extra_info) + return 1.0 + + config = DictConfig({"reward": {"max_concurrent": 10, "timeout": 10.0}}) + + RateLimitedRewardManager.init_class(config, tokenizer) + manager = RateLimitedRewardManager( + config=config, tokenizer=tokenizer, compute_score=mock_reward_with_extra_info + ) + + data = create_test_data_proto(tokenizer, "answer", "answer") + data.non_tensor_batch["extra_info"] = np.array([{"custom_field": "test_value"}], dtype=object) + + await manager.run_single(data) + + assert "custom_field" in received_extra_info + assert received_extra_info["custom_field"] == "test_value" + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/verl/tests/experimental/reward_loop/test_reward_model_disrm.py b/verl/tests/experimental/reward_loop/test_reward_model_disrm.py new file mode 100644 index 0000000000000000000000000000000000000000..2e42be92ab19294ff207e504f95daa8913986917 --- /dev/null +++ b/verl/tests/experimental/reward_loop/test_reward_model_disrm.py @@ -0,0 +1,155 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os + +import ray +import torch +from hydra import compose, initialize_config_dir + +from verl.experimental.reward_loop import RewardLoopManager +from verl.protocol import DataProto +from verl.utils import hf_tokenizer +from verl.utils.model import compute_position_id_with_mask +from verl.utils.tokenizer import normalize_token_ids + + +def create_data_samples(tokenizer) -> DataProto: + convs = [ + [ + { + "role": "user", + "content": "What is the range of the numeric output of a sigmoid node in a neural network?", + }, + {"role": "assistant", "content": "Between -1 and 1."}, + ], + [ + { + "role": "user", + "content": "What is the range of the numeric output of a sigmoid node in a neural network?", + }, + {"role": "assistant", "content": "Between 0 and 1."}, + ], + [ + {"role": "user", "content": "What is the capital of Australia?"}, + { + "role": "assistant", + "content": "Canberra is the capital city of Australia.", + }, + ], + [ + {"role": "user", "content": "What is the capital of Australia?"}, + { + "role": "assistant", + "content": "Sydney is the capital of Australia.", + }, + ], + ] + raw_prompt = [conv[:1] for conv in convs] + data_source = ["gsm8k"] * len(convs) + reward_info = [{"ground_truth": "Not Used"}] * len(convs) + extra_info = [{"question": conv[0]["content"]} for conv in convs] + + prompt_length, response_length = 1024, 4096 + pad_token_id = tokenizer.pad_token_id + prompts, responses, input_ids, attention_masks = [], [], [], [] + for conv in convs: + prompt_tokens = normalize_token_ids(tokenizer.apply_chat_template(conv[:1], tokenize=True)) + response_tokens = normalize_token_ids(tokenizer.apply_chat_template(conv, tokenize=True))[len(prompt_tokens) :] + + padded_prompt = [pad_token_id] * (prompt_length - len(prompt_tokens)) + prompt_tokens + padded_response = response_tokens + [pad_token_id] * (response_length - len(response_tokens)) + attention_mask = ( + [0] * (prompt_length - len(prompt_tokens)) + + [1] * len(prompt_tokens) + + [1] * len(response_tokens) + + [0] * (response_length - len(response_tokens)) + ) + prompts.append(torch.tensor(padded_prompt)) + responses.append(torch.tensor(padded_response)) + input_ids.append(torch.tensor(padded_prompt + padded_response)) + attention_masks.append(torch.tensor(attention_mask)) + + prompts = torch.stack(prompts) + responses = torch.stack(responses) + input_ids = torch.stack(input_ids) + attention_masks = torch.stack(attention_masks) + position_ids = compute_position_id_with_mask(attention_masks) + + data = DataProto.from_dict( + tensors={ + "prompts": prompts, + "responses": responses, + "input_ids": input_ids, + "attention_mask": attention_masks, + "position_ids": position_ids, + }, + non_tensors={ + "data_source": data_source, + "reward_model": reward_info, + "raw_prompt": raw_prompt, + "extra_info": extra_info, + }, + ) + return data, convs + + +def test_reward_model_manager(): + ray.init( + runtime_env={ + "env_vars": { + "TOKENIZERS_PARALLELISM": "true", + "NCCL_DEBUG": "WARN", + "VLLM_LOGGING_LEVEL": "INFO", + "VLLM_USE_V1": "1", + } + } + ) + with initialize_config_dir(config_dir=os.path.abspath("verl/trainer/config")): + config = compose(config_name="ppo_trainer") + + rollout_model_name = os.path.expanduser("~/models/Qwen/Qwen2.5-1.5B-Instruct") + reward_model_name = os.path.expanduser("~/models/Skywork/Skywork-Reward-V2-Llama-3.2-1B") + + config.actor_rollout_ref.model.path = rollout_model_name + config.reward.num_workers = 1 + config.reward.reward_manager.name = "dapo" + config.reward.reward_model.enable = True + config.reward.reward_model.enable_resource_pool = True + config.reward.reward_model.n_gpus_per_node = 8 + config.reward.reward_model.nnodes = 1 + config.reward.reward_model.model_path = reward_model_name + config.reward.reward_model.rollout.name = os.getenv("ROLLOUT_NAME", "vllm") + config.reward.reward_model.rollout.gpu_memory_utilization = 0.9 + config.reward.reward_model.rollout.tensor_model_parallel_size = 2 + config.reward.reward_model.rollout.skip_tokenizer_init = False + config.reward.reward_model.rollout.prompt_length = 2048 + config.reward.reward_model.rollout.response_length = 4096 + + # 1. init reward model manager + reward_loop_manager = RewardLoopManager(config) + + # 2. init test data + rollout_tokenizer = hf_tokenizer(rollout_model_name) + data, convs = create_data_samples(rollout_tokenizer) + + # 3. generate responses + outputs = reward_loop_manager.compute_rm_score(data) + + for idx, (conv, output) in enumerate(zip(convs, outputs, strict=True)): + print(f"Problem {idx}:\n{conv[0]['content']}\n") + print(f"AI Solution {idx}:\n{conv[1]['content']}\n") + print(f"DisRM Score {idx}:\n{output.batch['rm_scores'].sum(dim=-1).item()}\n") + print("=" * 50 + "\n") + + ray.shutdown() diff --git a/verl/tests/experimental/reward_loop/test_reward_model_genrm.py b/verl/tests/experimental/reward_loop/test_reward_model_genrm.py new file mode 100644 index 0000000000000000000000000000000000000000..d86e044d1294adffda35c0e9b8bcb3443088dad5 --- /dev/null +++ b/verl/tests/experimental/reward_loop/test_reward_model_genrm.py @@ -0,0 +1,158 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os + +import ray +import torch +from hydra import compose, initialize_config_dir + +from verl.experimental.reward_loop import RewardLoopManager +from verl.protocol import DataProto +from verl.utils import hf_tokenizer +from verl.utils.model import compute_position_id_with_mask +from verl.utils.tokenizer import normalize_token_ids + + +def create_data_samples(tokenizer) -> DataProto: + convs = [ + [ + { + "role": "user", + "content": "What is the range of the numeric output of a sigmoid node in a neural network?", + }, + {"role": "assistant", "content": "Between -1 and 1."}, + ], + [ + { + "role": "user", + "content": "What is the range of the numeric output of a sigmoid node in a neural network?", + }, + {"role": "assistant", "content": "Between 0 and 1."}, + ], + [ + {"role": "user", "content": "What is the capital of Australia?"}, + { + "role": "assistant", + "content": "Canberra is the capital city of Australia.", + }, + ], + [ + {"role": "user", "content": "What is the capital of Australia?"}, + { + "role": "assistant", + "content": "Sydney is the capital of Australia.", + }, + ], + ] + raw_prompt = [conv[:1] for conv in convs] + data_source = ["gsm8k"] * len(convs) + reward_info = [{"ground_truth": "Not Used"}] * len(convs) + extra_info = [{"question": conv[0]["content"]} for conv in convs] + + prompt_length, response_length = 1024, 4096 + pad_token_id = tokenizer.pad_token_id + prompts, responses, input_ids, attention_masks = [], [], [], [] + for conv in convs: + prompt_tokens = normalize_token_ids(tokenizer.apply_chat_template(conv[:1], tokenize=True)) + response_tokens = normalize_token_ids(tokenizer.apply_chat_template(conv, tokenize=True))[len(prompt_tokens) :] + + padded_prompt = [pad_token_id] * (prompt_length - len(prompt_tokens)) + prompt_tokens + padded_response = response_tokens + [pad_token_id] * (response_length - len(response_tokens)) + attention_mask = ( + [0] * (prompt_length - len(prompt_tokens)) + + [1] * len(prompt_tokens) + + [1] * len(response_tokens) + + [0] * (response_length - len(response_tokens)) + ) + prompts.append(torch.tensor(padded_prompt)) + responses.append(torch.tensor(padded_response)) + input_ids.append(torch.tensor(padded_prompt + padded_response)) + attention_masks.append(torch.tensor(attention_mask)) + + prompts = torch.stack(prompts) + responses = torch.stack(responses) + input_ids = torch.stack(input_ids) + attention_masks = torch.stack(attention_masks) + position_ids = compute_position_id_with_mask(attention_masks) + + data = DataProto.from_dict( + tensors={ + "prompts": prompts, + "responses": responses, + "input_ids": input_ids, + "attention_mask": attention_masks, + "position_ids": position_ids, + }, + non_tensors={ + "data_source": data_source, + "reward_model": reward_info, + "raw_prompt": raw_prompt, + "extra_info": extra_info, + }, + ) + return data, convs + + +def test_reward_model_manager(): + ray.init( + runtime_env={ + "env_vars": { + "TOKENIZERS_PARALLELISM": "true", + "NCCL_DEBUG": "WARN", + "VLLM_LOGGING_LEVEL": "INFO", + "VLLM_USE_V1": "1", + } + } + ) + with initialize_config_dir(config_dir=os.path.abspath("verl/trainer/config")): + config = compose(config_name="ppo_trainer") + + rollout_model_name = os.path.expanduser("~/models/Qwen/Qwen2.5-0.5B-Instruct") + reward_model_name = os.path.expanduser("~/models/Qwen/Qwen2.5-1.5B-Instruct") + + config.actor_rollout_ref.model.path = rollout_model_name + config.reward.custom_reward_function.path = "tests/experimental/reward_loop/reward_fn.py" + config.reward.custom_reward_function.name = "compute_score_gsm8k" + config.reward.num_workers = 1 + config.reward.reward_manager.name = "dapo" + config.reward.reward_model.enable = True + config.reward.reward_model.enable_resource_pool = True + config.reward.reward_model.n_gpus_per_node = 8 + config.reward.reward_model.nnodes = 1 + config.reward.reward_model.model_path = reward_model_name + config.reward.reward_model.rollout.name = os.getenv("ROLLOUT_NAME", "vllm") + config.reward.reward_model.rollout.gpu_memory_utilization = 0.9 + config.reward.reward_model.rollout.tensor_model_parallel_size = 2 + config.reward.reward_model.rollout.skip_tokenizer_init = False + config.reward.reward_model.rollout.prompt_length = 2048 + config.reward.reward_model.rollout.response_length = 4096 + + # 1. init reward model manager + reward_loop_manager = RewardLoopManager(config) + + # 2. init test data + rollout_tokenizer = hf_tokenizer(rollout_model_name) + data, convs = create_data_samples(rollout_tokenizer) + + # 3. generate responses + outputs = reward_loop_manager.compute_rm_score(data) + + for idx, (conv, output) in enumerate(zip(convs, outputs, strict=True)): + print(f"Problem {idx}:\n{conv[0]['content']}\n") + print(f"AI Solution {idx}:\n{conv[1]['content']}\n") + print(f"GRM Response {idx}:\n{output.non_tensor_batch['genrm_response']}\n") + print("=" * 50 + "\n") + + ray.shutdown() diff --git a/verl/tests/experimental/vla/test_sim_envs.py b/verl/tests/experimental/vla/test_sim_envs.py new file mode 100644 index 0000000000000000000000000000000000000000..c2e44feebb5364fb76003a901171caab40fe98b5 --- /dev/null +++ b/verl/tests/experimental/vla/test_sim_envs.py @@ -0,0 +1,113 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import unittest + +import pytest + +# verl/experimental/vla is temporarily unbuildable on this branch after the +# legacy fsdp_workers / actor / critic / sharding_manager modules were removed +# (see the top-level deprecation PR). VLA still imports from those paths, so we +# skip its tests wholesale in CI until VLA is ported to the unified model +# engine. Set VERL_RUN_VLA_TESTS=1 to force-run locally once the port lands. +if os.environ.get("VERL_RUN_VLA_TESTS") != "1": + pytest.skip( + "verl.experimental.vla tests are disabled in CI; set VERL_RUN_VLA_TESTS=1 to opt in.", + allow_module_level=True, + ) + +import numpy as np # noqa: E402 +from omegaconf import OmegaConf # noqa: E402 + + +# @pytest.mark.parametrize("simulator_type", ["libero", "isaac"]) +@pytest.mark.parametrize("simulator_type", ["isaac"]) +def test_sim_env_creation_and_step(simulator_type): + num_envs = 8 + actions = np.array( + [ + [5.59112417e-01, 8.06460073e-02, 1.36817226e-02, -4.64279854e-04, -1.72158767e-02, -6.57548380e-04, -1], + [2.12711899e-03, -3.13366604e-01, 3.41386353e-04, -4.64279854e-04, -8.76528812e-03, -6.57548380e-04, -1], + [7.38182960e-02, -4.64548351e-02, -6.63602950e-02, -4.64279854e-04, -2.32520114e-02, -6.57548380e-04, -1], + [7.38182960e-02, -1.60845593e-01, 3.41386353e-04, -4.64279854e-04, 1.05503430e-02, -6.57548380e-04, -1], + [7.38182960e-02, -3.95982152e-01, -7.97006313e-02, -5.10713711e-03, 3.22804279e-02, -6.57548380e-04, -1], + [2.41859427e-02, -3.64206941e-01, -6.63602950e-02, -4.64279854e-04, 1.05503430e-02, -6.57548380e-04, -1], + [4.62447664e-02, -5.16727952e-01, -7.97006313e-02, -4.64279854e-04, 1.05503430e-02, 8.73740975e-03, -1], + [4.62447664e-02, -5.73923331e-01, 3.41386353e-04, -4.64279854e-04, 6.92866212e-03, -6.57548380e-04, -1], + ] + ) + cfg = OmegaConf.create( + { + "max_episode_steps": 512, + "only_eval": False, + "reward_coef": 1.0, + "init_params": { + "camera_names": ["agentview"], + }, + "video_cfg": { + "save_video": True, + "video_base_dir": "/tmp/test_sim_env_creation_and_step", + }, + "task_suite_name": "libero_10", + "num_envs": num_envs, + "num_group": 1, + "group_size": num_envs, + "seed": 0, + }, + ) + + sim_env = None + if simulator_type == "isaac": + from verl.experimental.vla.envs.isaac_env.isaac_env import IsaacEnv + + sim_env = IsaacEnv(cfg, rank=0, world_size=1) + elif simulator_type == "libero": + from verl.experimental.vla.envs.libero_env.libero_env import LiberoEnv + + sim_env = LiberoEnv(cfg, rank=0, world_size=1) + else: + raise ValueError(f"simulator_type {simulator_type} is not supported") + + video_count = 0 + for i in [0]: + # The first call to step with actions=None will reset the environment + step = 0 + sim_env.reset_envs_to_state_ids([0] * num_envs, [i] * num_envs) + for action in actions: + obs_venv, reward_venv, terminated_venv, truncated_venv, info_venv = sim_env.step( + np.array([action] * num_envs) + ) + + assert isinstance(obs_venv, dict) + assert reward_venv.shape == (num_envs,) + assert terminated_venv.shape == (num_envs,) + assert truncated_venv.shape == (num_envs,) + assert isinstance(info_venv, dict) + + if terminated_venv.any() or truncated_venv.any(): + break + step += 1 + + sim_env.flush_video(video_sub_dir=f"task_{i}") + assert os.path.exists(os.path.join(cfg.video_cfg.video_base_dir, f"rank_0/task_{i}/{video_count}.mp4")) + os.remove(os.path.join(cfg.video_cfg.video_base_dir, f"rank_0/task_{i}/{video_count}.mp4")) + video_count += 1 + + print("test passed") + sim_env.close() + + +if __name__ == "__main__": + unittest.main() diff --git a/verl/tests/kill_github_tests.sh b/verl/tests/kill_github_tests.sh new file mode 100644 index 0000000000000000000000000000000000000000..5c76d7658d5373f4a5d73c9aa9c84a7d14b08402 --- /dev/null +++ b/verl/tests/kill_github_tests.sh @@ -0,0 +1,41 @@ +#!/bin/bash + +if [ "$#" -ne 1 ]; then + echo "Usage: $0 YOUR_GITHUB_TOKEN" + echo "Please provide exactly one input argument for your github token." + exit 1 +fi + +# Set your GitHub repository details +OWNER="volcengine" +REPO="verl" +TOKEN=$1 + +# API URL for workflow runs +API_URL="https://api.github.com/repos/$OWNER/$REPO/actions/runs?status=queued" + +# Check required commands +command -v jq >/dev/null 2>&1 || { echo "jq is required but not installed. Aborting."; exit 1; } + +# Get queued workflow runs +response=$(curl -s -H "Authorization: token $TOKEN" -H "Accept: application/vnd.github.v3+json" "$API_URL") + +# Run this for debugging +# echo $response + +# Extract run IDs +queued_run_ids=$(echo "$response" | jq -r '.workflow_runs[] | .id') + +if [ -z "$queued_run_ids" ]; then + echo "No queued workflow runs found." + exit 0 +fi + +# Cancel each queued run +for run_id in $queued_run_ids; do + echo "Cancelling run $run_id" + cancel_url="https://api.github.com/repos/$OWNER/$REPO/actions/runs/$run_id/cancel" + curl -s -X POST -H "Authorization: token $TOKEN" -H "Accept: application/vnd.github.v3+json" "$cancel_url" +done + +echo "Cancelled all queued workflow runs." diff --git a/verl/tests/models/test_engine.py b/verl/tests/models/test_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..73780921e7ca58633ae3d1c4f49468568ff52647 --- /dev/null +++ b/verl/tests/models/test_engine.py @@ -0,0 +1,524 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os + +os.environ["NCCL_DEBUG"] = "WARN" + +from functools import partial + +import numpy as np +import pytest +import ray +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from transformers import ( + AutoConfig, + AutoModelForCausalLM, + AutoModelForTokenClassification, + AutoTokenizer, + Qwen3Config, + Qwen3MoeConfig, +) + +try: + from transformers.core_model_loading import revert_weight_conversion +except ImportError: + revert_weight_conversion = None + pass + +from verl import DataProto +from verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup +from verl.trainer.config import CheckpointConfig +from verl.utils import tensordict_utils as tu +from verl.utils.model import compute_position_id_with_mask, create_random_mask +from verl.utils.torch_functional import logprobs_from_logits_naive +from verl.workers.config import ( + ActorConfig, + CriticConfig, + FSDPEngineConfig, + FSDPOptimizerConfig, + HFModelConfig, + McoreEngineConfig, + McoreOptimizerConfig, +) +from verl.workers.engine_workers import TrainingWorker, TrainingWorkerConfig +from verl.workers.utils.losses import ppo_loss, sft_loss, value_loss +from verl.workers.utils.padding import left_right_2_no_padding, no_padding_2_padding + + +def get_test_language_model(device_count): + if device_count == 1: + model = "~/models/HuggingFaceTB/SmolLM2-135M-Instruct" + else: + model = "~/models/Qwen/Qwen2.5-0.5B" + model = os.path.expanduser(model) + return model + + +def create_training_config(model_type, strategy, device_count, model): + if device_count == 1: + tp = pp = cp = fsdp_size = 1 + else: + tp = pp = cp = 2 + fsdp_size = 4 + + path = os.path.expanduser(model) + model_config = HFModelConfig(path=path, use_remove_padding=True) + + kwargs = dict( + param_offload=True, + optimizer_offload=True, + grad_offload=True, + use_dynamic_bsz=True, + use_remove_padding=True, + max_token_len_per_gpu=500, + infer_max_token_len_per_gpu=1000, + ) + + if strategy == "megatron": + engine_config = McoreEngineConfig( + forward_only=False, + use_mbridge=True, + tensor_model_parallel_size=tp, + pipeline_model_parallel_size=pp, + context_parallel_size=cp, + **kwargs, + ) + optimizer_config = McoreOptimizerConfig(lr_decay_steps=10) + elif strategy in ["fsdp", "fsdp2"]: + engine_config = FSDPEngineConfig( + forward_only=False, fsdp_size=fsdp_size, strategy=strategy, ulysses_sequence_parallel_size=cp, **kwargs + ) + optimizer_config = FSDPOptimizerConfig() + else: + raise NotImplementedError(f"strategy {strategy} is not supported") + + config = TrainingWorkerConfig( + model_type=model_type, + model_config=model_config, + engine_config=engine_config, + optimizer_config=optimizer_config, + checkpoint_config=None, + ) + return config + + +@pytest.mark.parametrize("strategy", ["fsdp", "fsdp2", "megatron"]) +def test_actor_engine(strategy): + ray.init() + device_count = torch.cuda.device_count() + config = create_training_config( + model_type="language_model", + strategy=strategy, + device_count=device_count, + model=get_test_language_model(device_count), + ) + ray_cls_with_init = RayClassWithInitArgs(cls=ray.remote(TrainingWorker), config=config) + resource_pool = RayResourcePool(process_on_nodes=[device_count]) + wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init) + # init model + wg.reset() + + sft_loss_ = partial(sft_loss, config=config) + + wg.set_loss_fn(sft_loss_) + + batch_size = 8 + seqlen = 32 + + response_length = seqlen // 2 + + torch.manual_seed(1) + np.random.seed(1) + + input_ids = torch.randint(0, config.model_config.hf_config.vocab_size, (batch_size, seqlen)) + attention_mask = create_random_mask( + input_ids=input_ids, max_ratio_of_valid_token=0.8, max_ratio_of_left_padding=0.2, min_ratio_of_valid_token=0.6 + ) + position_ids = compute_position_id_with_mask(attention_mask) + + global_token_num = torch.sum(attention_mask, dim=-1).tolist() + + print(input_ids.float().mean(), attention_mask.float().mean()) + + responses = input_ids[:, response_length:] + response_mask = attention_mask[:, response_length:] + + assert torch.all(response_mask[:, 0] == 1) + + data = DataProto.from_single_dict( + { + "input_ids": input_ids, + "prompts": input_ids[:, :response_length], + "attention_mask": attention_mask, + "position_ids": position_ids, + "responses": responses, + "response_mask": response_mask, + }, + meta_info={"temperature": 1.0, "global_token_num": global_token_num, "compute_loss": False}, + ) + + data_td = data.to_tensordict() + data_td = left_right_2_no_padding(data_td) + + # eval + output = wg.infer_batch(data_td) + output = output.get() + logprobs_unpad = tu.get(output, "log_probs").cpu() + logprobs = no_padding_2_padding(logprobs_unpad, data_td) + + output = DataProto.from_single_dict({"old_log_probs": logprobs}) + + # load hf model and compare results with hf model + path = config.model_config.path + hf_model = AutoModelForCausalLM.from_pretrained(path, torch_dtype=torch.bfloat16) + hf_output = hf_model(input_ids, attention_mask=attention_mask) + hf_logprobs = logprobs_from_logits_naive( + hf_output.logits[:, -response_length - 1 : -1, :].float(), input_ids[:, -response_length:] + ) + hf_logprobs_mean = torch.mean(hf_logprobs * response_mask) + mcore_logprobs_mean = torch.mean(output.batch["old_log_probs"] * response_mask) + + torch.testing.assert_close(hf_logprobs_mean, mcore_logprobs_mean, atol=1e-3, rtol=1e-2) + + data = data.union(output) + + # TODO: sft_loss_ is not compatible with ActorWorker until we replace DataProto with torch.jagged TensorDict + # wg.set_loss_fn(sft_loss_) + + # train for one step + # metrics = wg.update_actor(data) + # print(metrics) + + # add ppo data + data.batch["advantages"] = torch.rand_like(responses, dtype=torch.float32) + data.batch["ref_log_prob"] = torch.rand_like(responses, dtype=torch.float32) + + # construct actor config + actor_config = ActorConfig(strategy=strategy, rollout_n=1, ppo_micro_batch_size_per_gpu=-1) + + # set ppo loss + ppo_loss_ = partial(ppo_loss, config=actor_config) + wg.set_loss_fn(ppo_loss_) + + # update again + data_td = data.to_tensordict() + data_td = left_right_2_no_padding(data_td) + + # auto load/offload + tu.assign_non_tensor(data_td, global_batch_size=data_td.shape[0]) + ppo_metrics = wg.train_batch(data_td) + ppo_metrics = ppo_metrics.get() + ppo_metrics = tu.get(ppo_metrics, "metrics") + print(ppo_metrics) + + # test manual load/offload + tu.assign_non_tensor(data_td, disable_auto_offload=True) + wg.to("device") + ppo_metrics = wg.train_batch(data_td) + ppo_metrics = ppo_metrics.get() + ppo_metrics = tu.get(ppo_metrics, "metrics") + print(ppo_metrics) + wg.to("cpu") + + ray.shutdown() + + +def create_value_model(language_model_path, output_path): + config = AutoConfig.from_pretrained(language_model_path) + config.num_labels = 1 + config.classifier_dropout = 0 + config.tie_word_embeddings = False + model = AutoModelForTokenClassification.from_config(config) + tokenizer = AutoTokenizer.from_pretrained(os.path.expanduser(language_model_path)) + assert model.config.num_labels == 1 + path = os.path.expanduser(output_path) + model.save_pretrained(path) + tokenizer.save_pretrained(path) + config.save_pretrained(path) + return path + + +@pytest.mark.parametrize("strategy", ["fsdp", "fsdp2"]) +def test_critic_engine(strategy): + device_count = torch.cuda.device_count() + value_model_path = os.path.expanduser("~/models/test_model") + language_model_path = get_test_language_model(device_count=device_count) + create_value_model(language_model_path, value_model_path) + + torch.manual_seed(1) + np.random.seed(1) + + ray.init() + + config = create_training_config( + model_type="value_model", strategy=strategy, device_count=device_count, model=value_model_path + ) + ray_cls_with_init = RayClassWithInitArgs(cls=ray.remote(TrainingWorker), config=config) + resource_pool = RayResourcePool(process_on_nodes=[device_count]) + wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init) + # init model + wg.reset() + + batch_size = 8 + seqlen = 32 + + response_length = seqlen // 2 + input_ids = torch.randint(0, config.model_config.hf_config.vocab_size, (batch_size, seqlen)) + attention_mask = create_random_mask( + input_ids=input_ids, max_ratio_of_valid_token=0.8, max_ratio_of_left_padding=0.2, min_ratio_of_valid_token=0.6 + ) + position_ids = compute_position_id_with_mask(attention_mask) + + global_token_num = torch.sum(attention_mask, dim=-1).tolist() + + print(input_ids.float().mean(), attention_mask.float().mean()) + + responses = input_ids[:, response_length:] + response_mask = attention_mask[:, response_length:] + + assert torch.all(response_mask[:, 0] == 1) + + data = DataProto.from_single_dict( + { + "input_ids": input_ids, + "prompts": input_ids[:, :response_length], + "attention_mask": attention_mask, + "position_ids": position_ids, + "responses": responses, + "response_mask": response_mask, + }, + meta_info={"temperature": 1.0, "global_token_num": global_token_num, "compute_loss": False}, + ) + + data_td = data.to_tensordict() + data_td = left_right_2_no_padding(data_td) + + # eval + output = wg.infer_batch(data_td) + output = output.get() + + values_unpad = tu.get(output, "values").float().cpu() + values = no_padding_2_padding(values_unpad, data_td) + + output = DataProto.from_single_dict({"values": values}) + + # load hf model and compare results with hf model + with torch.device("cuda"), torch.autocast(device_type="cuda", dtype=torch.bfloat16): + hf_model = AutoModelForTokenClassification.from_pretrained( + value_model_path, torch_dtype=torch.float32, attn_implementation="flash_attention_2" + ) + hf_output = hf_model(input_ids.cuda(), attention_mask=attention_mask.cuda()) + hf_values = hf_output.logits[:, -response_length - 1 : -1, :].float().squeeze(-1).cpu() + + hf_values_mean = torch.mean(hf_values * response_mask) + engine_values = torch.mean(output.batch["values"] * response_mask) + + torch.testing.assert_close(hf_values_mean, engine_values, atol=1e-2, rtol=1e-2) + + data = data.union(output) + + # add ppo data + data.batch["returns"] = torch.rand_like(responses, dtype=torch.float32) + + # update again + # create critic config + critic_config = CriticConfig( + strategy=strategy, rollout_n=1, ppo_micro_batch_size_per_gpu=-1, model=config.model_config + ) + value_loss_ = partial(value_loss, config=critic_config) + wg.set_loss_fn(value_loss_) + + # update again + data_td = data.to_tensordict() + data_td = left_right_2_no_padding(data_td) + + # auto load/offload + tu.assign_non_tensor(data_td, global_batch_size=data_td.shape[0]) + ppo_metrics = wg.train_batch(data_td) + ppo_metrics = ppo_metrics.get() + ppo_metrics = tu.get(ppo_metrics, "metrics") + print(ppo_metrics) + + ray.shutdown() + + +def create_actor_model(tmp_path, config): + model = AutoModelForCausalLM.from_config(config) + path = os.path.join(tmp_path, "test_model") + model.save_pretrained(path) + config.save_pretrained(path) + return path + + +def _worker(rank: int, world_size: int, rendezvous_file: str, strategy: str, model_path: str): + torch.cuda.set_device(rank) + dist.init_process_group( + backend="nccl", + init_method=f"file://{rendezvous_file}", + rank=rank, + world_size=world_size, + ) + + ref_model_config = AutoConfig.from_pretrained(model_path) + with torch.device("meta"): + ref_model = AutoModelForCausalLM.from_config(ref_model_config) + + from verl.workers.engine import BaseEngine, EngineRegistry + + # construct configs + model_config = HFModelConfig(path=model_path, load_tokenizer=False) + + if strategy == "megatron": + engine_config = McoreEngineConfig( + forward_only=False, + use_mbridge=True, + tensor_model_parallel_size=2, + pipeline_model_parallel_size=2, + context_parallel_size=1, + ) + optimizer_config = McoreOptimizerConfig(lr_decay_steps=10) + elif strategy in ["fsdp", "fsdp2"]: + engine_config = FSDPEngineConfig( + forward_only=False, fsdp_size=4, strategy=strategy, ulysses_sequence_parallel_size=2 + ) + optimizer_config = FSDPOptimizerConfig() + else: + raise NotImplementedError(f"strategy {strategy} is not supported") + + checkpoint_config = CheckpointConfig() + + # build model engine + engine: BaseEngine = EngineRegistry.new( + model_type="language_model", + backend=engine_config.strategy, + model_config=model_config, + engine_config=engine_config, + optimizer_config=optimizer_config, + checkpoint_config=checkpoint_config, + ) + + engine.initialize() + + # get per tensor parameter + per_tensor_params, _ = engine.get_per_tensor_param() + + if strategy == "megatron" and revert_weight_conversion is not None: + ref_state_dict = revert_weight_conversion(ref_model, ref_model.state_dict()) + else: + ref_state_dict = ref_model.state_dict() + + # load ground truth and compare + for key, value in per_tensor_params: + assert key in ref_state_dict, f"{key} not in ref_state_dict" + assert value.shape == ref_state_dict[key].shape, ( + f"{key} shape not equal, {value.shape} != {ref_state_dict[key].shape}" + ) + if rank == 0: + print(key, value.shape) + + dist.barrier() + dist.destroy_process_group() + + +@pytest.mark.parametrize("world_size", [8]) +@pytest.mark.parametrize("config", [Qwen3Config(num_hidden_layers=2), Qwen3MoeConfig(num_hidden_layers=2)]) +@pytest.mark.parametrize("strategy", ["megatron", "fsdp", "fsdp2"]) +def test_per_tensor_generator(world_size, tmp_path, config, strategy): + rendezvous_file = str(tmp_path / "rdzv_mask") + os.makedirs(os.path.dirname(rendezvous_file), exist_ok=True) + # create a model + model_path = create_actor_model(tmp_path, config) + # spawn workers + mp.spawn( + fn=_worker, + args=(world_size, rendezvous_file, strategy, model_path), + nprocs=world_size, + join=True, + ) + + +def _autocast_dtype_worker(rank: int, world_size: int, rendezvous_file: str, model_path: str): + # Regression test for #5932: FSDP engine must resolve autocast dtype from + # mixed_precision.param_dtype rather than hardcoding bfloat16. + torch.cuda.set_device(rank) + dist.init_process_group( + backend="nccl", + init_method=f"file://{rendezvous_file}", + rank=rank, + world_size=world_size, + ) + + from verl.workers.engine import BaseEngine, EngineRegistry + + model_config = HFModelConfig( + path=model_path, + load_tokenizer=False, + override_config={"attn_implementation": "sdpa"}, + ) + + def build_engine(mixed_precision): + engine_config = FSDPEngineConfig( + forward_only=False, + fsdp_size=world_size, + strategy="fsdp2", + ulysses_sequence_parallel_size=1, + mixed_precision=mixed_precision, + ) + engine: BaseEngine = EngineRegistry.new( + model_type="language_model", + backend=engine_config.strategy, + model_config=model_config, + engine_config=engine_config, + optimizer_config=FSDPOptimizerConfig(), + checkpoint_config=CheckpointConfig(), + ) + engine.initialize() + return engine + + from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler + + # bf16 (default) resolves to torch.bfloat16, no scaler needed + engine = build_engine({"param_dtype": "bf16", "reduce_dtype": "fp32", "buffer_dtype": "fp32"}) + assert engine._autocast_dtype == torch.bfloat16, f"expected bf16, got {engine._autocast_dtype}" + assert engine.scaler is None, "bf16 should not create a scaler" + + # fp32 resolves to torch.float32 (forward_step will use nullcontext), no scaler + engine = build_engine({"param_dtype": "fp32", "reduce_dtype": "fp32", "buffer_dtype": "fp32"}) + assert engine._autocast_dtype == torch.float32, f"expected fp32, got {engine._autocast_dtype}" + assert engine.scaler is None, "fp32 should not create a scaler" + + # fp16 gets a ShardedGradScaler for loss scaling during backward + engine = build_engine({"param_dtype": "fp16", "reduce_dtype": "fp32", "buffer_dtype": "fp32"}) + assert engine._autocast_dtype == torch.float16, f"expected fp16, got {engine._autocast_dtype}" + assert isinstance(engine.scaler, ShardedGradScaler), "fp16 must create a ShardedGradScaler" + + dist.barrier() + dist.destroy_process_group() + + +@pytest.mark.parametrize("world_size", [8]) +@pytest.mark.parametrize("config", [Qwen3Config(num_hidden_layers=2)]) +def test_fsdp2_autocast_dtype_honors_mixed_precision(world_size, tmp_path, config): + rendezvous_file = str(tmp_path / "rdzv_autocast") + os.makedirs(os.path.dirname(rendezvous_file), exist_ok=True) + model_path = create_actor_model(tmp_path, config) + mp.spawn( + fn=_autocast_dtype_worker, + args=(world_size, rendezvous_file, model_path), + nprocs=world_size, + join=True, + ) diff --git a/verl/tests/models/test_fsdp_no_padding_on_gpu.py b/verl/tests/models/test_fsdp_no_padding_on_gpu.py new file mode 100644 index 0000000000000000000000000000000000000000..64cd4b046b88e05e435dd6e8762282b475d20c9f --- /dev/null +++ b/verl/tests/models/test_fsdp_no_padding_on_gpu.py @@ -0,0 +1,175 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""GPU regression coverage for issue #6278. + +These tests exercise the FSDP language-model engine path with +``pad_mode=NO_PADDING`` and ``use_remove_padding=False``. The issue was that +the model attention mask was built from response-only ``loss_mask`` rows, then +used as a full prompt+response sequence mask. +""" + +import warnings +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +if not torch.cuda.is_available(): + pytest.skip("Requires CUDA", allow_module_level=True) + + +def _nested(rows: list[list[int]], device: str) -> torch.Tensor: + return torch.nested.as_nested_tensor( + [torch.tensor(row, device=device, dtype=torch.long) for row in rows], + layout=torch.jagged, + ) + + +def _make_micro_batch(device: str): + from tensordict import TensorDict + + from verl.utils import tensordict_utils as tu + from verl.utils.dataset.dataset_utils import DatasetPadMode + + input_ids = _nested([[11, 12, 13, 14, 15], [21, 22, 23]], device) + position_ids = _nested([[0, 1, 2, 3, 4], [0, 1, 2]], device) + prompts = _nested([[11, 12, 13], [21, 22]], device) + responses = _nested([[14, 15], [23]], device) + + # This is intentionally response-only. The old attention-mask path used this + # tensor and produced masks of lengths [2, 1] instead of full sequence [5, 3]. + loss_mask = torch.nested.as_nested_tensor( + [ + torch.ones(2, device=device, dtype=torch.int64), + torch.ones(1, device=device, dtype=torch.int64), + ], + layout=torch.jagged, + ) + + micro_batch = TensorDict( + { + "input_ids": input_ids, + "position_ids": position_ids, + "prompts": prompts, + "responses": responses, + "loss_mask": loss_mask, + "temperature": torch.ones(2, device=device), + }, + batch_size=[2], + ) + tu.assign_non_tensor( + micro_batch, + use_remove_padding=False, + use_fused_kernels=False, + pad_mode=DatasetPadMode.NO_PADDING, + pad_token_id=0, + max_response_len=2, + max_response_length=2, + ) + return micro_batch + + +def _fsdp_engine_with_lm_head_cls(): + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message="NPU not support router replay for now.", category=UserWarning) + from verl.workers.engine.fsdp.transformer_impl import FSDPEngineWithLMHead + + return FSDPEngineWithLMHead + + +def _legacy_attention_mask_from_response_loss_mask( + loss_mask: torch.Tensor, batch_size: int, max_seq_len: int +) -> torch.Tensor: + """Mirror the pre-fix mask construction to make the regression explicit.""" + attention_mask_list = [torch.ones_like(t, dtype=torch.int32) for t in loss_mask] + attention_mask = torch.nested.as_nested_tensor(attention_mask_list, layout=torch.jagged) + return torch.nested.to_padded_tensor(attention_mask, padding=0, output_size=(batch_size, max_seq_len)) + + +def test_prepare_model_inputs_uses_full_sequence_attention_mask_on_gpu(): + device = "cuda" + FSDPEngineWithLMHead = _fsdp_engine_with_lm_head_cls() + engine = FSDPEngineWithLMHead.__new__(FSDPEngineWithLMHead) + micro_batch = _make_micro_batch(device) + + model_inputs, output_args = engine.prepare_model_inputs(micro_batch) + legacy_attention_mask = _legacy_attention_mask_from_response_loss_mask( + loss_mask=micro_batch["loss_mask"], + batch_size=micro_batch.batch_size[0], + max_seq_len=model_inputs["input_ids"].shape[1], + ) + + expected_attention_mask = torch.tensor( + [ + [1, 1, 1, 1, 1], + [1, 1, 1, 0, 0], + ], + device=device, + dtype=torch.int32, + ) + expected_legacy_attention_mask = torch.tensor( + [ + [1, 1, 0, 0, 0], + [1, 0, 0, 0, 0], + ], + device=device, + dtype=torch.int32, + ) + + torch.testing.assert_close(legacy_attention_mask, expected_legacy_attention_mask) + assert not torch.equal(legacy_attention_mask, expected_attention_mask) + torch.testing.assert_close(model_inputs["attention_mask"], expected_attention_mask) + assert model_inputs["attention_mask"].device.type == "cuda" + assert model_inputs["attention_mask"].dtype == torch.int32 + assert model_inputs["input_ids"].shape == (2, 5) + assert output_args["input_ids_rmpad_rolled"].shape == (8,) + + +def test_prepare_model_outputs_can_be_sliced_back_to_response_shape_on_gpu(): + from verl.utils.torch_functional import logprobs_from_logits + from verl.workers.utils.padding import no_padding_2_padding + + device = "cuda" + FSDPEngineWithLMHead = _fsdp_engine_with_lm_head_cls() + engine = FSDPEngineWithLMHead.__new__(FSDPEngineWithLMHead) + micro_batch = _make_micro_batch(device) + _, output_args = engine.prepare_model_inputs(micro_batch) + + torch.manual_seed(0) + vocab_size = 32 + max_seq_len = 5 + logits = torch.randn(2, max_seq_len, vocab_size, device=device) + + model_output = engine.prepare_model_outputs( + output=SimpleNamespace(logits=logits.clone()), + output_args=output_args, + micro_batch=micro_batch, + logits_processor_func=None, + ) + padded_log_probs = no_padding_2_padding(model_output["log_probs"], micro_batch) + + flat_logits = torch.cat([logits[0, :5], logits[1, :3]], dim=0) + expected_full_log_probs = logprobs_from_logits(flat_logits, output_args["input_ids_rmpad_rolled"]) + expected = torch.stack( + [ + expected_full_log_probs[2:4], + F.pad(expected_full_log_probs[6:7], (0, 1)), + ], + dim=0, + ) + + assert model_output["log_probs"].is_nested + assert padded_log_probs.shape == (2, 2) + torch.testing.assert_close(padded_log_probs, expected) diff --git a/verl/tests/models/test_fused_kernels_ulysses_sp_on_cpu.py b/verl/tests/models/test_fused_kernels_ulysses_sp_on_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..69155095b62b7d2c26fdb31b1c95240a4f317e7b --- /dev/null +++ b/verl/tests/models/test_fused_kernels_ulysses_sp_on_cpu.py @@ -0,0 +1,274 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Regression test for issue #6068 (use_fused_kernels=True + ulysses_sp>1). + +Under Ulysses sequence parallelism, the fused-kernel forward functions in +`verl/models/transformers/*.py` were computing +`rolled_labels = torch.roll(input_ids, shifts=-1, dims=-1)` *after* Ulysses had +already SP-sliced the input. `torch.roll` wraps around the local-shard +boundary rather than the global sequence, so the last position on every SP +rank ended up predicting the wrong label. This biased ~1 position per rank +per micro-batch and manifested as a slow training-quality regression at +SP > 1 (issue #6068). + +The fix plumbs the engine's pre-rolled `input_ids_rmpad_rolled` into the fused +forwards via a new `shift_labels` kwarg, mirroring what the veomni engine +already does at `verl/workers/engine/veomni/transformer_impl.py:659`. + +These tests: + 1. Demonstrate the root cause directly (no model needed) — slice-then-roll + diverges from roll-then-slice at every shard boundary. + 2. Verify each adapter's fused-forward (torch + triton backends) honors + the new `shift_labels` kwarg — i.e., the *routing* fix. + +The fused kernels themselves are Triton + CUDA only, so we patch them out +for the CPU tests; their numerical correctness is covered by +`tests/utils/test_linear_cross_entropy.py`. The end-to-end SP=2 integration +test in `tests/special_distributed/test_fused_kernels_ulysses_sp.py` +exercises the full path on 2 GPUs. +""" + +import contextlib +from unittest import mock + +import pytest +import torch + +from verl.models.transformers import dense_common, glm4v, qwen2_vl, qwen3_5, qwen3_vl + +# --------------------------------------------------------------------------- +# 1. Root-cause demonstration: slice-then-local-roll != global-roll-then-slice. +# --------------------------------------------------------------------------- + + +def _global_then_slice(input_ids: torch.Tensor, sp_size: int) -> list[torch.Tensor]: + """Correct behavior: roll on the full sequence, then SP-slice. Mirrors the + engine's `input_ids_rmpad_rolled` -> `ulysses_pad_and_slice_inputs` path. + """ + rolled = torch.roll(input_ids, shifts=-1, dims=-1) + return list(torch.chunk(rolled, sp_size, dim=-1)) + + +def _slice_then_local_roll(input_ids: torch.Tensor, sp_size: int) -> list[torch.Tensor]: + """Buggy behavior: SP-slice first, then roll on the local shard.""" + sliced = list(torch.chunk(input_ids, sp_size, dim=-1)) + return [torch.roll(s, shifts=-1, dims=-1) for s in sliced] + + +def test_local_roll_diverges_from_global_roll_under_sp(): + """For SP > 1, slice-then-local-roll produces different labels than + global-roll-then-slice at exactly the shard-boundary position on every rank. + """ + torch.manual_seed(0) + total_nnz, sp_size = 32, 4 + input_ids = torch.randint(0, 10000, (1, total_nnz)) + + correct_shards = _global_then_slice(input_ids, sp_size) + buggy_shards = _slice_then_local_roll(input_ids, sp_size) + + # Interior positions match — the bug is only at the shard boundary. + for correct, buggy in zip(correct_shards, buggy_shards, strict=True): + torch.testing.assert_close(correct[..., :-1], buggy[..., :-1]) + + # Last position of every shard differs — that's the bias term that + # accumulates across training steps. + for rank in range(sp_size): + correct_last = correct_shards[rank][..., -1] + buggy_last = buggy_shards[rank][..., -1] + assert not torch.equal(correct_last, buggy_last), ( + f"rank {rank}: expected divergence at shard boundary but got {correct_last}=={buggy_last}" + ) + + +# --------------------------------------------------------------------------- +# 2. Adapter routing tests: each fused-forward (torch + triton) must use +# `shift_labels` verbatim when the engine passes it, and must fall back to +# local roll when it's absent (preserves SP=1 behavior). +# --------------------------------------------------------------------------- + + +HIDDEN_SIZE = 8 +VOCAB_SIZE = 64 + + +class _FakeConfig: + """Minimal stand-in for `model.config` used by `forward_base_model`.""" + + output_attentions = False + output_hidden_states = False + + +class _FakeBaseOutput: + """Mimics the HF model output: tuple-indexable and attribute-accessible.""" + + def __init__(self, hidden: torch.Tensor): + self._hidden = hidden + self.hidden_states = None + self.past_key_values = None + self.attentions = None + + def __getitem__(self, idx): + if idx == 0: + return self._hidden + raise IndexError(idx) + + +def _make_fake_lm() -> torch.nn.Module: + """Minimal stand-in for the language-model wrapper used by every adapter. + + Provides `self.config`, `self.model(...)`, and `self.lm_head`. The fake + base model accepts the full kwarg set that `forward_base_model` passes. + """ + + class FakeBaseModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.embed = torch.nn.Embedding(VOCAB_SIZE, HIDDEN_SIZE) + + def forward(self, input_ids=None, **_kwargs): + hidden = self.embed(input_ids).to(torch.float32) + return _FakeBaseOutput(hidden) + + class FakeLM(torch.nn.Module): + def __init__(self): + super().__init__() + self.config = _FakeConfig() + self.model = FakeBaseModel() + self.lm_head = torch.nn.Linear(HIDDEN_SIZE, VOCAB_SIZE, bias=False) + + torch.manual_seed(42) + return FakeLM() + + +@contextlib.contextmanager +def _patch_fused_kernels(): + """Patch the fused kernels (Triton, CUDA-only) and the VLM body wrappers + so the CPU test exercises only the routing layer. Yields a dict that + captures the `input_ids` arg the kernel would have received. + """ + captured: dict[str, torch.Tensor] = {} + + def _fake_log_probs_and_entropy(input_ids: torch.Tensor): + if input_ids.dim() == 1: + shape = input_ids.shape + else: + shape = input_ids.shape # already (B, T) + log_probs = torch.zeros(shape, dtype=torch.float32) + entropy = torch.zeros(shape, dtype=torch.float32) + return log_probs, entropy + + def fake_fused_linear_forward(self, hidden_states, vocab_weights, input_ids, temperature=1.0): + captured["input_ids"] = input_ids.detach().clone() + return _fake_log_probs_and_entropy(input_ids) + + def fake_linear_ce(hidden_states, vocab_weights, input_ids, temperature, reduction): + captured["input_ids"] = input_ids.detach().clone() + return _fake_log_probs_and_entropy(input_ids) + + def fake_qwen2_vl_forward(self, input_ids, **_kwargs): + # Bypass `process_position_ids` and the real VLM body; we only care + # about the routing layer that runs *after* the model forward. + return _FakeBaseOutput(torch.zeros(1, input_ids.shape[-1], self.lm_head.weight.shape[1])) + + def fake_glm4v_forward(self, input_ids, **_kwargs): + return _FakeBaseOutput(torch.zeros(1, input_ids.shape[-1], self.lm_head.weight.shape[1])) + + with ( + mock.patch( + "verl.utils.experimental.torch_functional.FusedLinearForPPO.forward", + new=fake_fused_linear_forward, + ), + mock.patch( + "verl.utils.kernel.linear_cross_entropy.linear_cross_entropy", + new=fake_linear_ce, + ), + mock.patch( + "verl.models.transformers.qwen2_vl.qwen2_vl_forward", + new=fake_qwen2_vl_forward, + ), + mock.patch( + "verl.models.transformers.glm4v.glm4v_forward", + new=fake_glm4v_forward, + ), + ): + yield captured + + +ALL_ADAPTERS = [ + dense_common.forward_with_torch_backend, + dense_common.forward_with_triton_backend, + qwen3_5.forward_with_torch_backend, + qwen3_5.forward_with_triton_backend, + qwen3_vl.forward_with_torch_backend, + qwen3_vl.forward_with_triton_backend, + qwen2_vl.forward_with_torch_backend, + qwen2_vl.forward_with_triton_backend, + glm4v.forward_with_torch_backend, + glm4v.forward_with_triton_backend, +] + + +def _adapter_id(forward_fn) -> str: + return f"{forward_fn.__module__.rsplit('.', 1)[-1]}.{forward_fn.__name__}" + + +@pytest.mark.parametrize("forward_fn", ALL_ADAPTERS, ids=_adapter_id) +def test_adapter_honors_shift_labels(forward_fn): + """When `shift_labels` is provided, every fused adapter must pass it + through to the kernel verbatim — no local re-rolling. + """ + model = _make_fake_lm() + input_ids = torch.tensor([[10, 20, 30, 40]], dtype=torch.long) + # Deliberately != torch.roll(input_ids); the only way these labels reach + # the kernel is if the adapter honors `shift_labels`. + shift_labels = torch.tensor([[20, 30, 40, 50]], dtype=torch.long) + + with _patch_fused_kernels() as captured: + forward_fn( + model, + input_ids=input_ids, + labels=None, + temperature=1.0, + shift_labels=shift_labels, + return_dict=True, + ) + + assert "input_ids" in captured, "fused kernel was never called" + assert torch.equal(captured["input_ids"], shift_labels), ( + f"{_adapter_id(forward_fn)}: expected kernel to see {shift_labels.tolist()}, " + f"got {captured['input_ids'].tolist()}" + ) + + +@pytest.mark.parametrize("forward_fn", ALL_ADAPTERS, ids=_adapter_id) +def test_adapter_falls_back_to_local_roll_when_shift_labels_absent(forward_fn): + """Backward-compat: callers that don't pass `shift_labels` (e.g. the SP=1 + code path or any non-engine consumer) see unchanged behavior. + """ + model = _make_fake_lm() + input_ids = torch.tensor([[10, 20, 30, 40]], dtype=torch.long) + expected = torch.roll(input_ids, shifts=-1, dims=-1) + + with _patch_fused_kernels() as captured: + forward_fn( + model, + input_ids=input_ids, + labels=None, + temperature=1.0, + return_dict=True, + ) + + assert torch.equal(captured["input_ids"], expected), ( + f"{_adapter_id(forward_fn)}: expected fallback to torch.roll(input_ids), got {captured['input_ids'].tolist()}" + ) diff --git a/verl/tests/models/test_liger_vl_compat.py b/verl/tests/models/test_liger_vl_compat.py new file mode 100644 index 0000000000000000000000000000000000000000..2eca4e64c24c038a9d5d25b005b108a45f4b41d8 --- /dev/null +++ b/verl/tests/models/test_liger_vl_compat.py @@ -0,0 +1,86 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Test that Liger kernel integration doesn't break VL models. + +Regression test for https://github.com/verl-project/verl/issues/2609 +and https://github.com/verl-project/verl/issues/1720. + +The bug: _apply_liger_kernel_to_instance with default kwargs sets +fused_linear_cross_entropy=True, which replaces the model forward. +After verl's apply_monkey_patch patches the base model forward, +Liger's forward crashes with 'BaseModelOutputWithPast has no attribute rope_deltas'. +""" + +import pytest +import torch +from transformers import Qwen3VLConfig, Qwen3VLForConditionalGeneration + +from verl.models.transformers.monkey_patch import apply_monkey_patch + + +def create_tiny_qwen3_vl(): + """Create a minimal Qwen3-VL model (random weights, 1 layer) for testing.""" + config = Qwen3VLConfig( + text_config=dict( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + vocab_size=1000, + rope_scaling=dict( + type="mrope", + mrope_section=[4, 4, 4], + ), + ), + vision_config=dict( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=1, + num_attention_heads=4, + patch_size=16, + temporal_patch_size=2, + in_channels=3, + spatial_merge_size=2, + ), + ) + model = Qwen3VLForConditionalGeneration(config) + return model.to(dtype=torch.bfloat16, device="cuda") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_liger_vl_forward_with_monkey_patch(): + """Liger with fused_linear_cross_entropy=False + apply_monkey_patch works on VL models.""" + pytest.importorskip("liger_kernel") + from liger_kernel.transformers.monkey_patch import _apply_liger_kernel_to_instance + + model = create_tiny_qwen3_vl() + + _apply_liger_kernel_to_instance( + model=model, + fused_linear_cross_entropy=False, + swiglu=True, + ) + apply_monkey_patch(model, use_remove_padding=False, use_fused_kernels=False) + + input_ids = torch.randint(0, 500, (1, 32), device="cuda") + attention_mask = torch.ones_like(input_ids) + output = model(input_ids=input_ids, attention_mask=attention_mask) + + assert output.logits is not None + assert output.logits.shape[-1] == 1000 # vocab_size + + del model + torch.cuda.empty_cache() diff --git a/verl/tests/models/test_tiled_mlp_accuracy.py b/verl/tests/models/test_tiled_mlp_accuracy.py new file mode 100644 index 0000000000000000000000000000000000000000..6b022243ffe4ba15724fcf2c89f91a92e0b1e37c --- /dev/null +++ b/verl/tests/models/test_tiled_mlp_accuracy.py @@ -0,0 +1,218 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Test script to verify TiledMLP accuracy by comparing logits and gradients +between regular MLP and TiledMLP under FSDP2. +Run with: torchrun --nproc_per_node=2 tests/test_tiled_mlp_accuracy.py +""" + +import torch +import torch.distributed as dist +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.fsdp import fully_shard + + +def setup_distributed(): + dist.init_process_group(backend="nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + torch.cuda.set_device(rank) + return rank, world_size + + +def create_model(model_name="Qwen/Qwen3-1.7B", num_layers=2): + """Load a Qwen3-1.7B model with only 2 layers from pretrained weights.""" + from transformers import AutoConfig, AutoModelForCausalLM + + config = AutoConfig.from_pretrained(model_name, trust_remote_code=True) + config.num_hidden_layers = num_layers + + model = AutoModelForCausalLM.from_pretrained( + model_name, + config=config, + torch_dtype=torch.bfloat16, + trust_remote_code=True, + attn_implementation="flash_attention_2", + ) + return model + + +def apply_fsdp2(model, device_mesh): + """Apply FSDP2 sharding to model.""" + for layer in model.model.layers: + fully_shard(layer, mesh=device_mesh) + fully_shard(model, mesh=device_mesh) + return model + + +def run_forward_backward(model, input_ids, labels): + """Run forward and backward pass, return logits and gradients.""" + model.zero_grad() + + outputs = model(input_ids=input_ids, labels=labels) + logits = outputs.logits.clone().detach() + loss = outputs.loss + + loss.backward() + + # Collect MLP gradients + gradients = {} + for name, param in model.named_parameters(): + if "mlp" in name and param.grad is not None: + gradients[name] = param.grad.clone().detach() + + return logits, gradients, loss.item() + + +def compare_results(logits1, grads1, logits2, grads2, rank): + """Compare logits and gradients between two runs.""" + # Compare logits + logits_diff = (logits1 - logits2).abs() + logits_max_diff = logits_diff.max().item() + logits_mean_diff = logits_diff.mean().item() + + # Compare gradients (only for params that exist on this rank due to FSDP sharding) + all_pass = True + grad_results = [] + for name in sorted(grads1.keys()): + if name in grads2: + g1, g2 = grads1[name], grads2[name] + diff = (g1 - g2).abs() + max_diff = diff.max().item() + mean_diff = diff.mean().item() + + # Check if within tolerance (1e-2 for bf16) + passed = max_diff < 1e-2 + if not passed: + all_pass = False + grad_results.append((name, max_diff, mean_diff, passed)) + + # Only print on rank 0 to avoid duplicate output + if rank == 0: + print("\n=== Comparison Results ===") + print("\nLogits:") + print(f" Max diff: {logits_max_diff:.2e}") + print(f" Mean diff: {logits_mean_diff:.2e}") + + print("\nMLP Parameter Gradients:") + if grad_results: + for name, max_diff, mean_diff, passed in grad_results: + status = "✓" if passed else "✗" + print(f" {name}: max={max_diff:.2e}, mean={mean_diff:.2e} {status}") + else: + print(" (Gradients sharded to other ranks under FSDP2)") + + return all_pass + + +def main(): + rank, world_size = setup_distributed() + device_mesh = init_device_mesh("cuda", (world_size,)) + + model_name = "Qwen/Qwen3-1.7B" + num_layers = 2 + + if rank == 0: + print(f"Running TiledMLP accuracy test with {world_size} GPUs") + print(f"Model: {model_name} ({num_layers} layers, from pretrained)") + + dist.barrier() + + # ========== Create Model 1: WITHOUT TiledMLP ========== + if rank == 0: + print("\n" + "=" * 60) + print("Creating Model 1 (without TiledMLP)") + print("=" * 60) + + model1 = create_model(model_name, num_layers) + model1 = apply_fsdp2(model1, device_mesh) + model1 = model1.cuda() + + # Create deterministic input + torch.manual_seed(42) + batch_size, seq_len = 2, 256 + vocab_size = model1.config.vocab_size + input_ids = torch.randint(0, vocab_size, (batch_size, seq_len), device="cuda") + labels = input_ids.clone() + + # ========== Run Model 1: WITHOUT TiledMLP ========== + if rank == 0: + print("\n" + "=" * 60) + print("Running forward/backward on Model 1 (without TiledMLP)") + print("=" * 60) + + logits1, grads1, loss1 = run_forward_backward(model1, input_ids, labels) + if rank == 0: + print(f"Loss: {loss1:.4f}") + + # Free model1 memory before creating model2 + del model1 + torch.cuda.empty_cache() + + dist.barrier() + + # ========== Create Model 2, apply TiledMLP patch, then FSDP2 ========== + if rank == 0: + print("\n" + "=" * 60) + print("Creating Model 2 (with TiledMLP, patch before FSDP2)") + print("=" * 60) + + model2 = create_model(model_name, num_layers) + + # Apply TiledMLP patch AFTER model instantiation but BEFORE FSDP2 wrap + if rank == 0: + print("Applying TiledMLP monkey patch before FSDP2...") + + from verl.models.transformers.tiled_mlp import apply_tiled_mlp_monkey_patch + + apply_tiled_mlp_monkey_patch(num_shards=4, model_type="qwen3") + + model2 = apply_fsdp2(model2, device_mesh) + model2 = model2.cuda() + + dist.barrier() + + # ========== Run Model 2: WITH TiledMLP ========== + if rank == 0: + print("\n" + "=" * 60) + print("Running forward/backward on Model 2 (with TiledMLP)") + print("=" * 60) + + logits2, grads2, loss2 = run_forward_backward(model2, input_ids, labels) + if rank == 0: + print(f"Loss: {loss2:.4f}") + + dist.barrier() + + # ========== Compare Results ========== + all_pass = compare_results(logits1, grads1, logits2, grads2, rank) + + dist.barrier() + + if rank == 0: + print("\n" + "=" * 60) + print("SUMMARY") + print("=" * 60) + print(f"Loss diff: {abs(loss1 - loss2):.2e}") + print(f"All gradient checks: {'PASS' if all_pass else 'FAIL'}") + + # Cleanup + del model2 + torch.cuda.empty_cache() + + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/verl/tests/models/test_transformer.py b/verl/tests/models/test_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..ddd085497a16cd73e828bff596dd888d054827af --- /dev/null +++ b/verl/tests/models/test_transformer.py @@ -0,0 +1,239 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +from transformers import ( + ApertusConfig, + AutoModelForCausalLM, + AutoModelForTokenClassification, + GemmaConfig, + LlamaConfig, + MistralConfig, + Qwen2Config, +) + +from verl.utils.device import get_device_name + +if get_device_name() == "cuda": + from flash_attn.bert_padding import index_first_axis, pad_input, rearrange, unpad_input +elif get_device_name() == "npu": + from verl.utils.attention_utils import index_first_axis, pad_input, rearrange, unpad_input + +from verl.utils.model import compute_position_id_with_mask, create_random_mask +from verl.utils.torch_functional import log_probs_from_logits_all_rmpad, masked_mean + +# TODO(sgm): add more models for test +# we only need one scale for each model +test_configs = [ + LlamaConfig(num_hidden_layers=1), + MistralConfig(num_hidden_layers=1), + GemmaConfig(num_hidden_layers=1), + Qwen2Config(num_hidden_layers=1), + ApertusConfig(num_hidden_layers=1), +] + + +def test_hf_casual_models(): + batch_size = 4 + seqlen = 128 + response_length = 127 + + for config in test_configs: + # config = AutoConfig.from_pretrained(test_case) + with torch.device(get_device_name()): + model = AutoModelForCausalLM.from_config( + config=config, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2" + ) + model = model.to(device=get_device_name()) + input_ids = torch.randint(low=0, high=config.vocab_size, size=(batch_size, seqlen), device=get_device_name()) + attention_mask = create_random_mask( + input_ids=input_ids, + max_ratio_of_left_padding=0.1, + max_ratio_of_valid_token=0.8, + min_ratio_of_valid_token=0.5, + ) + position_ids = compute_position_id_with_mask( + attention_mask + ) # TODO(sgm): we can construct the position_ids_rmpad here + + input_ids_rmpad, indices, *_ = unpad_input( + input_ids.unsqueeze(-1), attention_mask + ) # input_ids_rmpad (total_nnz, ...) + input_ids_rmpad = input_ids_rmpad.transpose(0, 1) # (1, total_nnz) + + # unpad the position_ids to align the rotary + position_ids_rmpad = index_first_axis( + rearrange(position_ids.unsqueeze(-1), "b s ... -> (b s) ..."), indices + ).transpose(0, 1) + + # input with input_ids_rmpad and postition_ids to enable flash attention varlen + logits_rmpad = model( + input_ids_rmpad, position_ids=position_ids_rmpad, use_cache=False + ).logits # (1, total_nnz, vocab_size) + + origin_logits = model( + input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, use_cache=False + ).logits + origin_logits_rmpad, origin_logits_indices, *_ = unpad_input(origin_logits, attention_mask) + + logits_rmpad = logits_rmpad.squeeze(0) + log_probs = log_probs_from_logits_all_rmpad( + input_ids_rmpad=input_ids_rmpad, + logits_rmpad=logits_rmpad, + indices=indices, + batch_size=batch_size, + seqlen=seqlen, + response_length=response_length, + ) # (batch, seqlen) + origin_log_probs = log_probs_from_logits_all_rmpad( + input_ids_rmpad=input_ids_rmpad, + logits_rmpad=origin_logits_rmpad, + indices=origin_logits_indices, + batch_size=batch_size, + seqlen=seqlen, + response_length=response_length, + ) # (batch, seqlen) + + torch.testing.assert_close( + masked_mean(log_probs, attention_mask[:, -response_length - 1 : -1]), + masked_mean(origin_log_probs, attention_mask[:, -response_length - 1 : -1]), + atol=1e-2, + rtol=1e-5, + ) + print("Check pass") + + +def test_hf_value_models(): + batch_size = 4 + seqlen = 128 + + for config in test_configs: + # config = AutoConfig.from_pretrained(test_case) + config.num_labels = 1 + config.classifier_dropout = 0 + config.hidden_dropout = 0 + with torch.device(get_device_name()): + model = AutoModelForTokenClassification.from_config( + config=config, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2" + ) + model = model.to(device=get_device_name()) + input_ids = torch.randint(low=0, high=config.vocab_size, size=(batch_size, seqlen), device=get_device_name()) + attention_mask = create_random_mask( + input_ids=input_ids, + max_ratio_of_left_padding=0.1, + max_ratio_of_valid_token=0.8, + min_ratio_of_valid_token=0.5, + ) + position_ids = compute_position_id_with_mask( + attention_mask + ) # TODO(sgm): we can construct the position_ids_rmpad here + + input_ids_rmpad, indices, *_ = unpad_input( + input_ids.unsqueeze(-1), attention_mask + ) # input_ids_rmpad (total_nnz, ...) + input_ids_rmpad = input_ids_rmpad.transpose(0, 1) # (1, total_nnz) + + # unpad the position_ids to align the rotary + position_ids_rmpad = index_first_axis( + rearrange(position_ids.unsqueeze(-1), "b s ... -> (b s) ..."), indices + ).transpose(0, 1) + + origin_logits = model( + input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, use_cache=False + ).logits + + # input with input_ids_rmpad and postition_ids to enable flash attention varlen + rmpad_logits = model( + input_ids_rmpad, position_ids=position_ids_rmpad, use_cache=False + ).logits # (1, total_nnz, 1) + rmpad_logits = rmpad_logits.squeeze(0) + pad_logits = pad_input(rmpad_logits, indices, batch_size, seqlen=seqlen) + + torch.testing.assert_close( + masked_mean(pad_logits, attention_mask[:, :, None]), + masked_mean(origin_logits, attention_mask[:, :, None]), + atol=1e-2, + rtol=1e-5, + ) + print("Value model check pass") + + +def test_attn_implementation_override(): + """Test that attn_implementation override config is properly respected.""" + # Test case 1: Test the actual extraction logic (no network required) + test_cases = [ + ({}, "flash_attention_2"), # Default case + ({"attn_implementation": "eager"}, "eager"), # Override case + ({"attn_implementation": "sdpa"}, "sdpa"), # Another override + ({"other_config": "value"}, "flash_attention_2"), # No attn_implementation key + ] + + for override_config, expected in test_cases: + actual = override_config.get("attn_implementation", "flash_attention_2") + assert actual == expected, f"Expected {expected}, got {actual} for config {override_config}" + + # Test case 2: Test with local config creation (simulate FSDP worker behavior) + # Test default behavior + override_config_default = {} + attn_implementation_default = override_config_default.get("attn_implementation", "flash_attention_2") + assert attn_implementation_default == "flash_attention_2" + + # Test override behavior + override_config_eager = {"attn_implementation": "eager"} + attn_implementation_eager = override_config_eager.get("attn_implementation", "flash_attention_2") + assert attn_implementation_eager == "eager" + + # Test that we can create a config with specific attn_implementation + config_with_eager = LlamaConfig(num_hidden_layers=1, _attn_implementation="eager") + assert config_with_eager._attn_implementation == "eager" + + config_with_flash = LlamaConfig(num_hidden_layers=1, _attn_implementation="flash_attention_2") + assert config_with_flash._attn_implementation == "flash_attention_2" + + print("✓ All attn_implementation override config tests passed") + + +def test_fsdp_worker_attn_implementation_integration(): + """Test integration of attn_implementation with FSDP worker logic.""" + + # Mock the FSDP worker configuration scenario + mock_override_config = {"attn_implementation": "eager"} + + # Test the exact logic used in FSDP workers + attn_implementation = mock_override_config.get("attn_implementation", "flash_attention_2") + assert attn_implementation == "eager" + + # Test with empty config (should default) + mock_override_config_empty = {} + attn_implementation_default = mock_override_config_empty.get("attn_implementation", "flash_attention_2") + assert attn_implementation_default == "flash_attention_2" + + # Test that the parameter would be passed correctly to both AutoConfig and Model + expected_calls = [ + ("AutoConfig.from_pretrained", {"attn_implementation": attn_implementation}), + ("AutoModel.from_pretrained", {"attn_implementation": attn_implementation}), + ] + + # Verify the parameter extraction works as expected + for call_name, expected_params in expected_calls: + assert expected_params["attn_implementation"] == "eager" + + print("✓ FSDP worker integration test passed") + + +if __name__ == "__main__": + test_hf_casual_models() + test_hf_value_models() + test_attn_implementation_override() + test_fsdp_worker_attn_implementation_integration() diff --git a/verl/tests/models/test_transformers_ulysses.py b/verl/tests/models/test_transformers_ulysses.py new file mode 100644 index 0000000000000000000000000000000000000000..c3e6eef9021959840897f3d00371d4eaf90f112d --- /dev/null +++ b/verl/tests/models/test_transformers_ulysses.py @@ -0,0 +1,283 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import contextlib +import copy +from dataclasses import dataclass + +import pytest +import torch +import torch.distributed +import transformers +from packaging import version +from torch.distributed import init_device_mesh +from transformers import AutoModelForCausalLM, LlamaConfig, PretrainedConfig, Qwen2Config + +from verl.models.transformers.monkey_patch import apply_monkey_patch +from verl.protocol import DataProto +from verl.utils.device import get_device_name, get_torch_device +from verl.utils.distributed import initialize_global_process_group +from verl.utils.model import compute_position_id_with_mask, create_random_mask +from verl.utils.ulysses import ( + FSDPUlyssesShardingManager, + gather_outputs_and_unpad, + get_ulysses_sequence_parallel_world_size, + set_ulysses_sequence_parallel_group, + ulysses_pad_and_slice_inputs, +) + +if get_device_name() == "cuda": + from flash_attn.bert_padding import index_first_axis, rearrange, unpad_input +elif get_device_name() == "npu": + from verl.utils.attention_utils import index_first_axis, rearrange, unpad_input + +# TODO(sgm): add more models for test +# we only need one scale for each model + + +@dataclass +class SequenceParallelConfig: + config: PretrainedConfig + sp_size: int + is_valid: bool + + +def test_configs(): + configs = [ + SequenceParallelConfig( + LlamaConfig(num_hidden_layers=2, num_attention_heads=32, num_key_value_heads=32), sp_size=8, is_valid=True + ), + SequenceParallelConfig( + Qwen2Config(num_hidden_layers=2, num_attention_heads=28, num_key_value_heads=4, hidden_size=3584), + sp_size=4, + is_valid=True, + ), + SequenceParallelConfig( + Qwen2Config(num_hidden_layers=2, num_attention_heads=28, num_key_value_heads=4, hidden_size=3584), + sp_size=8, + is_valid=False, + ), + SequenceParallelConfig( + Qwen2Config(num_hidden_layers=2, num_attention_heads=32, num_key_value_heads=4), sp_size=4, is_valid=True + ), + SequenceParallelConfig( + Qwen2Config(num_hidden_layers=2, num_attention_heads=32, num_key_value_heads=4), sp_size=8, is_valid=True + ), + ] + + if version.parse(transformers.__version__) >= version.parse("4.56.0"): + from transformers import ApertusConfig + + configs.append( + SequenceParallelConfig( + ApertusConfig(num_hidden_layers=2, num_attention_heads=32, num_key_value_heads=32, hidden_size=4096), + sp_size=8, + is_valid=True, + ) + ) + + return configs + + +def sync_model_parameters_global(layer): + # synchronize weights + for p in layer.parameters(): + torch.distributed.broadcast(tensor=p.data, src=0) + + +@pytest.mark.parametrize("test_config", test_configs()) +def test_hf_casual_fwd_bwd(test_config): + if not torch.distributed.is_initialized(): + initialize_global_process_group() + + context = contextlib.nullcontext() if test_config.is_valid else pytest.raises(AssertionError) + with context: + world_size = torch.distributed.get_world_size() + _hf_casual_fwd_bwd(test_config.config, test_config.sp_size, world_size // test_config.sp_size) + + # TODO: seems not work, will cause `socketStartConnect: Connect to xxx failed : Software caused connection abort` + # torch.distributed.destroy_process_group() + + +def _hf_casual_fwd(config, sp_size, dp_size): + assert get_torch_device().device_count() >= 2, "need at least 2 gpus for test" + + ulysses_device_mesh = init_device_mesh( + device_type=get_device_name(), mesh_shape=(dp_size, sp_size), mesh_dim_names=("dp", "sp") + ) + sharding_manager = FSDPUlyssesShardingManager(ulysses_device_mesh) + + batch_size = 1 + seqlen = 128 + # response_length = 127 + + # patch before load + with torch.device(get_device_name()): + model = AutoModelForCausalLM.from_config( + config=config, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2" + ) + apply_monkey_patch(model, sp_size) + model = model.to(device=get_device_name()) + sync_model_parameters_global(model) + + # different rank will generate different input_ids following fsdp + input_ids = torch.randint(low=0, high=config.vocab_size, size=(batch_size, seqlen), device=get_device_name()) + attention_mask = create_random_mask( + input_ids=input_ids, max_ratio_of_left_padding=0, max_ratio_of_valid_token=0.9, min_ratio_of_valid_token=0.8 + ) + position_ids = compute_position_id_with_mask( + attention_mask + ) # TODO(sgm): we can construct the position_ids_rmpad here + + model_inputs = { + "input_ids": input_ids.to(get_device_name()), + "attention_mask": attention_mask.to(get_device_name()), + "position_ids": position_ids.int().to(get_device_name()), + } + + model_inputs = DataProto.from_dict(model_inputs) + + # 1. perform ulysses forward + with sharding_manager: + model_inputs = sharding_manager.preprocess_data(model_inputs) + input_ids = model_inputs.batch["input_ids"] + attention_mask = model_inputs.batch["attention_mask"] + position_ids = model_inputs.batch["position_ids"] + input_ids_rmpad, indices, *_ = unpad_input( + input_ids.unsqueeze(-1), attention_mask + ) # input_ids_rmpad (total_nnz, ...) + input_ids_rmpad = input_ids_rmpad.transpose(0, 1) # (1, total_nnz) + # unpad the position_ids to align the rotary + position_ids_rmpad = index_first_axis( + rearrange(position_ids.unsqueeze(-1), "b s ... -> (b s) ..."), indices + ).transpose(0, 1) + + # slice input tensor for ulysses + # input_ids are padded and sliced + # postition_ids are only padded but not sliced + input_ids_rmpad_sliced, position_ids_rmpad_padded, pad_size = ulysses_pad_and_slice_inputs( + input_ids_rmpad, position_ids_rmpad, sp_size=get_ulysses_sequence_parallel_world_size() + ) + + # input with input_ids_rmpad and postition_ids to enable flash attention varlen + logits_split_in_seq = model( + input_ids_rmpad_sliced, position_ids=position_ids_rmpad_padded, use_cache=False + ).logits # (1, total_nnz/n, vocab_size) + + # all_gather output + logits_full = gather_outputs_and_unpad(logits_split_in_seq, gather_dim=1, unpad_dim=1, padding_size=pad_size) + + # 2. perform normal forward + set_ulysses_sequence_parallel_group(None) + logits_rmpad_local = model( + input_ids_rmpad, position_ids=position_ids_rmpad, use_cache=False + ).logits # (1, total_nnz, vocab_size) + + mean_local = logits_rmpad_local.mean() + mean_full = logits_full.mean() + torch.testing.assert_close(mean_local, mean_full, rtol=1e-2, atol=1e-5) + + +def _hf_casual_fwd_bwd(config, sp_size, dp_size): + assert get_torch_device().device_count() >= 2, "need at least 2 gpus for test" + + ulysses_device_mesh = init_device_mesh( + device_type=get_device_name(), mesh_shape=(dp_size, sp_size), mesh_dim_names=("dp", "sp") + ) + sharding_manager = FSDPUlyssesShardingManager(ulysses_device_mesh) + + batch_size = 1 + seqlen = 128 + # response_length = 127 + + # patch before load + with torch.device(get_device_name()): + model = AutoModelForCausalLM.from_config( + config=config, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2" + ) + apply_monkey_patch(model, sp_size) + model = model.to(device=get_device_name()) + sync_model_parameters_global(model) + + # different rank will generate different input_ids following fsdp + input_ids = torch.randint(low=0, high=config.vocab_size, size=(batch_size, seqlen), device=get_device_name()) + attention_mask = create_random_mask( + input_ids=input_ids, max_ratio_of_left_padding=0, max_ratio_of_valid_token=0.9, min_ratio_of_valid_token=0.8 + ) + position_ids = compute_position_id_with_mask( + attention_mask + ) # TODO(sgm): we can construct the position_ids_rmpad here + + model_inputs = { + "input_ids": input_ids.to(get_device_name()), + "attention_mask": attention_mask.to(get_device_name()), + "position_ids": position_ids.int().to(get_device_name()), + } + + model_inputs = DataProto.from_dict(model_inputs) + + # 1. perform ulysses forward + with sharding_manager: + model_inputs = sharding_manager.preprocess_data(model_inputs) + input_ids = model_inputs.batch["input_ids"] + attention_mask = model_inputs.batch["attention_mask"] + position_ids = model_inputs.batch["position_ids"] + input_ids_rmpad, indices, *_ = unpad_input( + input_ids.unsqueeze(-1), attention_mask + ) # input_ids_rmpad (total_nnz, ...) + input_ids_rmpad = input_ids_rmpad.transpose(0, 1) # (1, total_nnz) + # unpad the position_ids to align the rotary + position_ids_rmpad = index_first_axis( + rearrange(position_ids.unsqueeze(-1), "b s ... -> (b s) ..."), indices + ).transpose(0, 1) + + # slice input tensor for ulysses + # input_ids are padded and sliced + # postition_ids are only padded but not sliced + input_ids_rmpad_sliced, position_ids_rmpad_padded, pad_size = ulysses_pad_and_slice_inputs( + input_ids_rmpad, position_ids_rmpad, sp_size=get_ulysses_sequence_parallel_world_size() + ) + + # input with input_ids_rmpad and postition_ids to enable flash attention varlen + logits_split_in_seq = model( + input_ids_rmpad_sliced, position_ids=position_ids_rmpad_padded, use_cache=False + ).logits # (1, total_nnz/n, vocab_size) + + # all_gather output + logits_full = gather_outputs_and_unpad(logits_split_in_seq, gather_dim=1, unpad_dim=1, padding_size=pad_size) + + # 2. perform normal forward + set_ulysses_sequence_parallel_group(None) + input_ids_full = copy.deepcopy(input_ids_rmpad) + position_ids_full = copy.deepcopy(position_ids_rmpad) + model_no_sp = copy.deepcopy(model) + logits_rmpad_local = model_no_sp( + input_ids_full, position_ids=position_ids_full, use_cache=False + ).logits # (1, total_nnz, vocab_size) + + mean_local = logits_rmpad_local.mean() + mean_full = logits_full.mean() + + mean_full.backward() + mean_local.backward() + + # 3. check the gradients + grad = model.model.layers[0].self_attn.q_proj.weight.grad + grad_full = model_no_sp.model.layers[0].self_attn.q_proj.weight.grad + torch.testing.assert_close(mean_local, mean_full, rtol=1e-2, atol=3e-5) + # The check should be less strict because the gradient is not an averaged value. + torch.testing.assert_close(grad, grad_full, rtol=1e-2, atol=1e-3) + + +if __name__ == "__main__": + pytest.main([__file__, "-svv"]) diff --git a/verl/tests/single_controller/__init__.py b/verl/tests/single_controller/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1cd1e8433dffa0b3ba420be3e346f4f5cd062014 --- /dev/null +++ b/verl/tests/single_controller/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/verl/tests/single_controller/base/test_decorator.py b/verl/tests/single_controller/base/test_decorator.py new file mode 100644 index 0000000000000000000000000000000000000000..5447d65ce0ecfad235d63c3c8ca02d88c4c7a9e7 --- /dev/null +++ b/verl/tests/single_controller/base/test_decorator.py @@ -0,0 +1,76 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +import verl.single_controller.base.decorator as decorator_module +from verl.single_controller.base.decorator import ( + DISPATCH_MODE_FN_REGISTRY, + Dispatch, + _check_dispatch_mode, + get_predefined_dispatch_fn, + register_dispatch_mode, + update_dispatch_mode, +) + + +@pytest.fixture +def reset_dispatch_registry(): + # Store original state + original_registry = DISPATCH_MODE_FN_REGISTRY.copy() + yield + # Reset registry after test + decorator_module.DISPATCH_MODE_FN_REGISTRY.clear() + decorator_module.DISPATCH_MODE_FN_REGISTRY.update(original_registry) + + +def test_register_new_dispatch_mode(reset_dispatch_registry): + # Test registration + def dummy_dispatch(worker_group, *args, **kwargs): + return args, kwargs + + def dummy_collect(worker_group, output): + return output + + register_dispatch_mode("TEST_MODE", dummy_dispatch, dummy_collect) + + # Verify enum extension + _check_dispatch_mode(Dispatch.TEST_MODE) + + # Verify registry update + assert get_predefined_dispatch_fn(Dispatch.TEST_MODE) == { + "dispatch_fn": dummy_dispatch, + "collect_fn": dummy_collect, + } + # Clean up + Dispatch.remove("TEST_MODE") + + +def test_update_existing_dispatch_mode(reset_dispatch_registry): + # Store original implementation + original_mode = Dispatch.ONE_TO_ALL + + # New implementations + def new_dispatch(worker_group, *args, **kwargs): + return args, kwargs + + def new_collect(worker_group, output): + return output + + # Test update= + update_dispatch_mode(original_mode, new_dispatch, new_collect) + + # Verify update + assert get_predefined_dispatch_fn(original_mode)["dispatch_fn"] == new_dispatch + assert get_predefined_dispatch_fn(original_mode)["collect_fn"] == new_collect diff --git a/verl/tests/single_controller/check_worker_alive/main.py b/verl/tests/single_controller/check_worker_alive/main.py new file mode 100644 index 0000000000000000000000000000000000000000..cbdee9a8d6cf98544efc8abeb9555a66a2fd70ee --- /dev/null +++ b/verl/tests/single_controller/check_worker_alive/main.py @@ -0,0 +1,64 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import sys +import time + +import ray + +from verl.single_controller.base.decorator import Dispatch, register +from verl.single_controller.base.worker import Worker +from verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup + + +@ray.remote +class TestActor(Worker): + def __init__(self) -> None: + super().__init__() + + @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False) + def foo(self, wait_time): + time.sleep(wait_time) + sys.exit(1) + + +if __name__ == "__main__": + wait_time = int(os.getenv("WAIT_TIME", "10")) + + ray.init() + + # test single-node-no-partition + print("test single-node-no-partition") + resource_pool = RayResourcePool([2], use_gpu=False) + class_with_args = RayClassWithInitArgs(cls=TestActor) + + print("create worker group") + wg = RayWorkerGroup(resource_pool, class_with_args, name_prefix="test") + + wg.start_worker_aliveness_check(1) + time.sleep(1) + + print(time.time(), "start foo") + + _ = wg.foo(wait_time) + print("foo started") + + print( + time.time(), + f"wait 6x wait time {wait_time * 6} to let signal returned to process but still not exceed process wait time", + ) + time.sleep(wait_time * 6) + + ray.shutdown() diff --git a/verl/tests/single_controller/detached_worker/README.md b/verl/tests/single_controller/detached_worker/README.md new file mode 100644 index 0000000000000000000000000000000000000000..b06c4c6143e01d071458f7416033872d41d71031 --- /dev/null +++ b/verl/tests/single_controller/detached_worker/README.md @@ -0,0 +1,14 @@ +# Detached Worker +## How to run (Only on a single node) +- Start a local ray cluster: +```bash +ray start --head --port=6379 +``` +- Run the server +```bash +python3 server.py +``` +- On another terminal, Run the client +```bash +python3 client.py +``` diff --git a/verl/tests/single_controller/detached_worker/client.py b/verl/tests/single_controller/detached_worker/client.py new file mode 100644 index 0000000000000000000000000000000000000000..8c78aaf5d37f6ca5aced3ba5a42b64218cb950e1 --- /dev/null +++ b/verl/tests/single_controller/detached_worker/client.py @@ -0,0 +1,56 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +In client, we can get the server handler and send RPC request +""" + +import ray +import torch +from server import Trainer +from tensordict import TensorDict + +from verl import DataProto +from verl.single_controller.ray import RayClassWithInitArgs, RayWorkerGroup + + +def compute_position_id_with_mask(mask): + return torch.clip(torch.cumsum(mask, dim=-1) - 1, min=0, max=None) + + +if __name__ == "__main__": + ray.init(address="auto", namespace="verl") + # get the worker group using names + worker_names = ["trainerTrainer_0:0", "trainerTrainer_0:1"] + cls_with_init_args = RayClassWithInitArgs(cls=Trainer) + worker_group = RayWorkerGroup.from_detached(worker_names=worker_names, ray_cls_with_init=cls_with_init_args) + + batch_size = 16 + sequence_length = 1024 + + # give Trainer some data to train + input_ids = torch.randint(low=0, high=256, size=(batch_size, sequence_length), dtype=torch.int64, device="cuda") + attention_mask = torch.ones_like(input_ids) + position_ids = compute_position_id_with_mask(attention_mask) + + data = DataProto( + batch=TensorDict( + {"input_ids": input_ids, "attention_mask": attention_mask, "position_ids": position_ids}, + batch_size=batch_size, + ), + meta_info={}, + ) + + output = worker_group.train_model(data) + + print(output) diff --git a/verl/tests/single_controller/detached_worker/run.sh b/verl/tests/single_controller/detached_worker/run.sh new file mode 100644 index 0000000000000000000000000000000000000000..a3c6387933262694bf3534066b4310fda0a9fea3 --- /dev/null +++ b/verl/tests/single_controller/detached_worker/run.sh @@ -0,0 +1,5 @@ +#!/bin/bash +ray start --head --port=6379 +python3 server.py +python3 client.py +ray stop --force \ No newline at end of file diff --git a/verl/tests/single_controller/detached_worker/server.py b/verl/tests/single_controller/detached_worker/server.py new file mode 100644 index 0000000000000000000000000000000000000000..f8a7f014d2317b2d1918b2fe9fd5d6b177e09317 --- /dev/null +++ b/verl/tests/single_controller/detached_worker/server.py @@ -0,0 +1,152 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Server starts a Trainer. Client sends data to the server to train. +""" + +import os + +os.environ["MEGATRON_USE_CUDA_TIMER"] = "0" +os.environ["MEGATRON_START_PROCESS_TIMER"] = "False" +os.environ["NCCL_DEBUG"] = "WARN" + +import ray +import torch +from megatron.core import parallel_state as mpu +from megatron.core import tensor_parallel +from megatron.core.models.gpt.gpt_model import ModelType +from omegaconf import OmegaConf +from tensordict import TensorDict +from torch import nn +from transformers import LlamaConfig + +from verl import DataProto +from verl.models.llama.megatron import ParallelLlamaForCausalLMRmPadPP +from verl.single_controller.base import Worker +from verl.single_controller.base.decorator import Dispatch, make_nd_compute_dataproto_dispatch_fn, register +from verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup +from verl.utils.megatron.optimizer import get_megatron_optimizer, init_megatron_optim_config +from verl.utils.megatron_utils import get_model, mcore_model_parallel_config + + +@ray.remote +class Trainer(Worker): + def __init__(self): + super().__init__() + + if not torch.distributed.is_initialized(): + rank = int(os.environ["LOCAL_RANK"]) + torch.distributed.init_process_group(backend="nccl") + torch.cuda.set_device(rank) + + mpu.initialize_model_parallel( + tensor_model_parallel_size=2, + pipeline_model_parallel_size=1, + virtual_pipeline_model_parallel_size=None, + use_sharp=False, + context_parallel_size=1, + expert_model_parallel_size=1, + nccl_communicator_config_path=None, + ) + tensor_parallel.model_parallel_cuda_manual_seed(10) + + is_collect = ( + mpu.get_tensor_model_parallel_rank() == 0 + and mpu.get_pipeline_model_parallel_rank() == mpu.get_pipeline_model_parallel_world_size() - 1 + and mpu.get_context_parallel_rank() == 0 + ) + self._register_dispatch_collect_info( + mesh_name="train", dp_rank=mpu.get_data_parallel_rank(), is_collect=is_collect + ) + + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def init_model(self): + actor_model_config = LlamaConfig( + vocab_size=256, + hidden_size=2048, + intermediate_size=5504, + num_hidden_layers=24, + num_attention_heads=16, + num_key_value_heads=16, + ) + + megatron_config = mcore_model_parallel_config(sequence_parallel=True, params_dtype=torch.bfloat16) + self.megatron_config = megatron_config + + def megatron_actor_model_provider(pre_process, post_process): + # vpp is not supported yet because it will hang for some reason. Need debugging + # this_megatron_config = copy.deepcopy(megatron_config) + # this_megatron_config.virtual_pipeline_model_parallel_rank = vpp_rank + parallel_model = ParallelLlamaForCausalLMRmPadPP( + config=actor_model_config, + megatron_config=megatron_config, + pre_process=pre_process, + post_process=post_process, + ) + parallel_model.cuda() + return parallel_model + + actor_module = get_model( + model_provider_func=megatron_actor_model_provider, + model_type=ModelType.encoder_or_decoder, + wrap_with_ddp=True, + ) + actor_module = nn.ModuleList(actor_module) + + optim_config = OmegaConf.create({"lr": 1e-6, "clip_grad": 1.0}) + + optim_config = init_megatron_optim_config(optim_config) + self.optimizer_config = optim_config + actor_optimizer = get_megatron_optimizer(model=actor_module, config=optim_config) + + self.model = actor_module[0] + self.optimizer = actor_optimizer + + @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="train")) + def train_model(self, data: DataProto) -> DataProto: + input_ids = data.batch["input_ids"] + attention_mask = data.batch["attention_mask"] + position_ids = data.batch["position_ids"] + + self.optimizer.zero_grad() + self.model.zero_grad_buffer( + zero_buffer=(not self.optimizer_config.use_distributed_optimizer) + ) # use use_contiguous_buffers_in_local_ddp and no overlap_dp_param_comm + # update for 1 iteration + output = self.model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids).logits + output.mean().backward() + + update_successful, grad_norm, num_zeros_in_grad = self.optimizer.step( + self.megatron_config, self.megatron_config.timers + ) + + return DataProto(batch=TensorDict({"loss": output.detach()}, batch_size=output.shape[0])) + + +if __name__ == "__main__": + ray.init(address="auto", namespace="verl") + + resource_pool = RayResourcePool(process_on_nodes=[2], detached=True) + cls_with_init_args = RayClassWithInitArgs(cls=Trainer) + worker_group = RayWorkerGroup( + resource_pool=resource_pool, + ray_cls_with_init=cls_with_init_args, + name_prefix="trainer", + detached=True, + ) + + worker_group.init_model() + + worker_names = worker_group.worker_names + print(worker_names) diff --git a/verl/tests/single_controller/test_auto_padding_on_cpu.py b/verl/tests/single_controller/test_auto_padding_on_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..b60e719c98a5ee42918b55326f2f98c443c7dd9d --- /dev/null +++ b/verl/tests/single_controller/test_auto_padding_on_cpu.py @@ -0,0 +1,152 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import ray +import torch + +from verl import DataProto +from verl.protocol import DataProtoConfig +from verl.single_controller.base import Worker +from verl.single_controller.base.decorator import Dispatch, register +from verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup + +# or set env var VERL_AUTO_PADDING = "1" / "true" +DataProtoConfig.auto_padding = True + + +@ray.remote +class Actor(Worker): + def __init__(self) -> None: + super().__init__() + + @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO) + def add(self, data: DataProto): + data.batch["a"] += self.rank + return data + + +def test_auto_padding(): + ray.init(num_cpus=100) + + chunk_size = 4 + actor_cls = RayClassWithInitArgs(cls=Actor) + resource_pool = RayResourcePool(process_on_nodes=[chunk_size], use_gpu=False) + actor_wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=actor_cls) + + # test locally first + for test_size in range(4, 20): + local_data = DataProto.from_dict({"a": torch.zeros(test_size)}, {"na": np.zeros(test_size, dtype=object)}) + # print(f"before padding, local_data = {local_data}") + padding_size = (chunk_size - (test_size % chunk_size)) if (test_size % chunk_size > 0) else 0 + local_data.padding(padding_size) + # print(f"after padding, local_data = {local_data}") + assert len(local_data) == len(local_data) + len(local_data) % chunk_size, ( + f"expecting padded length to be {len(local_data) + len(local_data) % chunk_size}, but got {len(local_data)}" + ) + chunked = local_data.chunk(chunk_size) + assert len(chunked) == chunk_size, f"during test_size = {test_size}, expecting {chunk_size}, got {chunked}" + for dp in chunked: + assert len(dp) == test_size // chunk_size + bool(test_size % chunk_size), ( + f"test size = {test_size}, expecting dp to be length of " + f"{test_size // chunk_size + bool(test_size % chunk_size)}, but got {len(dp)}: {dp} {chunked}" + ) + + # test with RayWorkerGroup method decorated as dispatch_mode=Dispatch.DP_COMPUTE_PROTO + data = DataProto.from_dict({"a": torch.zeros(10)}, {"na": np.array([str(i) for i in range(10)], dtype=object)}) + output = actor_wg.add(data) + + print(output.batch["a"]) + assert len(output) == 10, "Failed in args split and padding." + + data = DataProto.from_dict({"a": torch.zeros(10)}, {"na": np.array([str(i) for i in range(10)], dtype=object)}) + output = actor_wg.add(data=data) + + print(output.batch["a"]) + assert len(output) == 10, "Failed in kwargs split and padding." + + data = DataProto.from_dict({"a": torch.zeros(1)}, {"na": np.array([str(i) for i in range(1)], dtype=object)}) + output = actor_wg.add(data) + + print(output.batch["a"]) + assert len(output) == 1, "Failed in args split and padding." + + data = DataProto.from_dict({"a": torch.zeros(1)}, {"na": np.array([str(i) for i in range(1)], dtype=object)}) + output = actor_wg.add(data=data) + + print(output.batch["a"]) + assert len(output) == 1, "Failed in kwargs split and padding." + + data = DataProto.from_dict({"a": torch.zeros(8)}, {"na": np.array([str(i) for i in range(8)], dtype=object)}) + output = actor_wg.add(data) + + print(output.batch["a"]) + assert len(output) == 8, "Failed in args split and padding." + + data = DataProto.from_dict({"a": torch.zeros(8)}, {"na": np.array([str(i) for i in range(8)], dtype=object)}) + output = actor_wg.add(data=data) + + print(output.batch["a"]) + assert len(output) == 8, "Failed in kwargs split and padding." + + # test data proto specific config + DataProtoConfig.auto_padding = False + + data = DataProto.from_dict( + {"a": torch.zeros(10)}, {"na": np.array([str(i) for i in range(10)], dtype=object)}, auto_padding=True + ) + output = actor_wg.add(data) + print(output.batch["a"]) + assert len(output) == 10, "Failed in args split and padding." + + data = DataProto.from_dict( + {"a": torch.zeros(10)}, {"na": np.array([str(i) for i in range(10)], dtype=object)}, auto_padding=True + ) + output = actor_wg.add(data=data) + print(output.batch["a"]) + assert len(output) == 10, "Failed in kwargs split and padding." + + data = DataProto.from_single_dict( + {"a": torch.zeros(1), "na": np.array([str(i) for i in range(1)], dtype=object)}, auto_padding=True + ) + output = actor_wg.add(data) + + print(output.batch["a"]) + assert len(output) == 1, "Failed in args split and padding." + + data = DataProto.from_single_dict( + {"a": torch.zeros(1), "na": np.array([str(i) for i in range(1)], dtype=object)}, auto_padding=True + ) + output = actor_wg.add(data=data) + + print(output.batch["a"]) + assert len(output) == 1, "Failed in kwargs split and padding." + + data = DataProto.from_single_dict({"a": torch.zeros(8), "na": np.array([str(i) for i in range(8)], dtype=object)}) + output = actor_wg.add(data) + + print(output.batch["a"]) + assert len(output) == 8, "Failed in args split and padding." + + data = DataProto.from_single_dict({"a": torch.zeros(8), "na": np.array([str(i) for i in range(8)], dtype=object)}) + output = actor_wg.add(data=data) + + print(output.batch["a"]) + assert len(output) == 8, "Failed in kwargs split and padding." + + ray.shutdown() + + +if __name__ == "__main__": + test_auto_padding() diff --git a/verl/tests/single_controller/test_colocated_workers.py b/verl/tests/single_controller/test_colocated_workers.py new file mode 100644 index 0000000000000000000000000000000000000000..980362f3f8ea6657d2cd582cb7cb4e07ff32769a --- /dev/null +++ b/verl/tests/single_controller/test_colocated_workers.py @@ -0,0 +1,86 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import ray + +from verl import DataProto +from verl.single_controller.base import Worker +from verl.single_controller.base.decorator import Dispatch, register +from verl.single_controller.ray.base import ( + RayClassWithInitArgs, + RayResourcePool, + RayWorkerGroup, + create_colocated_worker_cls, +) +from verl.utils.device import get_device_name + + +@ray.remote +class Actor(Worker): + def __init__(self) -> None: + super().__init__() + + @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO) + def add(self, data: DataProto): + data.batch["a"] += self.rank + return data + + +@ray.remote +class Critic(Worker): + def __init__(self, config) -> None: + super().__init__() + self.config = config + + @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO) + async def sub(self, data: DataProto): + data.batch["a"] -= self.config["b"] + return data + + +def test_colocated_workers(): + ray.init() + + import torch + + data = DataProto.from_dict({"a": torch.zeros(10)}) + # create separate workers on the same resource pool + actor_cls = RayClassWithInitArgs(cls=Actor) + critic_cls = RayClassWithInitArgs(cls=Critic, config={"b": 10}) + resource_pool = RayResourcePool(process_on_nodes=[2]) + + actor_wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=actor_cls, device_name=get_device_name()) + critic_wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=critic_cls, device_name=get_device_name()) + + expected_actor_output = actor_wg.add(data) + expected_critic_output = critic_wg.sub(data) + + # create colocated workers + cls_dict = {"actor": actor_cls, "critic": critic_cls} + ray_cls_with_init = create_colocated_worker_cls(cls_dict) + wg_dict = RayWorkerGroup( + resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init, device_name=get_device_name() + ) + spawn_wg = wg_dict.spawn(prefix_set=cls_dict.keys()) + + colocated_actor_wg = spawn_wg["actor"] + colocated_critic_wg = spawn_wg["critic"] + + actor_output = colocated_actor_wg.add(data) + critic_output = colocated_critic_wg.sub(data) + + torch.testing.assert_close(expected_actor_output.batch, actor_output.batch, atol=0, rtol=0) + torch.testing.assert_close(expected_critic_output.batch, critic_output.batch, atol=0, rtol=0) + + ray.shutdown() diff --git a/verl/tests/single_controller/test_colocated_workers_fused.py b/verl/tests/single_controller/test_colocated_workers_fused.py new file mode 100644 index 0000000000000000000000000000000000000000..6597ff4f6f667a59f528ce22bf13ca9af300dfd7 --- /dev/null +++ b/verl/tests/single_controller/test_colocated_workers_fused.py @@ -0,0 +1,86 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import ray + +from verl import DataProto +from verl.single_controller.base import Worker +from verl.single_controller.base.decorator import Dispatch, register +from verl.single_controller.ray.base import ( + RayClassWithInitArgs, + RayResourcePool, + RayWorkerGroup, + create_colocated_worker_cls_fused, +) +from verl.utils.device import get_device_name + + +@ray.remote +class Actor(Worker): + def __init__(self) -> None: + super().__init__() + + @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO) + def add(self, data: DataProto): + data.batch["a"] += self.rank + return data + + +@ray.remote +class Critic(Worker): + def __init__(self, config) -> None: + super().__init__() + self.config = config + + @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO) + def sub(self, data: DataProto): + data.batch["a"] -= self.config["b"] + return data + + +def test_colocated_workers_fused(): + ray.init() + + import torch + + data = DataProto.from_dict({"a": torch.zeros(10)}) + # create separate workers on the same resource pool + actor_cls = RayClassWithInitArgs(cls=Actor) + critic_cls = RayClassWithInitArgs(cls=Critic, config={"b": 10}) + resource_pool = RayResourcePool(process_on_nodes=[2]) + + actor_wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=actor_cls, device_name=get_device_name()) + critic_wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=critic_cls, device_name=get_device_name()) + + expected_actor_output = actor_wg.add(data) + expected_critic_output = critic_wg.sub(data) + + # create colocated workers + cls_dict = {"actor": actor_cls, "critic": critic_cls} + ray_cls_with_init = create_colocated_worker_cls_fused(cls_dict) + wg_dict = RayWorkerGroup( + resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init, device_name=get_device_name() + ) + spawn_wg = wg_dict.spawn(prefix_set=cls_dict.keys()) + + colocated_actor_wg = spawn_wg["actor"] + colocated_critic_wg = spawn_wg["critic"] + + actor_output = colocated_actor_wg.add(data) + critic_output = colocated_critic_wg.sub(data) + + torch.testing.assert_close(expected_actor_output.batch, actor_output.batch, atol=0, rtol=0) + torch.testing.assert_close(expected_critic_output.batch, critic_output.batch, atol=0, rtol=0) + + ray.shutdown() diff --git a/verl/tests/single_controller/test_data_transfer.py b/verl/tests/single_controller/test_data_transfer.py new file mode 100644 index 0000000000000000000000000000000000000000..848a3000974c11f6e29c235c59ef64de61dfaf83 --- /dev/null +++ b/verl/tests/single_controller/test_data_transfer.py @@ -0,0 +1,110 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +In this test, we instantiate a data parallel worker with N GPUs (auto-detected). +""" + +import ray +import tensordict +import torch +from codetiming import Timer +from packaging import version +from torch import distributed as dist + +from verl import DataProto +from verl.single_controller.base import Worker +from verl.single_controller.base.decorator import Dispatch, register +from verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup +from verl.utils.device import get_device_name +from verl.utils.ray_utils import parallel_put + + +@ray.remote +class DummyWorker(Worker): + def __init__(self): + super().__init__() + dist.init_process_group() + + @register(dispatch_mode=Dispatch.DP_COMPUTE, blocking=False) + def do_nothing(self, data): + for key in data.batch.keys(): + data.batch[key] += 1 + if version.parse(tensordict.__version__) >= version.parse("0.5.0"): + data.batch = data.batch.consolidate() + return data + + +def test_data_transfer(): + ray.init() + # construct resource pool + ngpus = torch.cuda.device_count() + resource_pool = RayResourcePool([ngpus]) + cls_with_init = RayClassWithInitArgs(cls=DummyWorker) + # construct worker group + wg = RayWorkerGroup(resource_pool, cls_with_init, device_name=get_device_name()) + + # this is real dataset size + batch_size = 4096 + seqlen = 32768 + + data_dict = {} + + for i in range(2): + data_dict[str(i)] = torch.randint(0, 10000, (batch_size, seqlen)) + + data = DataProto.from_dict(tensors=data_dict) + + print(data) + + # we manually split data here and send to each worker + data_list = data.chunk(wg.world_size) + + for i in range(wg.world_size): + # consolidate is necessary + if version.parse(tensordict.__version__) >= version.parse("0.5.0"): + data_list[i].batch = data_list[i].batch.consolidate() + + with Timer(name="ray.pickle", initial_text=True): + for i in range(wg.world_size): + ray.cloudpickle.pickle.dumps(data_list[i]) + + with Timer(name="raw.pickle", initial_text=True): + import pickle + + for i in range(wg.world_size): + pickle.dumps(data_list[i]) + + # we put in advance + with Timer(name="put", initial_text=True): + # takes around 40 seconds + data_list_ref = parallel_put(data_list) + # for i in range(wg.world_size): + # data_list[i] = ray.put(data_list[i]) + + with Timer(name="launch", initial_text=True): + output_ref = wg.do_nothing(data_list_ref) + + with Timer(name="get", initial_text=True): + # takes around 40 seconds + output_lst = ray.get(output_ref) + + for input_data, output_data in zip(data_list, output_lst, strict=True): + for key in input_data.batch.keys(): + assert torch.all(torch.eq(input_data.batch[key] + 1, output_data.batch[key])), ( + input_data.batch[key], + output_data.batch[key], + key, + ) + + ray.shutdown() diff --git a/verl/tests/single_controller/test_decorator_on_cpu.py b/verl/tests/single_controller/test_decorator_on_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..cb67bd0d4cf4b78997127333089dc3e37a4b05a5 --- /dev/null +++ b/verl/tests/single_controller/test_decorator_on_cpu.py @@ -0,0 +1,200 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +import time + +import pytest +import ray +import torch +from tensordict import TensorDict + +from verl.protocol import DataProto, DataProtoFuture +from verl.single_controller.base.decorator import Dispatch, make_nd_compute_dataproto_dispatch_fn, register +from verl.single_controller.base.worker import Worker +from verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup +from verl.utils import tensordict_utils as tu + + +# Pytest fixture for Ray setup/teardown +@pytest.fixture +def ray_init_shutdown(): + ray.init(num_cpus=100) + yield + ray.shutdown() + + +# Define a simple worker for testing +@ray.remote +class DecoratorTestWorker(Worker): + def __init__(self, initial_value=0): + super().__init__() + self.value = initial_value + # Simulate some setup if needed + time.sleep(0.1) # Ensure worker init completes + + self._register_dispatch_collect_info(mesh_name="train", dp_rank=self.rank, is_collect=True) + + # Test method for synchronous DP compute (default behavior) + @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO) + def dp_compute(self, data: DataProto) -> DataProto: + time.sleep(0.1) # Simulate work + rank_value = torch.tensor(self.rank, device=data.batch["input"].device, dtype=data.batch["input"].dtype) + data.batch["output"] = data.batch["input"] + self.value + rank_value + return data + + # Test async def method with DP compute (default behavior) + @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO, blocking=False) + async def async_dp_compute(self, data: DataProto) -> DataProto: + # Simulate async work + await asyncio.sleep(0.1) # Simulate async work + rank_value = torch.tensor(self.rank, device=data.batch["input"].device, dtype=data.batch["input"].dtype) + data.batch["output_async"] = data.batch["input"] * 2 + self.value + rank_value + return data + + @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="train"), blocking=False) + def dp_compute_td(self, data: TensorDict) -> TensorDict: + # note that we have to call contiguous so that we can modify data in plac + data = tu.contiguous(data) + rank_value = torch.tensor(self.rank, device=data["input"].device, dtype=data["input"].dtype) + data["output"] = data["input"] + self.value + rank_value + position_ids = data.pop("position_ids") + position_ids._ragged_idx = 2 + + for i, position_id in enumerate(position_ids.unbind(dim=0)): + assert (position_id == torch.arange(4 + rank_value * 2 + i).expand(position_id.shape)).all() + + return data + + +# Test function for synchronous DP compute +def test_decorator_dp_compute(ray_init_shutdown): + """ + Tests the default behavior of a synchronous decorated method with DP_COMPUTE_PROTO. + Verifies the result correctness. + """ + num_workers = 2 + resource_pool = RayResourcePool([num_workers], use_gpu=False, max_colocate_count=1) # Use CPU for simplicity + cls_with_args = RayClassWithInitArgs(cls=DecoratorTestWorker, initial_value=10) + worker_group = RayWorkerGroup( + resource_pool, cls_with_args, name_prefix=f"decorator_test_sync_dp_{int(time.time())}" + ) + + # Prepare input data (size 4, for 2 workers) + input_tensor = torch.arange(4, dtype=torch.float32) + data = DataProto(batch=TensorDict({"input": input_tensor}, batch_size=[4])) + + # Call the decorated method + output = worker_group.dp_compute(data) + + # Assert the result correctness + assert isinstance(output, DataProto), "Expected DataProto result" + assert "output" in output.batch.keys() + assert len(output) == len(data), "Output length should match input length" + + # Expected output calculation for DP_COMPUTE_PROTO with 2 workers + # Worker 0 gets data[0:2], Worker 1 gets data[2:4] + # Worker 0 adds initial_value(10) + rank(0) = 10 + # Worker 1 adds initial_value(10) + rank(1) = 11 + expected_output_part1 = torch.tensor([0, 1], dtype=torch.float32) + 10 + 0 + expected_output_part2 = torch.tensor([2, 3], dtype=torch.float32) + 10 + 1 + expected_output = torch.cat([expected_output_part1, expected_output_part2]) + + torch.testing.assert_close(output.batch["output"], expected_output, msg="Sync DP compute output data mismatch") + + +# Test function for async def method with DP compute +def test_decorator_async_function(ray_init_shutdown): + """ + Tests the decorator with an `async def` method using DP_COMPUTE_PROTO. + Verifies that the call returns a future and the result is correct after .get(). + """ + num_workers = 2 + resource_pool = RayResourcePool([num_workers], use_gpu=False, max_colocate_count=1) + cls_with_args = RayClassWithInitArgs(cls=DecoratorTestWorker, initial_value=5) + worker_group = RayWorkerGroup( + resource_pool, cls_with_args, name_prefix=f"decorator_test_async_dp_{int(time.time())}" + ) + + # Prepare input data (size 4, for 2 workers) + input_tensor = torch.arange(4, dtype=torch.float32) + data = DataProto(batch=TensorDict({"input": input_tensor}, batch_size=[4])) + + # Call the async decorated method - this should return a future + future_output: DataProtoFuture = worker_group.async_dp_compute(data) + + # Assert that the call returned a future + assert isinstance(future_output, DataProtoFuture), "Expected DataProtoFuture for async def call" + + # Get the result (this should block) + result_data = future_output.get() + + # Assert the result correctness + assert isinstance(result_data, DataProto) + assert "output_async" in result_data.batch.keys() + assert len(result_data) == len(data), "Output length should match input length" + + # Expected output calculation for DP_COMPUTE_PROTO with 2 workers + # Worker 0 gets data[0:2], Worker 1 gets data[2:4] + # Worker 0 calculates: input * 2 + initial_value(5) + rank(0) + # Worker 1 calculates: input * 2 + initial_value(5) + rank(1) + expected_output_part1 = (torch.tensor([0, 1], dtype=torch.float32) * 2) + 5 + 0 + expected_output_part2 = (torch.tensor([2, 3], dtype=torch.float32) * 2) + 5 + 1 + expected_output = torch.cat([expected_output_part1, expected_output_part2]) + + torch.testing.assert_close( + result_data.batch["output_async"], expected_output, msg="Async DP compute output data mismatch" + ) + + +def test_decorator_dp_compute_td(ray_init_shutdown): + num_workers = 2 + resource_pool = RayResourcePool([num_workers], use_gpu=False, max_colocate_count=1) # Use CPU for simplicity + cls_with_args = RayClassWithInitArgs(cls=DecoratorTestWorker, initial_value=10) + worker_group = RayWorkerGroup( + resource_pool, cls_with_args, name_prefix=f"decorator_test_sync_dp_{int(time.time())}" + ) + + # Prepare input data (size 4, for 2 workers) + input_tensor = torch.arange(4, dtype=torch.float32) + position_ids = torch.nested.as_nested_tensor( + [ + torch.arange(4).expand(4, 4).contiguous(), + torch.arange(5).expand(4, 5).contiguous(), + torch.arange(6).expand(4, 6).contiguous(), + torch.arange(7).expand(4, 7).contiguous(), + ], + layout=torch.jagged, + ) + data = TensorDict({"input": input_tensor, "position_ids": position_ids}, batch_size=[4]) + + # Call the decorated method + output = worker_group.dp_compute_td(data) + + output = output.get() + + # Assert the result correctness + assert isinstance(output, TensorDict), "Expected DataProto result" + assert "output" in output.keys() + assert len(output) == len(data), "Output length should match input length" + + # Expected output calculation for DP_COMPUTE_PROTO with 2 workers + # Worker 0 gets data[0:2], Worker 1 gets data[2:4] + # Worker 0 adds initial_value(10) + rank(0) = 10 + # Worker 1 adds initial_value(10) + rank(1) = 11 + expected_output_part1 = torch.tensor([0, 1], dtype=torch.float32) + 10 + 0 + expected_output_part2 = torch.tensor([2, 3], dtype=torch.float32) + 10 + 1 + expected_output = torch.cat([expected_output_part1, expected_output_part2]) + + torch.testing.assert_close(output["output"], expected_output, msg="Sync DP compute output data mismatch") diff --git a/verl/tests/single_controller/test_device_mesh_register.py b/verl/tests/single_controller/test_device_mesh_register.py new file mode 100644 index 0000000000000000000000000000000000000000..0bf3846519b4835b4df4132d6751f175bb77b30c --- /dev/null +++ b/verl/tests/single_controller/test_device_mesh_register.py @@ -0,0 +1,164 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import numpy as np +import ray +import torch +from tensordict import TensorDict + +import verl.utils.tensordict_utils as tu +from verl import DataProto +from verl.single_controller.base import Worker +from verl.single_controller.base.decorator import make_nd_compute_dataproto_dispatch_fn, register +from verl.utils.device import get_device_name, get_nccl_backend + + +@ray.remote +class TestActor(Worker): + def __init__(self): + super().__init__() + + import torch.distributed + + torch.distributed.init_process_group(backend=get_nccl_backend()) + ngpus = torch.distributed.get_world_size() + self.infer_device_mesh = torch.distributed.device_mesh.init_device_mesh( + device_type=get_device_name(), mesh_shape=[2, ngpus // 2], mesh_dim_names=["dp", "tp"] + ) + self.train_device_mesh = torch.distributed.device_mesh.init_device_mesh( + device_type=get_device_name(), mesh_shape=[2, 2, ngpus // 4], mesh_dim_names=["pp", "dp", "tp"] + ) + + self._register_dispatch_collect_info( + "infer", + dp_rank=self.infer_device_mesh["dp"].get_local_rank(), + is_collect=self.infer_device_mesh["tp"].get_local_rank() == 0, + ) + self._register_dispatch_collect_info( + "train", + dp_rank=self.train_device_mesh["dp"].get_local_rank(), + is_collect=self.train_device_mesh["tp"].get_local_rank() == 0 + and self.train_device_mesh["pp"].get_local_rank() == 1, + ) + + @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="infer")) + def generate_data_proto(self, data: DataProto): + tp_rank = self.infer_device_mesh["tp"].get_local_rank() + dp_rank = self.infer_device_mesh["dp"].get_local_rank() + data.batch["a"] += (tp_rank + 1) * dp_rank + return data + + @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="infer")) + def generate_tensordict(self, data: TensorDict): + tp_rank = self.infer_device_mesh["tp"].get_local_rank() + dp_rank = self.infer_device_mesh["dp"].get_local_rank() + data["a"] += (tp_rank + 1) * dp_rank + return data + + @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="train")) + def train_data_proto(self, data: DataProto): + tp_rank = self.train_device_mesh["tp"].get_local_rank() + dp_rank = self.train_device_mesh["dp"].get_local_rank() + pp_rank = self.train_device_mesh["pp"].get_local_rank() + data.batch["a"] += (tp_rank + 1) * (dp_rank + 2) * (pp_rank + 3) + # tp rank 0, pp rank 1, dp rank 0, output data added: 8 + 3 = 11 + # tp rank 0, pp rank 1, dp rank 1, output data added: 12 + 4 = 16 + return data + + @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="train")) + def train_tensordict(self, data: TensorDict): + tp_rank = self.train_device_mesh["tp"].get_local_rank() + dp_rank = self.train_device_mesh["dp"].get_local_rank() + pp_rank = self.train_device_mesh["pp"].get_local_rank() + data["a"] += (tp_rank + 1) * (dp_rank + 2) * (pp_rank + 3) + # tp rank 0, pp rank 1, dp rank 0, output data added: 8 + 3 = 11 + # tp rank 0, pp rank 1, dp rank 1, output data added: 12 + 4 = 16 + return data + + @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="infer")) + def generate_nested_tensor(self, data: TensorDict): + tp_rank = self.infer_device_mesh["tp"].get_local_rank() + dp_rank = self.infer_device_mesh["dp"].get_local_rank() + assert data.shape[0] == 8 + data["input_ids"] += tp_rank + dp_rank + + print(data) + return data + + +def test_dist_global_info_wg(): + # create a worker group with size 8 + # register a infer dist info with tp=4, dp=2 + # register a train dist info with tp=2, dp=2, pp=2 + # test the correctness of data dispatch and computation + from verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup + + ray.init() + ngpus = torch.cuda.device_count() + infer_tp = ngpus // 2 + train_tp = ngpus // 4 + + ray_cls = RayClassWithInitArgs(TestActor) + resource_pool = RayResourcePool(process_on_nodes=[ngpus]) + wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=ray_cls, device_name=get_device_name()) + + infer_input_data_proto = DataProto.from_single_dict(data={"a": torch.tensor([1, 2])}) + infer_output_data_proto = wg.generate_data_proto(infer_input_data_proto) + + expected_infer_dispatch = [d for d in range(2) for _ in range(infer_tp)] + assert wg._dispatch_info["infer"] == expected_infer_dispatch + + assert torch.all(torch.eq(infer_output_data_proto.batch["a"], torch.tensor([1, 3]))) + + infer_input_tensordict = infer_input_data_proto.to_tensordict() + infer_output_tensordict = wg.generate_tensordict(infer_input_tensordict) + assert torch.all(torch.eq(infer_output_tensordict["a"], torch.tensor([1, 3]))) + + train_input_data_proto = DataProto.from_single_dict(data={"a": torch.tensor([3, 4])}) + train_output_data_proto = wg.train_data_proto(train_input_data_proto) + + expected_train_dispatch = [d for _ in range(2) for d in range(2) for _ in range(train_tp)] + assert wg._dispatch_info["train"] == expected_train_dispatch + + assert torch.all(torch.eq(train_output_data_proto.batch["a"], torch.tensor([11, 16]))) + + train_input_tensordict = train_input_data_proto.to_tensordict() + train_output_tensordict = wg.train_tensordict(train_input_tensordict) + assert torch.all(torch.eq(train_output_tensordict["a"], torch.tensor([11, 16]))) + + # create a batch size of input_ids + input_ids = [ + torch.randint(low=0, high=128, size=(np.random.randint(low=1, high=10, dtype=np.int64),)) for _ in range(16) + ] + input_ids = torch.nested.as_nested_tensor(input_ids, layout=torch.jagged) + data = tu.get_tensordict(tensor_dict={"input_ids": input_ids}) + output = wg.generate_nested_tensor(data) + + input_ids_chunked = list(input_ids.chunk(2)) + + print(input_ids_chunked) + + input_ids_chunked[0] += 0 + input_ids_chunked[1] += 1 + + expected = tu.concat_nested_tensors(input_ids_chunked) + + assert torch.all(torch.eq(output["input_ids"].values(), expected.values())) + + ray.shutdown() + + +if __name__ == "__main__": + test_dist_global_info_wg() diff --git a/verl/tests/single_controller/test_driverfunc_to_worker.py b/verl/tests/single_controller/test_driverfunc_to_worker.py new file mode 100644 index 0000000000000000000000000000000000000000..16f9976067113eb35a8e6f7442569f8e5ca287de --- /dev/null +++ b/verl/tests/single_controller/test_driverfunc_to_worker.py @@ -0,0 +1,85 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os + +import ray +import torch +from tensordict import TensorDict + +from verl import DataProto +from verl.single_controller.base.worker import Worker +from verl.single_controller.ray import RayWorkerGroup +from verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool +from verl.utils.device import get_device_name + +os.environ["RAY_DEDUP_LOGS"] = "0" +os.environ["NCCL_DEBUG"] = "WARN" + + +@ray.remote +class ModelActor(Worker): + def __init__(self): + pass + + +class HackSelf: + def __init__(self): + pass + + +def get_aux_metrics(self, test_proto): + sequence_ids = test_proto.batch["sequence_ids"] + decode_count = [] + for i in range(sequence_ids.size(0)): + decode_count.append(len(sequence_ids[i].tolist())) + ret_proto = DataProto( + batch=TensorDict( + {"sequence_ids": sequence_ids, "decode_count": torch.tensor(decode_count)}, batch_size=sequence_ids.size(0) + ) + ) + return ret_proto + + +def test(): + # construct model + ray.init() + + # create 2 workers, each hold a GPU + resource_pool = RayResourcePool([2], use_gpu=True, name_prefix="a") + + class_with_args = RayClassWithInitArgs(cls=ModelActor) + shard_wg = RayWorkerGroup(resource_pool, class_with_args, device_name=get_device_name()) + + test_bs = 8 + test_proto = DataProto( + TensorDict( + { + "sequence_ids": torch.ones([test_bs, 2048], dtype=torch.int64), + }, + batch_size=test_bs, + ), + meta_info={"query_length": 1536}, + ) + + # Sharding among different ranks + ret_proto1 = shard_wg.execute_with_func_generator(get_aux_metrics, test_proto) + + # compare execute on driver + hs = HackSelf() + ret_proto2 = get_aux_metrics(hs, test_proto) + + torch.testing.assert_close(ret_proto1.batch["decode_count"], ret_proto2.batch["decode_count"]) + + ray.shutdown() diff --git a/verl/tests/single_controller/test_fused_workers_on_cpu.py b/verl/tests/single_controller/test_fused_workers_on_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..527ddc102419bae10f01684a9b4e3e3b13530522 --- /dev/null +++ b/verl/tests/single_controller/test_fused_workers_on_cpu.py @@ -0,0 +1,90 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import ray + +from verl.single_controller.base import Worker +from verl.single_controller.base.decorator import Dispatch, register +from verl.single_controller.ray.base import ( + RayClassWithInitArgs, + RayResourcePool, + RayWorkerGroup, + create_colocated_worker_raw_cls, +) + + +@ray.remote +class Actor(Worker): + def __init__(self) -> None: + super().__init__() + + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def add(self, x): + x += self.rank + return x + + +@ray.remote +class Critic(Worker): + def __init__(self, val) -> None: + super().__init__() + self.val = val + + @register(dispatch_mode=Dispatch.ALL_TO_ALL) + def sub(self, x): + x -= self.val + return x + + +actor_cls = RayClassWithInitArgs(cls=Actor) +critic_cls = RayClassWithInitArgs(cls=Critic, val=10) +cls_dict = {"actor": actor_cls, "critic": critic_cls} +FusedBaseClass = create_colocated_worker_raw_cls(cls_dict) + + +@ray.remote +class HybridWorker(FusedBaseClass): + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def foo(self, x): + return self.critic.sub(self.actor.add(x)) + + +def test_fused_workers(): + ray.init(num_cpus=100) + + # create separate workers on the same resource pool + process_on_nodes = [2] + resource_pool = RayResourcePool(process_on_nodes=process_on_nodes, use_gpu=False) + + # create colocated workers + hybrid_cls_with_init = RayClassWithInitArgs(cls=HybridWorker) + hybrid_cls_with_init.fused_worker_used = True + + fused_wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=hybrid_cls_with_init) + fused_wg.fuse(cls_dict.keys()) + + x = fused_wg.actor.add(0.1) + print(x) + y = fused_wg.critic.sub(x) + print(y) + z = fused_wg.foo(0.1) + print(z) + for i, j in zip(y, z, strict=True): + assert i == j + + ray.shutdown() + + +if __name__ == "__main__": + test_fused_workers() diff --git a/verl/tests/single_controller/test_get_set_dispatch_collect_cpu.py b/verl/tests/single_controller/test_get_set_dispatch_collect_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..2b832da89910d1876fdaed7ad88e02170e5c35c1 --- /dev/null +++ b/verl/tests/single_controller/test_get_set_dispatch_collect_cpu.py @@ -0,0 +1,47 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os + +import pytest + +from verl.single_controller.base import Worker + + +def test_get_set_dispatch_collect_cpu(): + os.environ["RANK"] = "0" + os.environ["LOCAL_RANK"] = "0" + os.environ["WORLD_SIZE"] = "2" + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = "12345" + + ref = Worker() + ref._register_dispatch_collect_info(mesh_name="actor", dp_rank=0, is_collect=True) + + actor = Worker() + actor._register_dispatch_collect_info(mesh_name="actor", dp_rank=1, is_collect=False) + + actor_rollout_ref = Worker() + actor_rollout_ref.set_dispatch_collect(mesh_name="ref", **ref.get_dispatch_collect()) + actor_rollout_ref.set_dispatch_collect(mesh_name="actor", **actor.get_dispatch_collect()) + + assert actor_rollout_ref._query_dispatch_info("ref") == 0 + assert actor_rollout_ref._query_collect_info("ref") + assert actor_rollout_ref._query_dispatch_info("actor") == 1 + assert not actor_rollout_ref._query_collect_info("actor") + + # test conflict mesh_name + actor2 = Worker() + actor2._register_dispatch_collect_info(mesh_name="actor", dp_rank=1, is_collect=False) + with pytest.raises(AssertionError): + actor_rollout_ref.set_dispatch_collect(mesh_name="actor", **actor2.get_dispatch_collect()) diff --git a/verl/tests/single_controller/test_high_level_scheduling_api.py b/verl/tests/single_controller/test_high_level_scheduling_api.py new file mode 100644 index 0000000000000000000000000000000000000000..ce15b8b91266d043094ffc19a79d5ed7c47dcacd --- /dev/null +++ b/verl/tests/single_controller/test_high_level_scheduling_api.py @@ -0,0 +1,97 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import gc +import time + +import ray + +from verl.single_controller.base.worker import Worker +from verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup, merge_resource_pool +from verl.utils.device import get_device_name + + +@ray.remote +class TestActor(Worker): + # TODO: pass *args and **kwargs is bug prone and not very convincing + def __init__(self, cuda_visible_devices=None) -> None: + super().__init__(cuda_visible_devices) + + def get_node_id(self): + return ray.get_runtime_context().get_node_id() + + +def test(): + import torch + + ray.init() + ngpus = torch.cuda.device_count() + half = ngpus // 2 + + # test single-node-no-partition + print("test single-node-no-partition") + resource_pool = RayResourcePool([ngpus], use_gpu=True) + + class_with_args = RayClassWithInitArgs(cls=TestActor) + + print("create actor worker group") + actor_wg = RayWorkerGroup( + resource_pool, class_with_args, name_prefix="high_level_api_actor", device_name=get_device_name() + ) + print("create critic worker group") + critic_wg = RayWorkerGroup( + resource_pool, class_with_args, name_prefix="hight_level_api_critic", device_name=get_device_name() + ) + print("create ref worker group") + ref_wg = RayWorkerGroup( + resource_pool, class_with_args, name_prefix="high_level_api_ref", device_name=get_device_name() + ) + + assert actor_wg.execute_all_sync("get_cuda_visible_devices") == [str(i) for i in range(ngpus)] + assert critic_wg.execute_all_sync("get_cuda_visible_devices") == [str(i) for i in range(ngpus)] + assert ref_wg.execute_all_sync("get_cuda_visible_devices") == [str(i) for i in range(ngpus)] + + del actor_wg + del critic_wg + del ref_wg + gc.collect() # make sure ray actors are deleted + + [ray.util.remove_placement_group(pg) for pg in resource_pool.get_placement_groups()] + print("wait 5s to remove placemeng_group") + time.sleep(5) + # test single-node-multi-partition + + print("test single-node-multi-partition") + rm_resource_pool = RayResourcePool([half], use_gpu=True, name_prefix="rm") + ref_resource_pool = RayResourcePool([half], use_gpu=True, name_prefix="ref") + total_resource_pool = merge_resource_pool(rm_resource_pool, ref_resource_pool) + + assert rm_resource_pool.world_size == half + assert ref_resource_pool.world_size == half + assert total_resource_pool.world_size == ngpus + + actor_wg = RayWorkerGroup( + total_resource_pool, class_with_args, name_prefix="high_level_api_actor", device_name=get_device_name() + ) + critic_wg = RayWorkerGroup( + total_resource_pool, class_with_args, name_prefix="high_level_api_critic", device_name=get_device_name() + ) + ref_wg = RayWorkerGroup( + ref_resource_pool, class_with_args, name_prefix="high_level_api_ref", device_name=get_device_name() + ) + + assert actor_wg.execute_all_sync("get_cuda_visible_devices") == [str(i) for i in range(ngpus)] + assert critic_wg.execute_all_sync("get_cuda_visible_devices") == [str(i) for i in range(ngpus)] + assert ref_wg.execute_all_sync("get_cuda_visible_devices") == [str(i) for i in range(half, ngpus)] + + ray.shutdown() diff --git a/verl/tests/single_controller/test_nested_worker.py b/verl/tests/single_controller/test_nested_worker.py new file mode 100644 index 0000000000000000000000000000000000000000..99145e5949ee9bf03f85f4201f1e025b42b4e200 --- /dev/null +++ b/verl/tests/single_controller/test_nested_worker.py @@ -0,0 +1,75 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import ray + +from verl.single_controller.base.decorator import Dispatch, register +from verl.single_controller.base.worker import Worker +from verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup +from verl.utils.device import get_device_name + + +class TestActor(Worker): + # TODO: pass *args and **kwargs is bug prone and not very convincing + def __init__(self, x) -> None: + super().__init__() + self.a = x + + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def get(self): + return self.a + self.rank + + +class TestHighLevelActor(Worker): + def __init__(self, x=None) -> None: + super().__init__() + self.test_actor = TestActor(x=x) + + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def get(self): + return self.test_actor.get() + + +def test_nested_worker(): + ray.init(num_cpus=100) + + # create 4 workers, each hold a GPU + resource_pool = RayResourcePool([4], use_gpu=True) + class_with_args = RayClassWithInitArgs(cls=ray.remote(TestActor), x=2) + + worker_group = RayWorkerGroup( + resource_pool=resource_pool, + ray_cls_with_init=class_with_args, + name_prefix="worker_group_basic", + device_name=get_device_name(), + ) + + output = worker_group.get() + + assert output == [2, 3, 4, 5] + + class_with_args = RayClassWithInitArgs(cls=ray.remote(TestHighLevelActor), x=2) + high_level_worker_group = RayWorkerGroup( + resource_pool=resource_pool, + ray_cls_with_init=class_with_args, + name_prefix="worker_group_basic_2", + device_name=get_device_name(), + ) + + output_1 = high_level_worker_group.get() + + assert output_1 == [2, 3, 4, 5] + + ray.shutdown() diff --git a/verl/tests/single_controller/test_ray_collectives.py b/verl/tests/single_controller/test_ray_collectives.py new file mode 100644 index 0000000000000000000000000000000000000000..0096e9c15c488e4ec1069f6ee3aae478009eee58 --- /dev/null +++ b/verl/tests/single_controller/test_ray_collectives.py @@ -0,0 +1,113 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Test for using ray collective group. +Suppose we Actor and Rollout. Actor contains 4 workers and Rollout contains 2 workers. We established a Worker to +Rollout relationship by using collective groups +Actor: rank 0, 1 - Rollout rank 0 +Rollout rank 2, 3 - Rollout rank 1 +Then, we initiate 4 p2p comms from actor to rollout +""" + +import ray +import ray.util.collective as collective +import torch + +from verl.single_controller.base import Worker +from verl.single_controller.base.decorator import Dispatch, register +from verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup + + +@ray.remote +class Actor(Worker): + @register(Dispatch.ONE_TO_ALL) + def init(self): + remote_rank = self.rank // 2 + self.group_name = f"A{self.rank}_R{remote_rank}" + collective.init_collective_group(world_size=2, rank=0, backend="nccl", group_name=self.group_name) + + @register(Dispatch.ONE_TO_ALL, blocking=False) + def send_tensors(self): + tensor = torch.ones(size=(4,), dtype=torch.float32, device="cuda") * self.rank + collective.send(tensor=tensor, dst_rank=1, group_name=self.group_name) + + +@ray.remote +class Rollout(Worker): + @register(Dispatch.ONE_TO_ALL) + def init(self): + self.remote_first_rank = self.rank * 2 + self.remote_second_rank = self.remote_first_rank + 1 + self.first_group_name = f"A{self.remote_first_rank}_R{self.rank}" + self.second_group_name = f"A{self.remote_second_rank}_R{self.rank}" + + collective.init_collective_group(world_size=2, rank=1, backend="nccl", group_name=self.first_group_name) + collective.init_collective_group(world_size=2, rank=1, backend="nccl", group_name=self.second_group_name) + + @register(Dispatch.ONE_TO_ALL, blocking=False) + def receive_tensors(self): + self.tensor1 = torch.randn(size=(4,), dtype=torch.float32, device="cuda") + self.tensor2 = torch.randn(size=(4,), dtype=torch.float32, device="cuda") + + collective.recv(self.tensor1, src_rank=0, group_name=self.first_group_name) + collective.recv(self.tensor2, src_rank=0, group_name=self.second_group_name) + + @register(Dispatch.ONE_TO_ALL) + def get_tensors(self): + return {f"src_{self.remote_first_rank}": self.tensor1, f"src_{self.remote_second_rank}": self.tensor2} + + +def test_ray_collective_group(): + ray.init() + ngpus = torch.cuda.device_count() + n_rollout = max(1, ngpus // 3) # keep 2:1 actor:rollout ratio + n_actor = n_rollout * 2 + + actor_resource_pool = RayResourcePool([n_actor]) + rollout_resource_pool = RayResourcePool([n_rollout]) + + actor_cls = RayClassWithInitArgs(cls=Actor) + rollout_cls = RayClassWithInitArgs(cls=Rollout) + + actor_wg = RayWorkerGroup( + resource_pool=actor_resource_pool, ray_cls_with_init=actor_cls, name_prefix="collective_group_actor" + ) + rollout_wg = RayWorkerGroup( + resource_pool=rollout_resource_pool, ray_cls_with_init=rollout_cls, name_prefix="collective_group_rollout" + ) + + actor_wg.init() + rollout_wg.init() + + out1 = actor_wg.send_tensors() + out2 = rollout_wg.receive_tensors() + + # block to wait + ray.get(out1) + ray.get(out2) + + output = {} + for d in rollout_wg.get_tensors(): + output |= d + + print(output) + + for i in range(n_actor): + assert torch.sum(output[f"src_{i}"]).item() == 4 * i + + ray.shutdown() + + +if __name__ == "__main__": + test_ray_collective_group() diff --git a/verl/tests/single_controller/test_ray_local_envs_on_cpu.py b/verl/tests/single_controller/test_ray_local_envs_on_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..6c51beeaf3f8600387ce14fe63c97a5c804c4237 --- /dev/null +++ b/verl/tests/single_controller/test_ray_local_envs_on_cpu.py @@ -0,0 +1,91 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +e2e test verl.single_controller.ray +""" + +import os + +import ray + +from verl.single_controller.base.worker import Worker +from verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup + + +@ray.remote +class TestActor(Worker): + def __init__(self) -> None: + super().__init__() + + def getenv(self, key): + val = os.getenv(key, f"{key} not set") + return val + + +def test_basics(): + ray.init(num_cpus=100) + + # create 4 workers, each hold a GPU + resource_pool = RayResourcePool([4], use_gpu=False) + class_with_args = RayClassWithInitArgs(cls=TestActor) + + worker_group = RayWorkerGroup( + resource_pool=resource_pool, ray_cls_with_init=class_with_args, name_prefix="worker_group_basic" + ) + + output = worker_group.execute_all_sync("getenv", key="RAY_LOCAL_WORLD_SIZE") + assert output == ["4", "4", "4", "4"] + + ray.shutdown() + + +def test_customized_worker_env(): + ray.init(num_cpus=100) + + # create 4 workers, each hold a GPU + resource_pool = RayResourcePool([4], use_gpu=False) + class_with_args = RayClassWithInitArgs(cls=TestActor) + + worker_group = RayWorkerGroup( + resource_pool=resource_pool, + ray_cls_with_init=class_with_args, + name_prefix="worker_group_customized", + worker_env={ + "test_key": "test_value", # new key will be appended + }, + ) + + output = worker_group.execute_all_sync("getenv", key="test_key") + assert output == ["test_value", "test_value", "test_value", "test_value"] + + try: + worker_group = RayWorkerGroup( + resource_pool=resource_pool, + ray_cls_with_init=class_with_args, + name_prefix="worker_group_error", + worker_env={ + "WORLD_SIZE": "100", # override system env will result in error + }, + ) + except ValueError as e: + assert "WORLD_SIZE" in str(e) + else: + raise ValueError("test failed") + + ray.shutdown() + + +if __name__ == "__main__": + test_basics() + test_customized_worker_env() diff --git a/verl/tests/single_controller/test_ray_utils_on_cpu.py b/verl/tests/single_controller/test_ray_utils_on_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..e36497d210f6ec5daa8b9d559987f5dcc3974af2 --- /dev/null +++ b/verl/tests/single_controller/test_ray_utils_on_cpu.py @@ -0,0 +1,54 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import ray + +from verl.utils.ray_utils import parallel_put + + +# Initialize Ray for testing if not already done globally +@pytest.fixture() +def init_ray(): + ray.init(num_cpus=4) + yield + ray.shutdown() + + +def test_parallel_put_basic(init_ray): + data = [1, "hello", {"a": 2}, [3, 4]] + refs = parallel_put(data) + assert len(refs) == len(data) + retrieved_data = [ray.get(ref) for ref in refs] + assert retrieved_data == data + + +def test_parallel_put_empty(init_ray): + data = [] + with pytest.raises(AssertionError): + _ = parallel_put(data) + + +def test_parallel_put_workers(init_ray): + data = list(range(20)) + # Test with specific number of workers + refs = parallel_put(data, max_workers=4) + assert len(refs) == len(data) + retrieved_data = [ray.get(ref) for ref in refs] + assert retrieved_data == data + # Test with default workers (should cap) + refs_default = parallel_put(data) + assert len(refs_default) == len(data) + retrieved_data_default = [ray.get(ref) for ref in refs_default] + assert retrieved_data_default == data diff --git a/verl/tests/single_controller/test_rvdz.py b/verl/tests/single_controller/test_rvdz.py new file mode 100644 index 0000000000000000000000000000000000000000..7dea12f95cd5cb697f5fcfa20a844331bd46e8f3 --- /dev/null +++ b/verl/tests/single_controller/test_rvdz.py @@ -0,0 +1,51 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import ray + + +@ray.remote +class TestWorker: + def __init__(self, rank, world_size, group_name): + self.rank = rank + self.world_size = world_size + self.group_name = group_name + self.communicator = None + + def init(self): + from verl.utils.rendezvous.ray_backend import create_nccl_communicator_in_ray + + self.communicator = create_nccl_communicator_in_ray(self.rank, self.world_size, self.group_name) + + def test(self): + if self.communicator is None: + return None + return self.communicator.rank_id() + + +def test_rvdz(): + ray.init() + + group_name = "test_group" + world_size = 2 + + workers = [TestWorker.options(num_gpus=1).remote(rank, world_size, group_name) for rank in range(world_size)] + + ray.get([worker.init.remote() for worker in workers]) + + ranks = ray.get([worker.test.remote() for worker in workers]) + + assert ranks == [0, 1], f"expecting [0, 1], got {ranks}" + + ray.shutdown() diff --git a/verl/tests/single_controller/test_split_resource_pool.py b/verl/tests/single_controller/test_split_resource_pool.py new file mode 100644 index 0000000000000000000000000000000000000000..1201c0b3b96f0c3386c9764ff222b2dfc3c3141c --- /dev/null +++ b/verl/tests/single_controller/test_split_resource_pool.py @@ -0,0 +1,185 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os + +import ray +import torch + +from verl import DataProto +from verl.single_controller.base import Worker +from verl.single_controller.base.decorator import Dispatch, register +from verl.single_controller.ray.base import ( + RayClassWithInitArgs, + RayResourcePool, + RayWorkerGroup, + split_resource_pool, +) +from verl.utils.device import get_device_name, get_nccl_backend + + +def get_local_gpus_num(division=1): + return max(1, torch.cuda.device_count() // division) + + +@ray.remote +class Actor(Worker): + def __init__(self, worker_id) -> None: + super().__init__() + self.worker_id = worker_id + self.temp_tensor = torch.rand(4096, 4096).to(get_device_name()) + + if not torch.distributed.is_initialized(): + rank = int(os.environ.get("RANK", 0)) + world_size = int(os.environ.get("WORLD_SIZE", 1)) + torch.distributed.init_process_group(backend=get_nccl_backend(), world_size=world_size, rank=rank) + + @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO) + def add(self, data: DataProto): + data.batch["a"] += self.rank + self.worker_id + return data + + +def test_split_resource_pool_with_split_size(): + ray.init() + ngpus = torch.cuda.device_count() + half = get_local_gpus_num(2) + # simulate 2 nodes of half GPUs each + global_resource_pool = RayResourcePool(process_on_nodes=[half, half]) + global_resource_pool.get_placement_groups(device_name=get_device_name()) + + actor_1_resource_pool, actor_2_resource_pool = split_resource_pool( + resource_pool=global_resource_pool, split_size=half + ) + actor_cls_1 = RayClassWithInitArgs(cls=Actor, worker_id=0) + actor_cls_2 = RayClassWithInitArgs(cls=Actor, worker_id=100) + actor_worker_1 = RayWorkerGroup( + resource_pool=actor_1_resource_pool, ray_cls_with_init=actor_cls_1, device_name=get_device_name() + ) + actor_worker_2 = RayWorkerGroup( + resource_pool=actor_2_resource_pool, ray_cls_with_init=actor_cls_2, device_name=get_device_name() + ) + assert actor_worker_1.world_size == half + assert actor_worker_2.world_size == half + + data = DataProto.from_dict({"a": torch.zeros(ngpus)}) + actor_output_1 = actor_worker_1.add(data) + actor_output_2 = actor_worker_2.add(data) + assert actor_output_1.batch["a"].tolist() == [float(r) for r in range(half) for _ in range(2)] + assert actor_output_2.batch["a"].tolist() == [float(r + 100) for r in range(half) for _ in range(2)] + + ray.shutdown() + + +def test_split_resource_pool_with_split_size_list(): + ray.init() + quarter = get_local_gpus_num(4) + # simulate 4 nodes of quarter GPUs each + global_resource_pool = RayResourcePool(process_on_nodes=[quarter] * 4) + global_resource_pool.get_placement_groups(device_name=get_device_name()) + + actor_1_resource_pool, actor_2_resource_pool = split_resource_pool( + resource_pool=global_resource_pool, + split_size=[quarter, 3 * quarter], + ) + actor_cls_1 = RayClassWithInitArgs(cls=Actor, worker_id=0) + actor_cls_2 = RayClassWithInitArgs(cls=Actor, worker_id=100) + actor_worker_1 = RayWorkerGroup( + resource_pool=actor_1_resource_pool, ray_cls_with_init=actor_cls_1, device_name=get_device_name() + ) + actor_worker_2 = RayWorkerGroup( + resource_pool=actor_2_resource_pool, ray_cls_with_init=actor_cls_2, device_name=get_device_name() + ) + assert actor_worker_1.world_size == quarter + assert actor_worker_2.world_size == 3 * quarter + + data_1 = DataProto.from_dict({"a": torch.zeros(quarter)}) + data_2 = DataProto.from_dict({"a": torch.zeros(3 * quarter)}) + actor_output_1 = actor_worker_1.add(data_1) + actor_output_2 = actor_worker_2.add(data_2) + print(actor_output_1.batch["a"].tolist()) + print(actor_output_2.batch["a"].tolist()) + assert actor_output_1.batch["a"].tolist() == list(range(quarter)) + assert actor_output_2.batch["a"].tolist() == list(range(100, 100 + 3 * quarter)) + + ray.shutdown() + + +def test_split_resource_pool_with_split_size_list_cross_nodes(): + ray.init() + half = get_local_gpus_num(2) + quarter = get_local_gpus_num(4) + # simulate 2 nodes of half GPUs each (cross-node split) + global_resource_pool = RayResourcePool(process_on_nodes=[half, half]) + global_resource_pool.get_placement_groups(device_name=get_device_name()) + + actor_1_resource_pool, actor_2_resource_pool = split_resource_pool( + resource_pool=global_resource_pool, + split_size=[quarter, 3 * quarter], + ) + actor_cls_1 = RayClassWithInitArgs(cls=Actor, worker_id=0) + actor_cls_2 = RayClassWithInitArgs(cls=Actor, worker_id=100) + actor_worker_1 = RayWorkerGroup( + resource_pool=actor_1_resource_pool, ray_cls_with_init=actor_cls_1, device_name=get_device_name() + ) + actor_worker_2 = RayWorkerGroup( + resource_pool=actor_2_resource_pool, ray_cls_with_init=actor_cls_2, device_name=get_device_name() + ) + + assert actor_worker_1.world_size == quarter + assert actor_worker_2.world_size == 3 * quarter + + data_1 = DataProto.from_dict({"a": torch.zeros(quarter)}) + data_2 = DataProto.from_dict({"a": torch.zeros(3 * quarter)}) + actor_output_1 = actor_worker_1.add(data_1) + actor_output_2 = actor_worker_2.add(data_2) + print(actor_output_1.batch["a"].tolist()) + print(actor_output_2.batch["a"].tolist()) + assert actor_output_1.batch["a"].tolist() == list(range(quarter)) + assert actor_output_2.batch["a"].tolist() == list(range(100, 100 + 3 * quarter)) + + ray.shutdown() + + +def test_split_resource_pool_with_split_twice(): + ray.init() + ngpus = torch.cuda.device_count() + quarter = get_local_gpus_num(4) + mid = ngpus - 2 * quarter # middle pool size + # simulate ngpus//2 nodes of 2 GPUs each + global_resource_pool = RayResourcePool(process_on_nodes=[2] * (ngpus // 2)) + global_resource_pool.get_placement_groups(device_name=get_device_name()) + + rp_1, rp_2, rp_3 = split_resource_pool( + resource_pool=global_resource_pool, + split_size=[quarter, mid, quarter], + ) + rp_2_subs = split_resource_pool(resource_pool=rp_2, split_size=1) + fp_list = [rp_1] + list(rp_2_subs) + [rp_3] + + correct_world_size = [quarter] + [1] * mid + [quarter] + correct_output = [] + for ws in correct_world_size: + idx = len(correct_output) + correct_output.append([float(r + idx * 100) for r in range(ws) for _ in range(4 // ws)]) + + for idx, rp in enumerate(fp_list): + actor_cls = RayClassWithInitArgs(cls=Actor, worker_id=idx * 100) + actor_worker = RayWorkerGroup(resource_pool=rp, ray_cls_with_init=actor_cls, device_name=get_device_name()) + data = DataProto.from_dict({"a": torch.zeros(4)}) + actor_output = actor_worker.add(data) + assert actor_worker.world_size == correct_world_size[idx] + assert actor_output.batch["a"].tolist() == correct_output[idx] + + ray.shutdown() diff --git a/verl/tests/single_controller/test_worker_group_basics.py b/verl/tests/single_controller/test_worker_group_basics.py new file mode 100644 index 0000000000000000000000000000000000000000..40090e7b8a73fdbfa89d697f3a4a8025247e8a4a --- /dev/null +++ b/verl/tests/single_controller/test_worker_group_basics.py @@ -0,0 +1,147 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +e2e test verl.single_controller.ray +""" + +import ray +import torch + +from verl.single_controller.base.decorator import Dispatch, Execute, collect_all_to_all, register +from verl.single_controller.base.worker import Worker +from verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup +from verl.utils.device import get_device_name + + +def two_to_all_dispatch_fn(worker_group, *args, **kwargs): + """ + Assume the input is a list of 2. Duplicate the input interleaved and pass to each worker. + """ + for arg in args: + assert len(arg) == 2 + for i in range(worker_group.world_size - 2): + arg.append(arg[i % 2]) + for k, v in kwargs.items(): + assert len(v) == 2 + for i in range(worker_group.world_size - 2): + v.append(v[i % 2]) + return args, kwargs + + +def get_ray_remote_options() -> str: + """Function that gets the torch.device based on the current machine. + This currently only supports CPU, CUDA, NPU. + Returns: + device + """ + if get_device_name() == "cuda": + return dict(num_gpus=0.1) + elif get_device_name() == "npu": + return dict(resources={"NPU": 0.1}) + return dict(num_cpus=0.1) + + +@ray.remote +class TestActor(Worker): + # TODO: pass *args and **kwargs is bug prone and not very convincing + def __init__(self, x) -> None: + super().__init__() + self._x = x + + def foo(self, y): + return self._x + y + + @register(dispatch_mode=Dispatch.ALL_TO_ALL, execute_mode=Execute.RANK_ZERO) + def foo_rank_zero(self, x, y): + return self._x + y + x + + @register(Dispatch.ONE_TO_ALL, blocking=False) + def foo_one_to_all(self, x, y): + return self._x + y + x + + @register(Dispatch.ALL_TO_ALL, blocking=False) + def foo_all_to_all(self, x, y): + return self._x + y + x + + @register(dispatch_mode={"dispatch_fn": two_to_all_dispatch_fn, "collect_fn": collect_all_to_all}) + def foo_custom(self, x, y): + return self._x + y + x + + +@ray.remote(num_cpus=0.1) +def remote_call_wg(worker_names): + class_with_args = RayClassWithInitArgs(cls=TestActor, x=2) + worker_group = RayWorkerGroup.from_detached( + worker_names=worker_names, ray_cls_with_init=class_with_args, name_prefix=None + ) + print(worker_group.worker_names) + + output_ref = worker_group.foo_custom(x=[1, 2], y=[5, 6]) + assert output_ref == [8, 10, 8, 10] + + output_ref = worker_group.foo_rank_zero(x=1, y=2) + assert output_ref == 5 + + return worker_group.worker_names + + +def add_one(data): + data = data.to(get_device_name()) + data += 1 + data = data.to("cpu") + return data + + +def test_basics(): + ray.init(num_cpus=100) + + # create 4 workers, each hold a GPU + resource_pool = RayResourcePool([4], use_gpu=True) + class_with_args = RayClassWithInitArgs(cls=TestActor, x=2) + + worker_group = RayWorkerGroup( + resource_pool=resource_pool, + ray_cls_with_init=class_with_args, + name_prefix="worker_group_basic", + device_name=get_device_name(), + ) + + print(worker_group.worker_names) + + # this will wait for all the results + output = worker_group.execute_all_sync("foo", y=3) + assert output == [5, 5, 5, 5] + + # this is a list of object reference. It won't block. + output_ref = worker_group.execute_all_async("foo", y=4) + print(output_ref) + + assert ray.get(output_ref) == [6, 6, 6, 6] + + output_ref = worker_group.foo_one_to_all(x=1, y=2) + assert ray.get(output_ref) == [5, 5, 5, 5] + + output_ref = worker_group.foo_all_to_all(x=[1, 2, 3, 4], y=[5, 6, 7, 8]) + assert ray.get(output_ref) == [8, 10, 12, 14] + + print(ray.get(remote_call_wg.remote(worker_group.worker_names))) + + output = worker_group.execute_func_rank_zero(add_one, torch.ones(2, 2)) + torch.testing.assert_close(output, torch.ones(2, 2) + 1) + + ray.shutdown() + + +if __name__ == "__main__": + test_basics() diff --git a/verl/tests/single_controller/test_worker_group_torch.py b/verl/tests/single_controller/test_worker_group_torch.py new file mode 100644 index 0000000000000000000000000000000000000000..6a1b0f29ebec23936d911d43daa92612009af71c --- /dev/null +++ b/verl/tests/single_controller/test_worker_group_torch.py @@ -0,0 +1,116 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os + +os.environ["RAY_DEDUP_LOGS"] = "0" +os.environ["NCCL_DEBUG"] = "WARN" + +import ray +import torch +import torch.distributed + +from verl.single_controller.base.worker import Worker +from verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup +from verl.utils.device import get_device_name + + +@ray.remote +class TestAllGatherActor(Worker): + def __init__(self, size) -> None: + super().__init__() + self.size = size + + def init(self): + torch.distributed.init_process_group() + self.tensor = torch.zeros(size=(self.size,), dtype=torch.int64, device=get_device_name()) + self.tensor += self.rank + + def all_gather(self): + world_size = self._world_size + output = torch.zeros( + size=(self.tensor.shape[0] * world_size,), dtype=self.tensor.dtype, device=self.tensor.device + ) + torch.distributed.all_gather_into_tensor(output, self.tensor, async_op=False) + return output + + +@ray.remote +class TestAllGatherActorV2(Worker): + def __init__(self, size) -> None: + super().__init__() + self.size = size + + torch.distributed.init_process_group() + self.tensor = torch.zeros(size=(self.size,), dtype=torch.int64, device=get_device_name()) + self.tensor += self.rank + + def all_gather(self): + world_size = self._world_size + output = torch.zeros( + size=(self.tensor.shape[0] * world_size,), dtype=self.tensor.dtype, device=self.tensor.device + ) + torch.distributed.all_gather_into_tensor(output, self.tensor, async_op=False) + return output + + +def test_all_gather_torch(): + """ + In this test, we instantiate 4 GPUs in a group and test the all_gather + """ + ray.init() + + # create 4 workers, each hold a GPU + resource_pool = RayResourcePool([4], use_gpu=True) + class_with_args = RayClassWithInitArgs(cls=TestAllGatherActor, size=2) + + worker_group = RayWorkerGroup( + resource_pool, class_with_args, name_prefix="worker_group_torch", device_name=get_device_name() + ) + + worker_group.execute_all_sync("init") + output = worker_group.execute_all_sync("all_gather") + for i in range(1, len(output)): + assert torch.all(output[i] == output[0]) + + output = output[0].cpu() + print(output) + assert torch.all(output == torch.tensor([0, 0, 1, 1, 2, 2, 3, 3], dtype=torch.int64)) + + ray.shutdown() + + +def test_all_gather_torch_v2(): + """ + In this test, we instantiate 4 GPUs in a group and test the all_gather + """ + ray.init() + + # create 4 workers, each hold a GPU + resource_pool = RayResourcePool([4], use_gpu=True) + class_with_args = RayClassWithInitArgs(cls=TestAllGatherActorV2, size=2) + + worker_group = RayWorkerGroup( + resource_pool, class_with_args, name_prefix="worker_group_torch", device_name=get_device_name() + ) + + output = worker_group.execute_all_sync("all_gather") + for i in range(1, len(output)): + assert torch.all(output[i] == output[0]) + + output = output[0].cpu() + print(output) + assert torch.all(output == torch.tensor([0, 0, 1, 1, 2, 2, 3, 3], dtype=torch.int64)) + + ray.shutdown() diff --git a/verl/tests/special_distributed/README.md b/verl/tests/special_distributed/README.md new file mode 100644 index 0000000000000000000000000000000000000000..f2f865e8bf95a673a0d6f56b74c7a2c12535faf2 --- /dev/null +++ b/verl/tests/special_distributed/README.md @@ -0,0 +1 @@ +This folder is reserved for unit tests (instead of end-to-end tests) that require multiple GPUs. diff --git a/verl/tests/special_distributed/run_all.sh b/verl/tests/special_distributed/run_all.sh new file mode 100644 index 0000000000000000000000000000000000000000..3d6c5c71e54a1d6000025840b1abc783f56b60d5 --- /dev/null +++ b/verl/tests/special_distributed/run_all.sh @@ -0,0 +1,19 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +#!/usr/bin/env bash + +set -e -x +torchrun --nproc-per-node=4 --standalone tests/special_distributed/test_tensor_dict.py +torchrun --nproc-per-node=4 --standalone tests/special_distributed/test_torch_functional.py diff --git a/verl/tests/special_distributed/test_fsdp_ckpt.py b/verl/tests/special_distributed/test_fsdp_ckpt.py new file mode 100644 index 0000000000000000000000000000000000000000..4c9b497c47cb9359efb6c9c598391ffb0493cb40 --- /dev/null +++ b/verl/tests/special_distributed/test_fsdp_ckpt.py @@ -0,0 +1,165 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os +import shutil +import tempfile + +import torch +import torch.distributed +from torch.distributed import init_device_mesh +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +from torch.distributed.fsdp import MixedPrecision, ShardingStrategy +from transformers import AutoModelForCausalLM, AutoTokenizer, Qwen2Config + +from verl.utils.checkpoint.fsdp_checkpoint_manager import FSDPCheckpointManager +from verl.utils.device import get_device_name, get_torch_device +from verl.utils.distributed import initialize_global_process_group +from verl.utils.fsdp_utils import MixedPrecisionPolicy, apply_fsdp2 + + +def create_random_input_ids(batch_size, seq_len, vocab_size): + if get_device_name() == "cuda": + from flash_attn.bert_padding import unpad_input + elif get_device_name() == "npu": + from verl.utils.attention_utils import unpad_input + from verl.utils.model import compute_position_id_with_mask, create_random_mask + + input_ids = torch.randint(0, vocab_size, (batch_size, seq_len), device=get_device_name()) + + attention_mask = create_random_mask( + input_ids, max_ratio_of_left_padding=0.1, min_ratio_of_valid_token=0.5, max_ratio_of_valid_token=0.7 + ) + position_ids = compute_position_id_with_mask(attention_mask) + + input_ids = unpad_input(input_ids.unsqueeze(-1), attention_mask)[0].transpose(0, 1) + position_ids = unpad_input(position_ids.unsqueeze(-1), attention_mask)[0].transpose(0, 1) + return input_ids, position_ids + + +def test_fsdp_ckpt(strategy="fsdp"): + assert get_torch_device().device_count() >= 2, "need at least 2 gpus for test" + local_rank, rank, world_size = initialize_global_process_group() + device_mesh = init_device_mesh(get_device_name(), mesh_shape=(world_size,), mesh_dim_names=("dp",)) + + model_name = os.path.expanduser("~/models/Qwen/Qwen2.5-0.5B-Instruct") + config = Qwen2Config(num_hidden_layers=1) + + with torch.device(get_device_name()): + model = AutoModelForCausalLM.from_config( + config=config, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2" + ) + model = model.to(device=get_device_name()) + + # Wrap model with FSDP + if strategy == "fsdp": + mixed_precision = MixedPrecision( + param_dtype=torch.bfloat16, reduce_dtype=torch.float32, buffer_dtype=torch.float32 + ) + + model = FSDP( + model, + use_orig_params=False, + device_id=get_torch_device().current_device(), + sharding_strategy=ShardingStrategy.FULL_SHARD, + mixed_precision=mixed_precision, + device_mesh=device_mesh, + ) + else: + mp_policy = MixedPrecisionPolicy( + param_dtype=torch.bfloat16, reduce_dtype=torch.float32, cast_forward_inputs=True + ) + fsdp_kwargs = { + "mesh": device_mesh, + "mp_policy": mp_policy, + } + apply_fsdp2(model, fsdp_kwargs, {}) + + optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4) + lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.9) + + # Create checkpoint manager + tokenizer = AutoTokenizer.from_pretrained(model_name) + checkpoint_manager = FSDPCheckpointManager( + model=model, optimizer=optimizer, lr_scheduler=lr_scheduler, tokenizer=tokenizer + ) + + # Generate sample input + batch_size = 10 + seq_len = 1024 + vocab_size = config.vocab_size + # First input for initial update + input_ids1, position_ids1 = create_random_input_ids(batch_size, seq_len, vocab_size) + + # Second input for verification + input_ids2, position_ids2 = create_random_input_ids(batch_size, seq_len, vocab_size) + + # Step 1: Initial update and save checkpoint + outputs1 = model(input_ids=input_ids1, position_ids=position_ids1) + loss1 = outputs1.logits.mean() + loss1.backward() + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + # Save checkpoint after first update + temp_dir = tempfile.mkdtemp() + checkpoint_path = os.path.join(temp_dir, "checkpoint") + checkpoint_manager.save_checkpoint(local_path=checkpoint_path, hdfs_path=None, global_step=0) + saved_state_dict = model.state_dict() + + # Step 2: Second update and forward pass + outputs2 = model(input_ids=input_ids2, position_ids=position_ids2) + loss2 = outputs2.logits.mean() + loss2.backward() + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + # Record logits after second update + with torch.no_grad(): + logits_before_load = model(input_ids=input_ids2, position_ids=position_ids2).logits + + # Step 3: Load checkpoint and repeat second update + checkpoint_manager.load_checkpoint(checkpoint_path) + loaded_state_dict = model.state_dict() + for key in loaded_state_dict: + assert key in saved_state_dict, f"Key {key} not found in saved state dict" + torch.testing.assert_close(loaded_state_dict[key], saved_state_dict[key], atol=0.0, rtol=0.0) + + # Repeat the second update with same input + outputs3 = model(input_ids=input_ids2, position_ids=position_ids2) + loss3 = outputs3.logits.mean() + loss3.backward() + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + # Record logits after loaded checkpoint and update + with torch.no_grad(): + logits_after_load = model(input_ids=input_ids2, position_ids=position_ids2).logits + + # Step 4: Verify outputs match + torch.testing.assert_close(logits_before_load, logits_after_load, atol=0.0, rtol=0.0) + print("Checkpoint save/load test passed!") + + # Cleanup + shutil.rmtree(temp_dir) + torch.distributed.barrier() + torch.distributed.destroy_process_group() + + +if __name__ == "__main__": + strategy = os.environ.get("STRATEGY", "fsdp") + os.environ["FLASH_ATTENTION_DETERMINISTIC"] = "1" + test_fsdp_ckpt(strategy=strategy) diff --git a/verl/tests/special_distributed/test_fused_kernels_ulysses_sp.py b/verl/tests/special_distributed/test_fused_kernels_ulysses_sp.py new file mode 100644 index 0000000000000000000000000000000000000000..5767b2a4c4bec40826618ae63395edb773298d6e --- /dev/null +++ b/verl/tests/special_distributed/test_fused_kernels_ulysses_sp.py @@ -0,0 +1,169 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""End-to-end regression test for issue #6068 (use_fused_kernels=True + ulysses_sp>1). + +Construction: +1. Build a small Llama-style model patched with `use_fused_kernels=True` and SP=2. +2. Run the fused forward through Ulysses (rank-local input_ids_rmpad slice + the + engine-style globally-rolled `shift_labels`) and compare its `log_probs` + against a single-GPU reference (SP=1, same fused kernel) gathered across ranks. +3. Without the fix, every rank's last-position log_prob diverges from the + reference because the fused forward locally re-rolls the SP-sliced input. + With the fix, log_probs are bitwise close. + +Run on 2 GPUs: + torchrun --nproc_per_node=2 -m pytest -svv \ + tests/special_distributed/test_fused_kernels_ulysses_sp.py +""" + +import pytest +import torch +import torch.distributed +from torch.distributed import init_device_mesh +from transformers import AutoModelForCausalLM, LlamaConfig + +from verl.models.transformers.monkey_patch import apply_monkey_patch +from verl.utils.device import get_device_name, get_torch_device +from verl.utils.distributed import initialize_global_process_group +from verl.utils.ulysses import ( + FSDPUlyssesShardingManager, + gather_outputs_and_unpad, + set_ulysses_sequence_parallel_group, + ulysses_pad_and_slice_inputs, +) + + +def _make_model(seed: int = 0): + config = LlamaConfig( + vocab_size=256, + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=8, + num_key_value_heads=8, + ) + torch.manual_seed(seed) + model = AutoModelForCausalLM.from_config( + config, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2" + ) + return model, config + + +def _broadcast_params(model): + for p in model.parameters(): + torch.distributed.broadcast(p.data, src=0) + + +@pytest.fixture(scope="module", autouse=True) +def _init_dist(): + if not torch.distributed.is_initialized(): + initialize_global_process_group() + yield + + +def test_fused_kernels_log_probs_match_under_sp(): + """log_probs from fused-kernel + SP=2 must equal fused-kernel + SP=1. + + This is the direct regression test for #6068. The bug was that the + fused-forward functions re-rolled the SP-sliced `input_ids` locally, + producing wrong labels at the shard boundary. The fix passes + `shift_labels` (engine-rolled, then SP-sliced) so the fused forward + uses correct labels. + """ + assert get_torch_device().device_count() >= 2, "need at least 2 gpus" + sp_size = 2 + device = get_device_name() + + # Build identical model on every rank. + model, config = _make_model(seed=42) + apply_monkey_patch( + model, + ulysses_sp_size=sp_size, + use_remove_padding=True, + use_fused_kernels=True, + fused_kernels_backend="torch", + ) + model = model.to(device=device, dtype=torch.bfloat16) + _broadcast_params(model) + + # rank-0-authoritative input. total_nnz must be divisible by sp_size. + total_nnz = 32 + if torch.distributed.get_rank() == 0: + input_ids = torch.randint(0, config.vocab_size, (1, total_nnz), device=device) + else: + input_ids = torch.empty((1, total_nnz), dtype=torch.long, device=device) + torch.distributed.broadcast(input_ids, src=0) + position_ids = torch.arange(total_nnz, device=device).unsqueeze(0) # (1, total_nnz) + + # ---- SP=1 reference: full sequence, fused forward, on rank 0 only ---- + set_ulysses_sequence_parallel_group(None) + if torch.distributed.get_rank() == 0: + ref_out = model( + input_ids=input_ids, + position_ids=position_ids, + use_cache=False, + return_dict=True, + temperature=1.0, + ) + ref_log_probs = ref_out.log_probs.detach().clone() # (1, total_nnz) + else: + ref_log_probs = torch.empty((1, total_nnz), device=device, dtype=torch.float32) + torch.distributed.broadcast(ref_log_probs, src=0) + + # ---- SP=2 path: each rank gets a slice; engine plumbs shift_labels ---- + world_size = torch.distributed.get_world_size() + assert world_size % sp_size == 0 + dp_size = world_size // sp_size + device_mesh = init_device_mesh(device_type=device, mesh_shape=(dp_size, sp_size), mesh_dim_names=("dp", "sp")) + sharding_manager = FSDPUlyssesShardingManager(device_mesh) + + with sharding_manager: + # Mimic the engine's preparation (transformer_impl.py ~951-984). + input_ids_rmpad = input_ids # (1, total_nnz), already "rmpadded" + input_ids_rmpad_rolled = torch.roll(input_ids_rmpad, shifts=-1, dims=1) + # Slice both inputs and rolled labels exactly like the engine. + # `position_ids` is already (1, total_nnz); helper expects 2D, not 3D. + sliced_input, sliced_pos, pad_size = ulysses_pad_and_slice_inputs( + input_ids_rmpad, + position_ids_rmpad=position_ids, + sp_size=sp_size, + ) + sliced_rolled, _, _ = ulysses_pad_and_slice_inputs( + input_ids_rmpad_rolled, position_ids_rmpad=None, sp_size=sp_size + ) + sp_out = model( + input_ids=sliced_input, + position_ids=sliced_pos, + use_cache=False, + return_dict=True, + temperature=1.0, + shift_labels=sliced_rolled, + ) + # log_probs is (1, total_nnz/sp + pad). Gather across SP ranks → (1, total_nnz). + sp_log_probs_local = sp_out.log_probs # (1, local_len) + sp_log_probs = gather_outputs_and_unpad( + sp_log_probs_local.squeeze(0), + gather_dim=0, + unpad_dim=0, + padding_size=pad_size, + ).unsqueeze(0) # (1, total_nnz) + + # bf16 matmul + chunked fused kernel — keep tolerances loose enough to absorb + # numerical noise but strict enough to catch the off-by-one label bug, which + # produces O(1) errors in the affected positions. + torch.testing.assert_close(sp_log_probs, ref_log_probs, atol=1e-3, rtol=1e-3) + + +if __name__ == "__main__": + pytest.main([__file__, "-svv"]) diff --git a/verl/tests/special_distributed/test_mcore_config_converter.py b/verl/tests/special_distributed/test_mcore_config_converter.py new file mode 100644 index 0000000000000000000000000000000000000000..d8f24c49911ed7b1fb1d73740dfc150e57dade0d --- /dev/null +++ b/verl/tests/special_distributed/test_mcore_config_converter.py @@ -0,0 +1,100 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os + +import megatron.core.parallel_state as mpu +import torch +from megatron.core.transformer import MLATransformerConfig, TransformerConfig +from transformers import AutoConfig, PretrainedConfig + +from verl.models.mcore import hf_to_mcore_config +from verl.utils.distributed import destroy_global_process_group, initialize_global_process_group + +TEST_MODELS = [ + "Qwen/Qwen2.5-7B", # Qwen2 dense + "Qwen/Qwen3-8B", # Qwen3 dense + "deepseek-ai/deepseek-coder-1.3b-instruct", # deepseek dense + "Qwen/Qwen2-57B-A14B", # Qwen2 moe + "Qwen/Qwen3-30B-A3B", # Qwen3 moe + # "mistralai/Mixtral-8x7B-v0.1", # Mixtral # require authentication + "deepseek-ai/DeepSeek-V3-Base", # Deepseek V3 +] + + +def check_config_converter_results(tf_config: TransformerConfig | MLATransformerConfig, hf_config: PretrainedConfig): + assert tf_config.num_layers == hf_config.num_hidden_layers, ( + f"Number of layers mismatch: {tf_config.num_layers} != {hf_config.num_hidden_layers}" + ) + assert tf_config.hidden_size == hf_config.hidden_size, ( + f"Hidden size mismatch: {tf_config.hidden_size} != {hf_config.hidden_size}" + ) + assert tf_config.num_attention_heads == hf_config.num_attention_heads, ( + f"Number of attention heads mismatch: {tf_config.num_attention_heads} != {hf_config.num_attention_heads}" + ) + assert tf_config.num_query_groups == hf_config.num_key_value_heads, ( + f"Number of query groups mismatch: {tf_config.num_query_groups} != {hf_config.num_key_value_heads}" + ) + assert tf_config.ffn_hidden_size == hf_config.intermediate_size, ( + f"FFN hidden size mismatch: {tf_config.ffn_hidden_size} != {hf_config.intermediate_size}" + ) + assert tf_config.attention_dropout == hf_config.attention_dropout, ( + f"Attention dropout mismatch: {tf_config.attention_dropout} != {hf_config.attention_dropout}" + ) + assert tf_config.hidden_dropout == getattr(hf_config, "hidden_dropout", 0.0), ( + f"Hidden dropout mismatch: {tf_config.hidden_dropout} != {getattr(hf_config, 'hidden_dropout', 0.0)}" + ) + if getattr(hf_config, "head_dim", None) is not None: + assert tf_config.kv_channels == getattr(hf_config, "head_dim", None), ( + f"Head dim mismatch: {tf_config.kv_channels} != {getattr(hf_config, 'head_dim', None)}" + ) + assert tf_config.layernorm_epsilon == hf_config.rms_norm_eps, ( + f"Layernorm epsilon mismatch: {tf_config.layernorm_epsilon} != {hf_config.rms_norm_eps}" + ) + + +def modify_hf_config(name: str, hf_config: PretrainedConfig): + if name == "deepseek-ai/DeepSeek-V3-Base": + hf_config.num_nextn_predict_layers = 0 + hf_config.quantization_config = None + return hf_config + + +def test_mcore_config_converter(): + """ + Test the conversion of Hugging Face model configurations to MCore configurations. + """ + local_rank, rank, world_size = initialize_global_process_group() + mpu.initialize_model_parallel( + tensor_model_parallel_size=2, + pipeline_model_parallel_size=2, + virtual_pipeline_model_parallel_size=None, + use_sharp=False, + context_parallel_size=2, + expert_model_parallel_size=1, + expert_tensor_parallel_size=None, + nccl_communicator_config_path=None, + ) + for model_name in TEST_MODELS: + print(f"testing {model_name}") + hf_config = AutoConfig.from_pretrained(os.path.expanduser(f"~/models/configs/{model_name}/config.json")) + hf_config = modify_hf_config(model_name, hf_config) + tf_config = hf_to_mcore_config(hf_config, torch.bfloat16) + check_config_converter_results(tf_config, hf_config) + + destroy_global_process_group() + + +if __name__ == "__main__": + test_mcore_config_converter() diff --git a/verl/tests/special_distributed/test_tensor_dict.py b/verl/tests/special_distributed/test_tensor_dict.py new file mode 100644 index 0000000000000000000000000000000000000000..8f380c7d542fc88cf8439dcad62ee88de798a7dc --- /dev/null +++ b/verl/tests/special_distributed/test_tensor_dict.py @@ -0,0 +1,171 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os + +os.environ["NCCL_DEBUG"] = "WARN" + +import numpy as np +import torch +import torch.distributed + +from verl.protocol import DataProto, all_gather_data_proto +from verl.utils.device import get_device_name +from verl.utils.distributed import initialize_global_process_group + + +def test_all_gather_data_proto(): + device_mesh = torch.distributed.device_mesh.init_device_mesh( + get_device_name(), mesh_shape=[2, 2], mesh_dim_names=["dp", "tp"] + ) + + global_rank = torch.distributed.get_rank() + + obs = torch.tensor([[1 * global_rank, 2 * global_rank + 1], [3 * global_rank, 4 * global_rank + 1]]) + + labels = ["a", "b"] if global_rank % 2 == 0 else ["b", "a"] + labels = np.array(labels, dtype=object) + data = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"labels": labels}, meta_info={"info": "test_info"}) + + all_gather_data_proto(data=data, process_group=device_mesh.get_group("dp")) + + if global_rank == 0: + expected_obs = torch.tensor([[0, 1], [0, 1], [2, 5], [6, 9]], device=get_device_name()) + expected_labels = ["a", "b", "a", "b"] + elif global_rank == 1: + expected_obs = torch.tensor([[1, 3], [3, 5], [3, 7], [9, 13]], device=get_device_name()) + expected_labels = ["b", "a", "b", "a"] + elif global_rank == 2: + expected_obs = torch.tensor([[0, 1], [0, 1], [2, 5], [6, 9]], device=get_device_name()) + expected_labels = ["a", "b", "a", "b"] + elif global_rank == 3: + expected_obs = torch.tensor([[1, 3], [3, 5], [3, 7], [9, 13]], device=get_device_name()) + expected_labels = ["b", "a", "b", "a"] + + torch.testing.assert_close(data.batch["obs"], expected_obs, atol=0, rtol=0) + assert (data.non_tensor_batch["labels"] == expected_labels).all() + assert data.meta_info == {"info": "test_info"} + + +def test_vocab_parallel_entropy(): + from megatron.core import parallel_state as mpu + + from verl.utils.megatron.tensor_parallel import vocab_parallel_entropy + from verl.utils.profiler import log_gpu_memory_usage + from verl.utils.torch_functional import entropy_from_logits + + mpu.initialize_model_parallel( + tensor_model_parallel_size=2, pipeline_model_parallel_size=1, virtual_pipeline_model_parallel_size=None + ) + + batch_size = 2 + seqlen = 128 + vocab_size = 155136 + + logits = torch.randn(batch_size * seqlen, vocab_size, device=get_device_name(), requires_grad=True) + target = torch.randint( + low=0, high=vocab_size, size=(batch_size * seqlen,), device=get_device_name(), dtype=torch.int64 + ) + + # broadcast across tp + torch.distributed.broadcast( + logits, mpu.get_tensor_model_parallel_src_rank(), group=mpu.get_tensor_model_parallel_group() + ) + torch.distributed.broadcast( + target, mpu.get_tensor_model_parallel_src_rank(), group=mpu.get_tensor_model_parallel_group() + ) + + tp_rank = mpu.get_tensor_model_parallel_rank() + vocab_size_per_tp = vocab_size // mpu.get_tensor_model_parallel_world_size() + + # get the local logits of each tp + vocab_parallel_logits = ( + logits.clone().detach()[:, tp_rank * vocab_size_per_tp : (tp_rank + 1) * vocab_size_per_tp].requires_grad_() + ) + logits.grad = None + vocab_parallel_logits.grad = None + + log_gpu_memory_usage("begin") + output_entropy = vocab_parallel_entropy(vocab_parallel_logits) + log_gpu_memory_usage("after forward") + grad_output = torch.randn_like(output_entropy) + output_entropy.backward(grad_output) + log_gpu_memory_usage("after backward") + + target_entropy = entropy_from_logits(logits) + torch.testing.assert_close(output_entropy, target_entropy) + target_entropy.backward(grad_output) + torch.testing.assert_close( + logits.grad[:, tp_rank * vocab_size_per_tp : (tp_rank + 1) * vocab_size_per_tp], vocab_parallel_logits.grad + ) + # make sure logits is not altered + torch.testing.assert_close( + logits[:, tp_rank * vocab_size_per_tp : (tp_rank + 1) * vocab_size_per_tp], vocab_parallel_logits + ) + + if mpu.get_tensor_model_parallel_rank() == 0: + print("test_vocab_parallel_entropy passes") + + mpu.destroy_model_parallel() + + +def test_vocab_parallel_sum_pi_squared(): + from megatron.core import parallel_state as mpu + + from verl.utils.megatron.tensor_parallel import vocab_parallel_sum_pi_squared + from verl.utils.torch_functional import calculate_sum_pi_squared_from_logits + + if not mpu.model_parallel_is_initialized(): + mpu.initialize_model_parallel( + tensor_model_parallel_size=2, pipeline_model_parallel_size=1, virtual_pipeline_model_parallel_size=None + ) + + batch_size = 2 + seqlen = 64 + vocab_size = 4096 + + logits = torch.randn(batch_size * seqlen, vocab_size, device=get_device_name()) + + # broadcast across tp so every rank shares the same full logits to shard from + torch.distributed.broadcast( + logits, mpu.get_tensor_model_parallel_src_rank(), group=mpu.get_tensor_model_parallel_group() + ) + + tp_rank = mpu.get_tensor_model_parallel_rank() + vocab_size_per_tp = vocab_size // mpu.get_tensor_model_parallel_world_size() + + vocab_parallel_logits = logits[:, tp_rank * vocab_size_per_tp : (tp_rank + 1) * vocab_size_per_tp].clone() + pre_call = vocab_parallel_logits.clone() + + output = vocab_parallel_sum_pi_squared(vocab_parallel_logits) + target = calculate_sum_pi_squared_from_logits(logits) + + torch.testing.assert_close(output, target, atol=1e-5, rtol=1e-5) + # non-destructive: input shard must not be mutated + torch.testing.assert_close(vocab_parallel_logits, pre_call) + # sanity: Σπ² ∈ (0, 1] + assert torch.all(output > 0) + assert torch.all(output <= 1.0 + 1e-5) + + if mpu.get_tensor_model_parallel_rank() == 0: + print("test_vocab_parallel_sum_pi_squared passes") + + mpu.destroy_model_parallel() + + +if __name__ == "__main__": + local_rank, rank, world_size = initialize_global_process_group() + test_all_gather_data_proto() + test_vocab_parallel_entropy() + test_vocab_parallel_sum_pi_squared() diff --git a/verl/tests/special_distributed/test_torch_functional.py b/verl/tests/special_distributed/test_torch_functional.py new file mode 100644 index 0000000000000000000000000000000000000000..d07d335f5a313e6557e72e2331c88176486fc016 --- /dev/null +++ b/verl/tests/special_distributed/test_torch_functional.py @@ -0,0 +1,35 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os + +import torch + +from verl.utils.torch_functional import allgather_dict_into_dict + +if __name__ == "__main__": + torch.distributed.init_process_group(backend="gloo") + + local_rank = int(os.environ["LOCAL_RANK"]) + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + + metrics_dict = {"loss": [0 + rank, 1 + rank, 2 + rank], "grad_norm": rank} + + result = allgather_dict_into_dict(data=metrics_dict, group=None) + + assert result["loss"] == [[0, 1, 2], [1, 2, 3], [2, 3, 4], [3, 4, 5]] + assert result["grad_norm"] == [0, 1, 2, 3] + + print(result) diff --git a/verl/tests/special_e2e/README.md b/verl/tests/special_e2e/README.md new file mode 100644 index 0000000000000000000000000000000000000000..3c295e844ceb11ee132564ab2949a05a2a066b3e --- /dev/null +++ b/verl/tests/special_e2e/README.md @@ -0,0 +1 @@ +This folder is reserved for end-to-end tests that typically require multiple GPUs. diff --git a/verl/tests/special_e2e/__init__.py b/verl/tests/special_e2e/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1ce90c5eb352d85c59105c0dc85b5f1dd576f095 --- /dev/null +++ b/verl/tests/special_e2e/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/verl/tests/special_e2e/check_custom_rwd_fn.py b/verl/tests/special_e2e/check_custom_rwd_fn.py new file mode 100644 index 0000000000000000000000000000000000000000..8d77a53729bd96b153f004eb230df85f1d32f890 --- /dev/null +++ b/verl/tests/special_e2e/check_custom_rwd_fn.py @@ -0,0 +1,33 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse + + +def check_congratulations_in_file(output_file): + with open(output_file) as f: + output = f.read() + + success_message = "Congratulations!!! You have called my_reward_function successfully!!!" + assert success_message in output, f"Success message of my_reward_function not found in {output_file}" + print("Check passes") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--output_file", required=True, type=str) + + args = parser.parse_args() + + check_congratulations_in_file(args.output_file) diff --git a/verl/tests/special_e2e/check_results.py b/verl/tests/special_e2e/check_results.py new file mode 100644 index 0000000000000000000000000000000000000000..9453282fbc80c88a12429369647208347d35491b --- /dev/null +++ b/verl/tests/special_e2e/check_results.py @@ -0,0 +1,53 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse + +import numpy as np + + +def extract_reward_from_line(line): + # TODO: this function needs error handling + try: + key_vals = line.split(" - ") + for key_val in key_vals: + key, val = key_val.split(":") + if key == "critic/rewards/mean": + reward = float(val) + return reward + return -np.inf + except Exception: + return -np.inf + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--output_file", required=True, type=str) + parser.add_argument("--target", type=float, default=0.2, help="target reward score") + + args = parser.parse_args() + + with open(args.output_file) as f: + output = f.read().split("\n") + + best_reward = -np.inf + for line in output: + if line.startswith("step"): + reward = extract_reward_from_line(line) + if reward > best_reward: + best_reward = reward + + print(f"Best reward is {best_reward}") + assert best_reward > args.target, f"Best reward must be greater than {args.target}. best_reward: {best_reward}" + print("Check passes") diff --git a/verl/tests/special_e2e/envs/__init__.py b/verl/tests/special_e2e/envs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..eb85e22f361e4af4635bda991ff12a1ed4911eec --- /dev/null +++ b/verl/tests/special_e2e/envs/__init__.py @@ -0,0 +1,17 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .digit_completion import DigitCompletion + +__all__ = ["DigitCompletion"] diff --git a/verl/tests/special_e2e/envs/digit_completion/__init__.py b/verl/tests/special_e2e/envs/digit_completion/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..80893ae41d6669f4f7265ce76d7ac28579b30b6f --- /dev/null +++ b/verl/tests/special_e2e/envs/digit_completion/__init__.py @@ -0,0 +1,22 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from transformers import AutoTokenizer, LlamaConfig + +from .task import DigitCompletion, generate_ground_truth_response +from .tokenizer import CharTokenizer + +AutoTokenizer.register(LlamaConfig, CharTokenizer, exist_ok=True) + +__all__ = ["DigitCompletion", "generate_ground_truth_response", "CharTokenizer"] diff --git a/verl/tests/special_e2e/envs/digit_completion/task.py b/verl/tests/special_e2e/envs/digit_completion/task.py new file mode 100644 index 0000000000000000000000000000000000000000..c3643a86b867b440352ed55dc0f978135ac79bcf --- /dev/null +++ b/verl/tests/special_e2e/envs/digit_completion/task.py @@ -0,0 +1,179 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Task and environment definition for digit completion.""" + +import numpy as np + + +class DigitCompletion: + """ + The implementation of a simple digit completion task. + The prompt is a sequence of numbers with fixed difference. The task is to complete the next N numbers. + If the max number is reached, the next number should be modulo with max number. + + For example, + - prompt = [1, 2, 3] + - N = 5 + - max_number = 6 + + the response should be [4, 5, 6, 7%6, 8%6] = [4, 5, 6, 0, 1] + + Note that the tokenizer is char-level to increase the difficulty. + """ + + def __init__(self, max_number: int, max_diff: int, max_num_in_response: int, seed=0): + """ + + Args: + max_number: the maximum number allowed in the arithmetic sequence + max_diff: the maximum diff. The actual common diff will be sampled from [0, max_diff] + max_num_in_response: the maximum number in the response + """ + super().__init__() + self.max_number = max_number + self.max_diff = max_diff + self.max_num_in_response = max_num_in_response + assert self.max_num_in_response < 10 + assert self.max_number > 0 + assert self.max_diff > 0 + self.max_number_length = len(str(max_number)) + # {num1},{num2}:{max_num_in_response},{max_number} + self._prompt_length = self.max_number_length * 2 + 4 + self.max_number_length # no negative is allowed + + self.np_rng = np.random.default_rng(seed=seed) + + def __str__(self): + return ( + f"Prompt length: {self.prompt_length}. Response length: {self.response_length}, " + f"Max number: {self.max_number}. Max diff: {self.max_diff}, " + f"Max number in response: {self.max_num_in_response}" + ) + + def get_state(self): + return {"rng": self.np_rng} + + def set_state(self, state): + assert "rng" in state, "rng must be inside state" + self.np_rng = state["rng"] + + @property + def prompt_length(self): + return self._prompt_length + + @property + def response_length(self): + # number length + comma length + [EOS] + # The actual number times 1.5 to allow 'U' + return (self.max_num_in_response * self.max_number_length + (self.max_num_in_response - 1) + 1) * 2 + + def add(self, a, b): + return (a + b) % self.max_number + + def get_all_prompts(self): + all_prompts = [] + for first_num in range(self.max_number + 1): + for diff in range(0, self.max_diff + 1): + second_num = self.add(first_num, diff) + for num_to_complete in range(self.max_num_in_response + 1): + prompt = str(first_num) + "," + str(second_num) + f":{self.max_number},{num_to_complete}" + all_prompts.append(prompt) + return all_prompts + + def sample_str_prompts(self): + # step 1: sample initial numbers + first_num = self.np_rng.integers(self.max_number + 1) + diff = self.np_rng.integers(self.max_diff + 1) + second_num = self.add(first_num, diff) + num_to_complete = self.np_rng.integers(self.max_num_in_response + 1) + prompt = str(first_num) + "," + str(second_num) + f":{self.max_number},{num_to_complete}" + return prompt + + def sample_batch_str_prompts(self, batch_size): + str_prompts = [] + for _ in range(batch_size): + str_prompts.append(self.sample_str_prompts()) + return str_prompts + + +def compute_attention_mask(prompts, pad_token_id): + mask = np.ones_like(prompts) + mask[prompts == pad_token_id] = 0 + return mask + + +def compute_position_id_with_mask(mask): + return np.clip(np.cumsum(mask, axis=-1) - 1, a_min=0, a_max=None) + + +def generate_ground_truth_response(prompt: str): + """Generate ground truth response given a prompt.""" + num, info = prompt.split(":") + num1, num2 = num.split(",") + max_number, num_to_gen = info.split(",") + num1 = int(num1) + num2 = int(num2) + max_number = int(max_number) + num_to_gen = int(num_to_gen) + diff = (num2 - num1) % max_number + results = [] + last_num = num2 + for _ in range(num_to_gen): + curr = (last_num + diff) % max_number + results.append(str(curr)) + last_num = curr + response = ",".join(results) + return response + + +def compute_reward(prompt: str, response: str, sequence_reward=1.0): + """We compute dense reward here so that we can directly train RL without SFT""" + response_length = len(response) + ground_truth_response = generate_ground_truth_response(prompt) + per_token_reward = sequence_reward / (len(ground_truth_response) + 1) # including [EOS] + + # pad + reward = np.zeros(response_length, dtype=np.float32) # this assumes that each char is a token + # assign reward until mismatches + ground_truth_idx = 0 + for i in range(response_length): + if ground_truth_idx == len(ground_truth_response): + break + + ground_truth_response_token = ground_truth_response[ground_truth_idx] + response_token = response[i] + if ground_truth_response_token == response_token: + reward[i] = per_token_reward + ground_truth_idx += 1 + else: + # no matches + break + + return reward, {"ground_truth_response": ground_truth_response} + + +if __name__ == "__main__": + task = DigitCompletion(max_number=20, max_diff=3, max_num_in_response=5) + print(task.sample_str_prompts()) + + prompt = "7,8:20,0" + response = "" + print(compute_reward(prompt, response)) + + prompt = "7,8:20,0" + response = "E000" + print(compute_reward(prompt, response)) + + prompt = "9,10:20,2" + response = "11,12,13" + print(compute_reward(prompt, response)) diff --git a/verl/tests/special_e2e/envs/digit_completion/tokenizer.py b/verl/tests/special_e2e/envs/digit_completion/tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..6ff471938937dc55ab528cb883e4ba2e03b35416 --- /dev/null +++ b/verl/tests/special_e2e/envs/digit_completion/tokenizer.py @@ -0,0 +1,155 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Copied from https://github.com/dariush-bahrami/character-tokenizer/blob/master/charactertokenizer/core.py + +CharacterTokenzier for Hugging Face Transformers. + +This is heavily inspired from CanineTokenizer in transformers package. +""" + +import json +import os +from pathlib import Path +from typing import Optional, Sequence + +from transformers.tokenization_utils import AddedToken, PreTrainedTokenizer + + +class CharTokenizer(PreTrainedTokenizer): + def __init__(self, characters: Sequence[str], model_max_length: int, chat_template, **kwargs): + """Character tokenizer for Hugging Face transformers. + + Args: + characters (Sequence[str]): List of desired characters. Any character which + is not included in this list will be replaced by a special token called + [UNK] with id=6. Following are list of all of the special tokens with + their corresponding ids: + "[CLS]": 0 + "[SEP]": 1 + "[BOS]": 2 + "[MASK]": 3 + "[PAD]": 4 + "[RESERVED]": 5 + "[UNK]": 6 + an id (starting at 7) will be assigned to each character. + + model_max_length (int): Model maximum sequence length. + """ + eos_token_str = "E" + sep_token_str = "S" + pad_token_str = "P" + unk_token_str = "U" + + self.characters = characters + self.model_max_length = model_max_length + eos_token = AddedToken(eos_token_str, lstrip=False, rstrip=False) + sep_token = AddedToken(sep_token_str, lstrip=False, rstrip=False) + pad_token = AddedToken(pad_token_str, lstrip=False, rstrip=False) + unk_token = AddedToken(unk_token_str, lstrip=False, rstrip=False) + + self._vocab_str_to_int = { + sep_token_str: 0, + eos_token_str: 1, + pad_token_str: 2, + unk_token_str: 3, + **{ch: i + 4 for i, ch in enumerate(characters)}, + } + self._vocab_int_to_str = {v: k for k, v in self._vocab_str_to_int.items()} + + super().__init__( + eos_token=eos_token, + sep_token=sep_token, + pad_token=pad_token, + unk_token=unk_token, + add_prefix_space=False, + model_max_length=model_max_length, + **kwargs, + ) + + self.chat_template = chat_template + + @property + def vocab_size(self) -> int: + return len(self._vocab_str_to_int) + + def get_vocab(self): + return self._vocab_str_to_int + + def _tokenize(self, text: str) -> list[str]: + return list(text) + + def _convert_token_to_id(self, token: str) -> int: + return self._vocab_str_to_int.get(token, self._vocab_str_to_int["U"]) + + def _convert_id_to_token(self, index: int) -> str: + return self._vocab_int_to_str[index] + + def convert_tokens_to_string(self, tokens): + return "".join(tokens) + + def build_inputs_with_special_tokens( + self, token_ids_0: list[int], token_ids_1: Optional[list[int]] = None + ) -> list[int]: + sep = [self.sep_token_id] + cls = [self.cls_token_id] + result = cls + token_ids_0 + sep + if token_ids_1 is not None: + result += token_ids_1 + sep + return result + + def get_special_tokens_mask( + self, + token_ids_0: list[int], + token_ids_1: Optional[list[int]] = None, + already_has_special_tokens: bool = False, + ) -> list[int]: + if already_has_special_tokens: + return super().get_special_tokens_mask( + token_ids_0=token_ids_0, + token_ids_1=token_ids_1, + already_has_special_tokens=True, + ) + + result = [1] + ([0] * len(token_ids_0)) + [1] + if token_ids_1 is not None: + result += ([0] * len(token_ids_1)) + [1] + return result + + def get_config(self) -> dict: + return { + "char_ords": [ord(ch) for ch in self.characters], + "model_max_length": self.model_max_length, + "chat_template": self.chat_template, + } + + @classmethod + def from_config(cls, config: dict): + cfg = {} + cfg["characters"] = [chr(i) for i in config["char_ords"]] + cfg["model_max_length"] = config["model_max_length"] + cfg["chat_template"] = config["chat_template"] + return cls(**cfg) + + def save_pretrained(self, save_directory: str | os.PathLike, **kwargs): + cfg_file = Path(save_directory) / "tokenizer_config.json" + cfg = self.get_config() + with open(cfg_file, "w") as f: + json.dump(cfg, f, indent=4) + + @classmethod + def from_pretrained(cls, save_directory: str | os.PathLike, **kwargs): + cfg_file = Path(save_directory) / "tokenizer_config.json" + with open(cfg_file) as f: + cfg = json.load(f) + return cls.from_config(cfg) diff --git a/verl/tests/special_e2e/ppo_trainer/expert_parallel/qwen2moe_minimal.json b/verl/tests/special_e2e/ppo_trainer/expert_parallel/qwen2moe_minimal.json new file mode 100644 index 0000000000000000000000000000000000000000..c215fa4f7ccb777035e4be513045fb6ddb204b8f --- /dev/null +++ b/verl/tests/special_e2e/ppo_trainer/expert_parallel/qwen2moe_minimal.json @@ -0,0 +1,4 @@ +{ + "num_hidden_layers": 2, + "max_window_layers": 2 +} \ No newline at end of file diff --git a/verl/tests/special_e2e/ppo_trainer/expert_parallel/qwen3moe_minimal.json b/verl/tests/special_e2e/ppo_trainer/expert_parallel/qwen3moe_minimal.json new file mode 100644 index 0000000000000000000000000000000000000000..c215fa4f7ccb777035e4be513045fb6ddb204b8f --- /dev/null +++ b/verl/tests/special_e2e/ppo_trainer/expert_parallel/qwen3moe_minimal.json @@ -0,0 +1,4 @@ +{ + "num_hidden_layers": 2, + "max_window_layers": 2 +} \ No newline at end of file diff --git a/verl/tests/special_e2e/ppo_trainer/run_function_reward.sh b/verl/tests/special_e2e/ppo_trainer/run_function_reward.sh new file mode 100644 index 0000000000000000000000000000000000000000..50cbb539aefa771a1eaced5c43dbf80c567d0066 --- /dev/null +++ b/verl/tests/special_e2e/ppo_trainer/run_function_reward.sh @@ -0,0 +1,166 @@ +#!/usr/bin/env bash +set -xeuo pipefail + +NUM_GPUS=${NUM_GPUS:-8} + +MODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-0.5B} +MODEL_PATH=${MODEL_PATH:-${HOME}/models/${MODEL_ID}} +#hf download "${MODEL_ID}" --local-dir "${MODEL_PATH}" + +TRAIN_FILES=${TRAIN_FILES:-$HOME/data/gsm8k/train.parquet} +VAL_FILES=${VAL_FILES:-$HOME/data/gsm8k/test.parquet} +MAX_PROMPT_LEN=${MAX_PROMPT_LEN:-512} +MAX_RESPONSE_LEN=${MAX_RESPONSE_LEN:-512} + +ENGINE=${ENGINE:-vllm} +if [ "$ENGINE" = "vllm" ]; then + export VLLM_USE_V1=1 +fi +ROLLOUT_MODE="async" + +RETURN_RAW_CHAT="True" +SKIP_TOKENIZER_INIT="True" + +GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.7} +ACTOR_FSDP_PARAM_OFFLOAD=${ACTOR_FSDP_PARAM_OFFLOAD:-True} +ACTOR_FSDP_OPTIMIZER_OFFLOAD=${ACTOR_FSDP_OPTIMIZER_OFFLOAD:-True} +REF_FSDP_PARAM_OFFLOAD=${REF_FSDP_PARAM_OFFLOAD:-True} +RM_PAD=${RM_PAD:-True} +FUSED_KERNELS=${FUSED_KERNELS:-False} +FUSED_KERNEL_BACKEND=${FUSED_KERNEL_BACKEND:-torch} # or 'triton' for triton backend +ADV_ESTIMATOR=${ADV_ESTIMATOR:-gae} +LOSS_MODE=${LOSS_MODE:-vanilla} +USE_KL=${USE_KL:-False} +CUSTOM_REWARD_FN=${CUSTOM_REWARD_FN:-False} +ENABLE_CHUNKED_PREFILL=${ENABLE_CHUNKED_PREFILL:-True} # For vLLM VLM placeholder issue: https://github.com/vllm-project/vllm/issues/15185 +STRATEGY=${STRATEGY:-fsdp} +# LoRA config +LORA_RANK=${LORA_RANK:-0} +LORA_ALPHA=${LORA_ALPHA:-${LORA_RANK}} +LORA_TARGET=${LORA_TARGET:-"all-linear"} +LORA_EXCLUDE=${LORA_EXCLUDE:-"DONT_EXCLUDE"} +USE_SHM=${USE_SHM:-False} +LOAD_FORMAT=${LOAD_FORMAT:-dummy} +LAYERED_SUMMON=${LAYERED_SUMMON:-False} +# Validation +VAL_BEFORE_TRAIN=${VAL_BEFORE_TRAIN:-False} +TEST_FREQ=${TEST_FREQ:--1} +# Save & Resume +RESUME_MODE=${RESUME_MODE:-disable} +SAVE_FREQ=${SAVE_FREQ:--1} +TOTAL_TRAIN_STEPS=${TOTAL_TRAIN_STEPS:-1} + +# whether to save hf_model +SAVE_HF_MODEL=${SAVE_HF_MODEL:-False} +FSDP_SIZE=${FSDP_SIZE:--1} +SP_SIZE=${SP_SIZE:-1} + +if [ "${SAVE_HF_MODEL}" = "True" ]; then + CHECKPOINT_CONTENTS="['model','hf_model','optimizer','extra']" +else + CHECKPOINT_CONTENTS="['model','optimizer','extra']" +fi + +train_traj_micro_bsz_per_gpu=2 # b +n_resp_per_prompt=4 # g + +train_traj_micro_bsz=$((train_traj_micro_bsz_per_gpu * 1)) # b * n +train_traj_mini_bsz=$((train_traj_micro_bsz * 2)) # 2 * b * n +train_prompt_mini_bsz=$((train_traj_mini_bsz * 2)) # 2 * b * n / g +train_prompt_bsz=$((train_prompt_mini_bsz * 2)) # 4 * b * n / g + +reward_fn_name=null +reward_fn_file_path=null +output_file="$(pwd)/output.txt" +if [ "${CUSTOM_REWARD_FN}" = "True" ]; then + reward_fn_name="my_reward_function" + reward_fn_file_path="$(pwd)/my_reward_function.py" + rm -rf "${reward_fn_file_path}" + cat < "$reward_fn_file_path" +def ${reward_fn_name}(data_source, solution_str, ground_truth, extra_info=None): + print(f"Congratulations!!! You have called ${reward_fn_name} successfully!!!") + return 0.1 +EOF + + rm -rf "${output_file}" +fi + +exp_name="${VERL_EXP_NAME:-$(basename "${MODEL_ID,,}")-function-reward-minimal}" + +python3 -m verl.trainer.main_ppo \ + algorithm.adv_estimator="${ADV_ESTIMATOR}" \ + data.train_files="${TRAIN_FILES}" \ + data.val_files="${VAL_FILES}" \ + data.train_batch_size="${train_prompt_bsz}" \ + data.max_prompt_length="${MAX_PROMPT_LEN}" \ + data.max_response_length="${MAX_RESPONSE_LEN}" \ + data.return_raw_chat=${RETURN_RAW_CHAT} \ + actor_rollout_ref.model.path="${MODEL_PATH}" \ + actor_rollout_ref.model.use_shm=${USE_SHM} \ + actor_rollout_ref.model.lora_rank=${LORA_RANK} \ + actor_rollout_ref.model.lora_alpha=${LORA_ALPHA} \ + actor_rollout_ref.model.target_modules=${LORA_TARGET} \ + actor_rollout_ref.model.exclude_modules=${LORA_EXCLUDE} \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.model.use_remove_padding="${RM_PAD}" \ + actor_rollout_ref.model.use_fused_kernels=${FUSED_KERNELS} \ + actor_rollout_ref.model.fused_kernel_options.impl_backend=${FUSED_KERNEL_BACKEND} \ + actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${train_traj_micro_bsz_per_gpu} \ + actor_rollout_ref.actor.strategy=${STRATEGY} \ + actor_rollout_ref.actor.fsdp_config.param_offload=${ACTOR_FSDP_PARAM_OFFLOAD} \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=${ACTOR_FSDP_OPTIMIZER_OFFLOAD} \ + actor_rollout_ref.actor.fsdp_config.fsdp_size=${FSDP_SIZE} \ + actor_rollout_ref.actor.ulysses_sequence_parallel_size="${SP_SIZE}" \ + actor_rollout_ref.actor.checkpoint.save_contents=${CHECKPOINT_CONTENTS} \ + actor_rollout_ref.actor.use_kl_loss="${USE_KL}" \ + actor_rollout_ref.actor.policy_loss.loss_mode="${LOSS_MODE}" \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${train_traj_micro_bsz_per_gpu} \ + actor_rollout_ref.rollout.tensor_model_parallel_size=2 \ + actor_rollout_ref.rollout.n=${n_resp_per_prompt} \ + actor_rollout_ref.rollout.name="${ENGINE}" \ + actor_rollout_ref.rollout.mode="${ROLLOUT_MODE}" \ + actor_rollout_ref.rollout.load_format=${LOAD_FORMAT} \ + actor_rollout_ref.rollout.layered_summon=${LAYERED_SUMMON} \ + actor_rollout_ref.rollout.skip_tokenizer_init="${SKIP_TOKENIZER_INIT}" \ + actor_rollout_ref.rollout.gpu_memory_utilization="${GPU_MEMORY_UTILIZATION}" \ + actor_rollout_ref.rollout.enable_chunked_prefill="${ENABLE_CHUNKED_PREFILL}" \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${train_traj_micro_bsz_per_gpu} \ + actor_rollout_ref.ref.fsdp_config.param_offload="${REF_FSDP_PARAM_OFFLOAD}" \ + critic.optim.lr=1e-5 \ + critic.model.use_remove_padding="${RM_PAD}" \ + critic.model.path="${MODEL_PATH}" \ + critic.model.enable_gradient_checkpointing=False \ + critic.ppo_micro_batch_size_per_gpu=${train_traj_micro_bsz_per_gpu} \ + critic.fsdp.param_offload=True \ + critic.fsdp.optimizer_offload=True \ + reward.custom_reward_function.path="${reward_fn_file_path}"\ + reward.custom_reward_function.name="${reward_fn_name}"\ + algorithm.use_kl_in_reward="${USE_KL}" \ + algorithm.kl_penalty=kl \ + algorithm.kl_ctrl.kl_coef=0.001 \ + trainer.critic_warmup=0 \ + trainer.logger=console \ + trainer.project_name='verl-test' \ + trainer.experiment_name="${exp_name}" \ + trainer.nnodes=1 \ + trainer.n_gpus_per_node="${NUM_GPUS}" \ + trainer.val_before_train="${VAL_BEFORE_TRAIN}" \ + trainer.test_freq="${TEST_FREQ}" \ + trainer.save_freq="${SAVE_FREQ}" \ + trainer.resume_mode="${RESUME_MODE}" \ + trainer.total_epochs=2 \ + trainer.device=cuda \ + trainer.total_training_steps="${TOTAL_TRAIN_STEPS}" $@ \ + | tee "${output_file}" + +if [ "${CUSTOM_REWARD_FN}" = "True" ]; then + python3 tests/special_e2e/check_custom_rwd_fn.py --output_file="${output_file}" + check_exit_code=$? + rm -rf "${reward_fn_file_path}" + rm -rf "${output_file}" + # Return the exit code of check_custom_rwd_fn.py if it fails + if [ $check_exit_code -ne 0 ]; then + exit $check_exit_code + fi +fi diff --git a/verl/tests/special_e2e/run_fully_async_policy.sh b/verl/tests/special_e2e/run_fully_async_policy.sh new file mode 100644 index 0000000000000000000000000000000000000000..1e05341a34c5cf4c3d16582b8512b9b71fae84fd --- /dev/null +++ b/verl/tests/special_e2e/run_fully_async_policy.sh @@ -0,0 +1,233 @@ +#!/usr/bin/env bash +set -xeuo pipefail + +# Test script for fully_async_policy E2E regression testing +# This script runs fully async PPO training with both FSDP2 and Megatron backends +# to ensure the asynchronous training mechanism works correctly + +NUM_GPUS=${NUM_GPUS:-8} +ACTOR_STRATEGY=${ACTOR_STRATEGY:-"fsdp2"} # fsdp2 or megatron +ROLLOUT_NAME=${ROLLOUT_NAME:-"vllm"} # vllm, sglang, or trtllm + +# Download model if not exists +MODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-0.5B-Instruct} +MODEL_PATH=${MODEL_PATH:-${HOME}/models/${MODEL_ID}} +# hf download "${MODEL_ID}" --local-dir "${MODEL_PATH}" + + +rollout_mode="async" +rollout_name="${ROLLOUT_NAME}" +return_raw_chat="True" +if [ "$rollout_name" = "vllm" ]; then + export VLLM_USE_V1=1 +fi + +# Algorithm parameters +adv_estimator=grpo + +use_kl_in_reward=False +kl_coef=0.0 +use_kl_loss=False +kl_loss_coef=0.0 + +clip_ratio_low=0.2 +clip_ratio_high=0.28 + +# Response length parameters +max_prompt_length=1024 +max_response_length=2048 +enable_overlong_buffer=True +overlong_buffer_len=128 +overlong_penalty_factor=1.0 + +# Training parameters +loss_agg_mode="token-mean" + +# Temperature parameters +temperature=1.0 +top_p=1.0 +top_k=-1 +val_top_p=0.7 + +# Fully async specific parameters +# Split GPUs evenly between rollout and training. +n_gpus_rollout=${N_GPUS_ROLLOUT:-$((NUM_GPUS / 2))} +n_gpus_training=${N_GPUS_TRAINING:-$((NUM_GPUS / 2))} + +train_prompt_bsz=0 +gen_prompt_bsz=1 +n_resp_per_prompt=16 +train_prompt_mini_bsz=16 +total_rollout_steps=$(((128))) +test_freq=-1 +staleness_threshold=0.5 +trigger_parameter_sync_step=4 +partial_rollout=True +use_trainer_do_validate=False + +exp_name="$(basename "${MODEL_ID,,}")-fully-async-policy-${rollout_name}-${ACTOR_STRATEGY}-minimal" + +echo "Running fully_async_policy with ${ACTOR_STRATEGY} strategy" +echo "Total GPUs: ${NUM_GPUS}, Rollout GPUs: ${n_gpus_rollout}, Training GPUs: ${n_gpus_training}" + +# Common parameters for both FSDP2 and Megatron +common_params=( + data.train_files="${HOME}/data/gsm8k/train.parquet" + data.val_files="${HOME}/data/gsm8k/test.parquet" + data.prompt_key=prompt + data.truncation='left' + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.train_batch_size=${train_prompt_bsz} + data.gen_batch_size=${gen_prompt_bsz} + data.return_raw_chat=${return_raw_chat} + actor_rollout_ref.rollout.n=${n_resp_per_prompt} + actor_rollout_ref.rollout.calculate_log_probs=True + algorithm.adv_estimator=${adv_estimator} + algorithm.use_kl_in_reward=${use_kl_in_reward} + algorithm.kl_ctrl.kl_coef=${kl_coef} + actor_rollout_ref.hybrid_engine=False + actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} + actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} + actor_rollout_ref.actor.clip_ratio_c=10.0 + actor_rollout_ref.model.path="${MODEL_PATH}" + actor_rollout_ref.actor.optim.lr=1e-6 + actor_rollout_ref.actor.optim.lr_warmup_steps=-1 + actor_rollout_ref.actor.optim.weight_decay=0.1 + actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} + actor_rollout_ref.actor.entropy_coeff=0 + actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} + actor_rollout_ref.rollout.gpu_memory_utilization=0.80 + actor_rollout_ref.rollout.temperature=${temperature} + actor_rollout_ref.rollout.top_p=${top_p} + actor_rollout_ref.rollout.top_k=${top_k} + actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} + actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} + actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} + actor_rollout_ref.rollout.val_kwargs.do_sample=True + actor_rollout_ref.rollout.val_kwargs.n=1 + actor_rollout_ref.rollout.enable_chunked_prefill=True + actor_rollout_ref.rollout.name=${rollout_name} + actor_rollout_ref.rollout.mode=${rollout_mode} + actor_rollout_ref.rollout.disable_log_stats=False + reward.reward_manager.name=dapo + +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} + +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} + +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} + +reward.reward_kwargs.overlong_buffer_cfg.log=False + +reward.reward_kwargs.max_resp_len=${max_response_length} + trainer.logger=['console'] + trainer.project_name='verl-test-fully-async' + trainer.experiment_name="${exp_name}" + trainer.val_before_train=True + trainer.save_freq=-1 + trainer.resume_mode=disable + trainer.nnodes=1 + trainer.n_gpus_per_node=${n_gpus_training} + trainer.log_val_generations=10 + rollout.nnodes=1 + rollout.n_gpus_per_node=${n_gpus_rollout} + rollout.total_rollout_steps=${total_rollout_steps} + trainer.total_epochs=2 + trainer.test_freq=${test_freq} + # Fully async specific configurations + async_training.staleness_threshold=${staleness_threshold} + async_training.partial_rollout="${partial_rollout}" + async_training.trigger_parameter_sync_step="${trigger_parameter_sync_step}" + async_training.use_trainer_do_validate=${use_trainer_do_validate} + actor_rollout_ref.rollout.checkpoint_engine.backend='nccl' + actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=1024 +) + + # Detect device + device_name=$(python3 - <<'EOF' +from verl.utils.device import get_device_name +print(get_device_name()) +EOF +) + +if [ "${ACTOR_STRATEGY}" == "fsdp2" ]; then + echo "Running fully async training with FSDP2 strategy..." + # FSDP2 specific parameters + # trtllm: one replica uses all rollout GPUs as a single TP group. + # vllm/sglang: TP=1, rely on data parallelism across replicas. + if [ "${rollout_name}" = "trtllm" ]; then + gen_tp=${GEN_TP:-${n_gpus_rollout}} + else + gen_tp=1 + fi + sp_size=1 + fsdp_size=1 + ref_offload=True + actor_offload=False + + if [ -n "$device_name" ] && [ "$device_name" == "npu" ]; then + common_params+=( + # Todo The checkpoint_engine.backend should be unified to nccl + # actor_rollout_ref.rollout.checkpoint_engine.backend='hccl' + actor_rollout_ref.rollout.gpu_memory_utilization=0.70 + ) + actor_offload=True + fi + python3 -m verl.experimental.fully_async_policy.fully_async_main \ + "${common_params[@]}" \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.fsdp_config.strategy=fsdp2 \ + critic.strategy=fsdp2 \ + actor_rollout_ref.actor.grad_clip=1.0 \ + actor_rollout_ref.model.use_remove_padding=True \ + actor_rollout_ref.actor.use_dynamic_bsz=True \ + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True \ + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \ + actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \ + actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \ + actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \ + actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \ + actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \ + actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} $@ + +elif [ "${ACTOR_STRATEGY}" == "megatron" ]; then + echo "Running fully async training with Megatron strategy..." + # Megatron specific parameters + if [ "${rollout_name}" = "trtllm" ]; then + gen_tp=${GEN_TP:-${n_gpus_rollout}} + else + gen_tp=2 + fi + train_tp=2 + train_pp=$((n_gpus_training / train_tp)) + ref_offload=True + actor_offload=True + common_params+=( + actor_rollout_ref.rollout.gpu_memory_utilization=0.60 + ) + + python3 -m verl.experimental.fully_async_policy.fully_async_main \ + --config-path=config \ + --config-name='fully_async_ppo_megatron_trainer.yaml' \ + "${common_params[@]}" \ + actor_rollout_ref.actor.strategy=megatron \ + critic.strategy=megatron \ + actor_rollout_ref.actor.optim.lr_decay_steps=10000000 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.actor.megatron.param_offload=${actor_offload} \ + actor_rollout_ref.actor.megatron.optimizer_offload=${actor_offload} \ + actor_rollout_ref.actor.megatron.grad_offload=${actor_offload} \ + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${train_pp} \ + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${train_tp} \ + actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \ + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${train_pp} \ + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${train_tp} \ + actor_rollout_ref.ref.megatron.param_offload=${ref_offload} $@ +else + echo "Error: Unknown strategy ${ACTOR_STRATEGY}. Please use 'fsdp2' or 'megatron'" + exit 1 +fi + +echo "Fully async policy E2E test completed successfully with ${ACTOR_STRATEGY} strategy" + diff --git a/verl/tests/special_e2e/run_fully_async_policy_genrm.sh b/verl/tests/special_e2e/run_fully_async_policy_genrm.sh new file mode 100644 index 0000000000000000000000000000000000000000..06a6e339ca6cf4184b8a0cc7dcb07e2daa31c9c4 --- /dev/null +++ b/verl/tests/special_e2e/run_fully_async_policy_genrm.sh @@ -0,0 +1,191 @@ +#!/usr/bin/env bash +set -xeuo pipefail + +# Test script for fully_async_policy + GenRM E2E regression testing +# This script runs fully async PPO training with a standalone GenRM reward model +# to verify that GenRM works correctly in async mode. +# +# GPU allocation (3 GPUs minimum): +# - 1 GPU: Rollout (vLLM async) +# - 1 GPU: Training (FSDP2, offload enabled) +# - 1 GPU: GenRM (vLLM standalone) + +NUM_GPUS=${NUM_GPUS:-3} + +# Model paths (use HF repo ID by default, auto-downloaded by transformers/vLLM) +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen2.5-0.5B-Instruct} +# Use the same small model as GenRM judge for testing (production would use a larger model) +GRM_PATH=${GRM_PATH:-Qwen/Qwen2.5-3B-Instruct} + +rollout_mode="async" +rollout_name="vllm" +export VLLM_USE_V1=1 +export GENRM_MODEL_NAME="${GRM_PATH}" +return_raw_chat="True" + +# Algorithm parameters +adv_estimator=grpo + +use_kl_in_reward=False +kl_coef=0.0 +use_kl_loss=False +kl_loss_coef=0.0 + +clip_ratio_low=0.2 +clip_ratio_high=0.28 + +# Response length parameters +max_prompt_length=256 +max_response_length=512 +max_num_tokens=$(( max_prompt_length + max_response_length + 1 )) +enable_overlong_buffer=True +overlong_buffer_len=128 +overlong_penalty_factor=1.0 + +# Training parameters +loss_agg_mode="token-mean" + +# Temperature parameters +temperature=1.0 +top_p=1.0 +top_k=-1 +val_top_p=0.7 + +# Fully async specific parameters +n_gpus_rollout=1 +n_gpus_training=1 +n_gpus_genrm=1 + +train_prompt_bsz=0 +gen_prompt_bsz=1 +n_resp_per_prompt=4 +train_prompt_mini_bsz=4 +total_rollout_steps=3200 # ~200 wandb data points on 3x H100 +test_freq=-1 +staleness_threshold=0.5 +trigger_parameter_sync_step=4 +partial_rollout=True +use_trainer_do_validate=False + +exp_name="$(basename "${MODEL_PATH,,}")-fully-async-policy-genrm-fsdp2-minimal" + +echo "Running fully_async_policy + GenRM with FSDP2 strategy" +echo "Total GPUs: ${NUM_GPUS}, Rollout GPUs: ${n_gpus_rollout}, Training GPUs: ${n_gpus_training}, GenRM GPUs: ${n_gpus_genrm}" + +# Detect device +device_name=$(python3 - <<'EOF' +from verl.utils.device import get_device_name +print(get_device_name()) +EOF +) + +gen_tp=1 +sp_size=1 +fsdp_size=1 +ref_offload=True +actor_offload=False + +if [ -n "$device_name" ] && [ "$device_name" == "npu" ]; then + actor_offload=True +fi + +python3 -m verl.experimental.fully_async_policy.fully_async_main \ + data.train_files="${HOME}/data/gsm8k/train.parquet" \ + data.val_files="${HOME}/data/gsm8k/test.parquet" \ + data.prompt_key=prompt \ + data.truncation='left' \ + data.max_prompt_length=${max_prompt_length} \ + data.max_response_length=${max_response_length} \ + data.train_batch_size=${train_prompt_bsz} \ + data.gen_batch_size=${gen_prompt_bsz} \ + data.return_raw_chat=${return_raw_chat} \ + actor_rollout_ref.rollout.n=${n_resp_per_prompt} \ + actor_rollout_ref.rollout.calculate_log_probs=True \ + algorithm.adv_estimator=${adv_estimator} \ + algorithm.use_kl_in_reward=${use_kl_in_reward} \ + algorithm.kl_ctrl.kl_coef=${kl_coef} \ + actor_rollout_ref.hybrid_engine=False \ + actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \ + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \ + actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \ + actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \ + actor_rollout_ref.actor.clip_ratio_c=10.0 \ + actor_rollout_ref.model.path="${MODEL_PATH}" \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.actor.optim.lr_warmup_steps=-1 \ + actor_rollout_ref.actor.optim.weight_decay=0.1 \ + actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \ + actor_rollout_ref.actor.entropy_coeff=0 \ + actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.50 \ + actor_rollout_ref.rollout.temperature=${temperature} \ + actor_rollout_ref.rollout.top_p=${top_p} \ + actor_rollout_ref.rollout.top_k=${top_k} \ + actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \ + actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \ + actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \ + actor_rollout_ref.rollout.val_kwargs.do_sample=True \ + actor_rollout_ref.rollout.val_kwargs.n=1 \ + actor_rollout_ref.rollout.enable_chunked_prefill=True \ + actor_rollout_ref.rollout.name=${rollout_name} \ + actor_rollout_ref.rollout.mode=${rollout_mode} \ + actor_rollout_ref.rollout.disable_log_stats=False \ + actor_rollout_ref.rollout.max_model_len=${max_num_tokens} \ + actor_rollout_ref.rollout.max_num_batched_tokens=${max_num_tokens} \ + actor_rollout_ref.rollout.max_num_seqs=${max_num_tokens} \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.fsdp_config.strategy=fsdp2 \ + critic.strategy=fsdp2 \ + actor_rollout_ref.actor.grad_clip=1.0 \ + actor_rollout_ref.model.use_remove_padding=True \ + actor_rollout_ref.actor.use_dynamic_bsz=True \ + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True \ + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \ + actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \ + actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \ + actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \ + actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \ + actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \ + actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \ + reward.reward_manager.name=dapo \ + +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \ + +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \ + +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \ + +reward.reward_kwargs.overlong_buffer_cfg.log=False \ + +reward.reward_kwargs.max_resp_len=${max_response_length} \ + reward.reward_model.enable=True \ + reward.reward_model.enable_resource_pool=True \ + reward.reward_model.n_gpus_per_node=${n_gpus_genrm} \ + reward.reward_model.nnodes=1 \ + reward.reward_model.model_path="${GRM_PATH}" \ + reward.reward_model.rollout.name=vllm \ + reward.reward_model.rollout.tensor_model_parallel_size=1 \ + reward.reward_model.rollout.gpu_memory_utilization=0.5 \ + reward.reward_model.rollout.skip_tokenizer_init=False \ + reward.custom_reward_function.path=tests/experimental/reward_loop/reward_fn.py \ + reward.custom_reward_function.name=compute_score_gsm8k \ + trainer.logger='["console","wandb"]' \ + trainer.project_name='verl-test-fully-async-genrm' \ + trainer.experiment_name="${exp_name}" \ + trainer.val_before_train=False \ + trainer.save_freq=-1 \ + trainer.resume_mode=disable \ + trainer.nnodes=1 \ + trainer.n_gpus_per_node=${n_gpus_training} \ + trainer.log_val_generations=10 \ + trainer.use_legacy_worker_impl=disable \ + rollout.nnodes=1 \ + rollout.n_gpus_per_node=${n_gpus_rollout} \ + rollout.total_rollout_steps=${total_rollout_steps} \ + trainer.total_epochs=2 \ + trainer.test_freq=${test_freq} \ + async_training.staleness_threshold=${staleness_threshold} \ + async_training.partial_rollout="${partial_rollout}" \ + async_training.trigger_parameter_sync_step="${trigger_parameter_sync_step}" \ + async_training.use_trainer_do_validate=${use_trainer_do_validate} \ + actor_rollout_ref.rollout.checkpoint_engine.backend='nccl' \ + actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=1024 \ + "$@" + +echo "Fully async policy + GenRM E2E test completed successfully" diff --git a/verl/tests/special_e2e/run_fully_async_policy_opd.sh b/verl/tests/special_e2e/run_fully_async_policy_opd.sh new file mode 100644 index 0000000000000000000000000000000000000000..eecfb90275c96b84bee4b3e7d5c0b32e788cdf23 --- /dev/null +++ b/verl/tests/special_e2e/run_fully_async_policy_opd.sh @@ -0,0 +1,279 @@ +#!/usr/bin/env bash +set -xeuo pipefail + +# Workaround for NVIDIA driver bug (r560-r575) causing SIGSEGV in ncclCuMemHostEnable() +# on PCIe machines without P2P access. See: https://github.com/NVIDIA/nccl/issues/1838 +export NCCL_CUMEM_ENABLE=0 +export NCCL_CUMEM_HOST_ENABLE=0 + +# Test script for fully_async_policy + Multi-Teacher Online Policy Distillation (OPD) +# This script runs fully async training with Megatron backend and multiple standalone +# teacher models to verify that multi-teacher distillation works correctly in async mode. +# +# Follows PR #6051 (Multi-Teacher OPD) setup: +# Student: Qwen3-VL-2B-Instruct (Megatron, tp=2, pp=2) +# GSM8K Teacher: Qwen3-4B-Instruct-2507 +# Geo3K Teacher: Qwen3-VL-4B-Instruct +# +# GPU allocation (8 GPUs): +# - 2 GPU: Rollout (student vLLM async, gen_tp=2) +# - 4 GPU: Training (student Megatron, tp=2 x pp=2) +# - 1 GPU: Teacher GSM8K (standalone vLLM) +# - 1 GPU: Teacher Geo3K (standalone vLLM) +# +# Usage: +# cd /root/verl && bash tests/special_e2e/run_fully_async_policy_opd.sh + +############################ Quick Config ############################ + +ROLLOUT_NAME="vllm" +export VLLM_USE_V1=1 + +STUDENT_MODEL_ID=${STUDENT_MODEL_ID:-Qwen/Qwen3-VL-2B-Instruct} +GSM8K_TEACHER_MODEL_ID=${GSM8K_TEACHER_MODEL_ID:-Qwen/Qwen3-4B-Instruct-2507} +GEO3K_TEACHER_MODEL_ID=${GEO3K_TEACHER_MODEL_ID:-Qwen/Qwen3-VL-4B-Instruct} + +STUDENT_MODEL=${STUDENT_MODEL:-${HOME}/models/${STUDENT_MODEL_ID}} +GSM8K_TEACHER_MODEL=${GSM8K_TEACHER_MODEL:-${HOME}/models/${GSM8K_TEACHER_MODEL_ID}} +GEO3K_TEACHER_MODEL=${GEO3K_TEACHER_MODEL:-${HOME}/models/${GEO3K_TEACHER_MODEL_ID}} + +DISTILLATION_LOSS_MODE="k1" +USE_POLICY_GRADIENT=True + +MAX_PROMPT=1024 +MAX_RESPONSE_LENGTH=2048 +MAX_NUM_TOKENS=$(( MAX_PROMPT + MAX_RESPONSE_LENGTH + 1 )) + +# Fully async specific +N_GPUS_ROLLOUT=2 +N_GPUS_TRAINING=4 +N_GPUS_TEACHER_TOTAL=2 # 1 per teacher +TOTAL_ROLLOUT_STEPS=${TOTAL_ROLLOUT_STEPS:-128} + +# Megatron parallelism +GEN_TP=2 +TRAIN_TP=2 +TRAIN_PP=2 + +STALENESS_THRESHOLD=0.5 +TRIGGER_PARAMETER_SYNC_STEP=4 + +############################ Data Preparation ############################ + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +VERL_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" + +GSM8K_DIR="${HOME}/data/gsm8k" +GEO3K_DIR="${HOME}/data/geo3k" + +# Prepare GSM8K (idempotent) +if [ ! -f "${GSM8K_DIR}/train.parquet" ]; then + echo "Preparing GSM8K dataset..." + python3 "${VERL_ROOT}/examples/data_preprocess/gsm8k.py" --local_save_dir "$GSM8K_DIR" +fi + +# Prepare Geo3K (idempotent) +if [ ! -f "${GEO3K_DIR}/train.parquet" ]; then + echo "Preparing Geo3K dataset..." + python3 "${VERL_ROOT}/examples/data_preprocess/geo3k.py" --local_save_dir "$GEO3K_DIR" +fi + +GSM8K_TRAIN="${GSM8K_DIR}/train.parquet" +GSM8K_TEST="${GSM8K_DIR}/test.parquet" +GEO3K_TRAIN="${GEO3K_DIR}/train.parquet" +GEO3K_TEST="${GEO3K_DIR}/test.parquet" + +TRAIN_FILES="['${GSM8K_TRAIN}','${GEO3K_TRAIN}']" +TEST_FILES="['${GSM8K_TEST}','${GEO3K_TEST}']" + +############################ Detect Device ############################ + +device_name=$(python3 - <<'EOF' +from verl.utils.device import get_device_name +print(get_device_name()) +EOF +) + +ACTOR_OFFLOAD=True + +############################ Parameter Groups ############################ + +DATA=( + data.train_files="$TRAIN_FILES" + data.val_files="$TEST_FILES" + data.prompt_key=prompt + data.truncation='left' + data.max_prompt_length=$MAX_PROMPT + data.max_response_length=$MAX_RESPONSE_LENGTH + data.train_batch_size=0 + data.gen_batch_size=1 + data.return_raw_chat=True + data.image_key=images +) + +MODEL=( + actor_rollout_ref.model.path="${STUDENT_MODEL}" + actor_rollout_ref.model.enable_gradient_checkpointing=True + actor_rollout_ref.model.use_remove_padding=True +) + +STUDENT=( + actor_rollout_ref.actor.strategy=megatron + actor_rollout_ref.actor.optim.lr=1e-6 + actor_rollout_ref.actor.optim.lr_warmup_steps=-1 + actor_rollout_ref.actor.optim.lr_decay_steps=10000000 + actor_rollout_ref.actor.optim.weight_decay=0.1 + actor_rollout_ref.actor.ppo_mini_batch_size=16 + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 + actor_rollout_ref.actor.entropy_coeff=0 + actor_rollout_ref.actor.loss_agg_mode="token-mean" + actor_rollout_ref.actor.clip_ratio_low=0.2 + actor_rollout_ref.actor.clip_ratio_high=0.28 + actor_rollout_ref.actor.clip_ratio_c=10.0 + actor_rollout_ref.actor.use_kl_loss=False + actor_rollout_ref.actor.kl_loss_coef=0.0 + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.megatron.param_offload=${ACTOR_OFFLOAD} + actor_rollout_ref.actor.megatron.optimizer_offload=${ACTOR_OFFLOAD} + actor_rollout_ref.actor.megatron.grad_offload=${ACTOR_OFFLOAD} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${TRAIN_PP} + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${TRAIN_TP} + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${TRAIN_PP} + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${TRAIN_TP} + actor_rollout_ref.ref.megatron.param_offload=True + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=$ROLLOUT_NAME + actor_rollout_ref.rollout.mode=async + actor_rollout_ref.rollout.n=4 + actor_rollout_ref.rollout.calculate_log_probs=True + actor_rollout_ref.rollout.gpu_memory_utilization=0.60 + actor_rollout_ref.rollout.temperature=1.0 + actor_rollout_ref.rollout.top_p=1.0 + actor_rollout_ref.rollout.top_k=-1 + actor_rollout_ref.rollout.enable_chunked_prefill=True + actor_rollout_ref.rollout.disable_log_stats=False + actor_rollout_ref.rollout.max_model_len=$MAX_NUM_TOKENS + actor_rollout_ref.rollout.max_num_batched_tokens=$MAX_NUM_TOKENS + actor_rollout_ref.rollout.max_num_seqs=$MAX_NUM_TOKENS + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.tensor_model_parallel_size=${GEN_TP} + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 + actor_rollout_ref.rollout.val_kwargs.temperature=1.0 + actor_rollout_ref.rollout.val_kwargs.top_p=0.7 + actor_rollout_ref.rollout.val_kwargs.top_k=-1 + actor_rollout_ref.rollout.val_kwargs.do_sample=True + actor_rollout_ref.rollout.val_kwargs.n=1 + actor_rollout_ref.rollout.agent.num_workers=1 + actor_rollout_ref.rollout.checkpoint_engine.backend='nccl' + actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=1024 + actor_rollout_ref.rollout.enforce_eager=False + +actor_rollout_ref.rollout.engine_kwargs.vllm.mm_processor_cache_gb=0 +) + +# Multi-teacher: one teacher per dataset, routed by the sample's `data_source` value. +DISTILLATION=( + distillation.enabled=True + distillation.teacher_key=data_source + distillation.n_gpus_per_node=${N_GPUS_TEACHER_TOTAL} + distillation.nnodes=1 + # --- gsm8k teacher (text-only) --- + +distillation.teacher_models.gsm8k.key="openai/gsm8k" + +distillation.teacher_models.gsm8k.model_path="${GSM8K_TEACHER_MODEL}" + +distillation.teacher_models.gsm8k.num_replicas=1 + +distillation.teacher_models.gsm8k.inference.name=$ROLLOUT_NAME + +distillation.teacher_models.gsm8k.inference.tensor_model_parallel_size=1 + +distillation.teacher_models.gsm8k.inference.gpu_memory_utilization=0.7 + +distillation.teacher_models.gsm8k.inference.enforce_eager=False + +distillation.teacher_models.gsm8k.inference.max_model_len=$MAX_NUM_TOKENS + +distillation.teacher_models.gsm8k.inference.max_num_batched_tokens=$MAX_NUM_TOKENS + +distillation.teacher_models.gsm8k.inference.max_num_seqs=$MAX_NUM_TOKENS + # --- geo3k teacher (vision-language) --- + +distillation.teacher_models.geo3k.key="hiyouga/geometry3k" + +distillation.teacher_models.geo3k.model_path="${GEO3K_TEACHER_MODEL}" + +distillation.teacher_models.geo3k.num_replicas=1 + +distillation.teacher_models.geo3k.inference.name=$ROLLOUT_NAME + +distillation.teacher_models.geo3k.inference.tensor_model_parallel_size=1 + +distillation.teacher_models.geo3k.inference.gpu_memory_utilization=0.7 + +distillation.teacher_models.geo3k.inference.enforce_eager=False + +distillation.teacher_models.geo3k.inference.max_model_len=$MAX_NUM_TOKENS + +distillation.teacher_models.geo3k.inference.max_num_batched_tokens=$MAX_NUM_TOKENS + +distillation.teacher_models.geo3k.inference.max_num_seqs=$MAX_NUM_TOKENS + +distillation.teacher_models.geo3k.inference.engine_kwargs.vllm.mm_processor_cache_gb=0 + # --- loss --- + distillation.distillation_loss.loss_mode=$DISTILLATION_LOSS_MODE + distillation.distillation_loss.topk=64 + distillation.distillation_loss.use_task_rewards=False + distillation.distillation_loss.use_policy_gradient=$USE_POLICY_GRADIENT + distillation.distillation_loss.loss_max_clamp=10.0 + distillation.distillation_loss.log_prob_min_clamp=-10.0 +) + +ALGORITHM=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + algorithm.kl_ctrl.kl_coef=0.0 +) + +REWARD=( + reward.reward_manager.name=dapo + +reward.reward_kwargs.overlong_buffer_cfg.enable=False + +reward.reward_kwargs.overlong_buffer_cfg.len=128 + +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=1.0 + +reward.reward_kwargs.overlong_buffer_cfg.log=False + +reward.reward_kwargs.max_resp_len=${MAX_RESPONSE_LENGTH} +) + +TRAINER=( + trainer.logger='["console"]' + trainer.project_name='verl-test-fully-async-opd' + trainer.experiment_name="qwen3-vl-2b-fully-async-multi-teacher-opd" + trainer.val_before_train=False + trainer.save_freq=-1 + trainer.resume_mode=disable + trainer.nnodes=1 + trainer.n_gpus_per_node=${N_GPUS_TRAINING} + trainer.log_val_generations=10 + +trainer.use_legacy_worker_impl=disable + trainer.total_epochs=2 + trainer.test_freq=-1 +) + +ASYNC_TRAINING=( + rollout.nnodes=1 + rollout.n_gpus_per_node=${N_GPUS_ROLLOUT} + rollout.total_rollout_steps=${TOTAL_ROLLOUT_STEPS} + async_training.staleness_threshold=${STALENESS_THRESHOLD} + async_training.partial_rollout=True + async_training.trigger_parameter_sync_step=${TRIGGER_PARAMETER_SYNC_STEP} + async_training.use_trainer_do_validate=False +) + +############################ Launch ############################ + +echo "Running fully_async_policy + Multi-Teacher OPD" +echo "Student: ${STUDENT_MODEL}" +echo "Teacher GSM8K: ${GSM8K_TEACHER_MODEL}" +echo "Teacher Geo3K: ${GEO3K_TEACHER_MODEL}" +echo "GPUs: ${N_GPUS_ROLLOUT} rollout + ${N_GPUS_TRAINING} training + ${N_GPUS_TEACHER_TOTAL} teachers" + +python3 -m verl.experimental.fully_async_policy.fully_async_main \ + --config-path=config \ + --config-name='fully_async_ppo_megatron_trainer.yaml' \ + actor_rollout_ref.hybrid_engine=False \ + critic.strategy=megatron \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${STUDENT[@]}" \ + "${ROLLOUT[@]}" \ + "${DISTILLATION[@]}" \ + "${ALGORITHM[@]}" \ + "${REWARD[@]}" \ + "${TRAINER[@]}" \ + "${ASYNC_TRAINING[@]}" \ + "$@" + +echo "Fully async policy + Multi-Teacher OPD E2E test completed successfully" diff --git a/verl/tests/special_e2e/run_one_step_off_policy.sh b/verl/tests/special_e2e/run_one_step_off_policy.sh new file mode 100644 index 0000000000000000000000000000000000000000..df195062f0f50ca615a75fd673a3fa4d6caad813 --- /dev/null +++ b/verl/tests/special_e2e/run_one_step_off_policy.sh @@ -0,0 +1,210 @@ +#!/usr/bin/env bash +set -xeuo pipefail + +# Test script for one_step_off_policy E2E regression testing +# This script runs one_step_off_policy with both FSDP2 and Megatron backends +# to ensure the asynchronous training mechanism works correctly + +NUM_GPUS=${NUM_GPUS:-8} +ACTOR_STRATEGY=${ACTOR_STRATEGY:-"fsdp2"} # fsdp2 or megatron + +# Download model if not exists +MODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-0.5B-Instruct} +MODEL_PATH=${MODEL_PATH:-${HOME}/models/${MODEL_ID}} +#hf download "${MODEL_ID}" --local-dir "${MODEL_PATH}" + +# Algorithm parameters +adv_estimator=grpo + +use_kl_in_reward=False +kl_coef=0.0 +use_kl_loss=False +kl_loss_coef=0.0 + +clip_ratio_low=0.2 +clip_ratio_high=0.28 + +# Response length parameters +max_prompt_length=1024 +max_response_length=2048 +enable_overlong_buffer=True +overlong_buffer_len=128 +overlong_penalty_factor=1.0 + +# Training parameters +loss_agg_mode="token-mean" +train_prompt_bsz=8 +n_resp_per_prompt=3 +train_prompt_mini_bsz=4 + +# Temperature parameters +temperature=1.0 +top_p=1.0 +top_k=-1 +val_top_p=0.7 + +# One-step-off-policy specific parameters +# Allocate 2 GPUs for rollout, remaining for training +n_gpus_rollout=2 +n_gpus_training=$((NUM_GPUS - n_gpus_rollout)) + +exp_name="$(basename "${MODEL_ID,,}")-one-step-off-policy-${ACTOR_STRATEGY}-minimal" + +echo "Running one_step_off_policy with ${ACTOR_STRATEGY} strategy" +echo "Total GPUs: ${NUM_GPUS}, Rollout GPUs: ${n_gpus_rollout}, Training GPUs: ${n_gpus_training}" + +# Common parameters for both FSDP2 and Megatron +common_params=( + data.train_files="${HOME}/data/gsm8k/train.parquet" + data.val_files="${HOME}/data/gsm8k/test.parquet" + data.prompt_key=prompt + data.truncation='left' + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.train_batch_size=${train_prompt_bsz} + actor_rollout_ref.rollout.n=${n_resp_per_prompt} + algorithm.adv_estimator=${adv_estimator} + algorithm.use_kl_in_reward=${use_kl_in_reward} + algorithm.kl_ctrl.kl_coef=${kl_coef} + actor_rollout_ref.hybrid_engine=False \ + actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} + actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} + actor_rollout_ref.actor.clip_ratio_c=10.0 + actor_rollout_ref.model.path="${MODEL_PATH}" + actor_rollout_ref.actor.optim.lr=1e-6 + actor_rollout_ref.actor.optim.lr_warmup_steps=-1 + actor_rollout_ref.actor.optim.weight_decay=0.1 + actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} + actor_rollout_ref.actor.entropy_coeff=0 + actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} + actor_rollout_ref.rollout.gpu_memory_utilization=0.80 + actor_rollout_ref.rollout.temperature=${temperature} + actor_rollout_ref.rollout.top_p=${top_p} + actor_rollout_ref.rollout.top_k=${top_k} + actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} + actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} + actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} + actor_rollout_ref.rollout.val_kwargs.do_sample=True + actor_rollout_ref.rollout.val_kwargs.n=1 + actor_rollout_ref.rollout.enable_chunked_prefill=True + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.checkpoint_engine.backend='nccl' + actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=1024 + reward.reward_manager.name=dapo + +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} + +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} + +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} + +reward.reward_kwargs.overlong_buffer_cfg.log=False + +reward.reward_kwargs.max_resp_len=${max_response_length} + trainer.logger=['console'] + trainer.project_name='verl-test' + trainer.experiment_name="${exp_name}" + trainer.val_before_train=True + trainer.test_freq=-1 + trainer.save_freq=-1 + trainer.total_epochs=2 + trainer.total_training_steps=2 + trainer.resume_mode=disable + trainer.nnodes=1 + trainer.n_gpus_per_node=${n_gpus_training} + rollout.nnodes=1 + rollout.n_gpus_per_node=${n_gpus_rollout} + +) + + # Detect device + device_name=$(python3 - <<'EOF' +from verl.utils.device import get_device_name +print(get_device_name()) +EOF +) + +if [ "${ACTOR_STRATEGY}" == "fsdp2" ]; then + echo "Running with FSDP2 strategy..." + # FSDP2 specific parameters + gen_tp=2 + sp_size=2 + fsdp_size=2 + ref_offload=True + actor_offload=False + + if [ "$device_name" ] && [ "$device_name" == "npu" ]; then + common_params+=( + # Todo The checkpoint_engine.backend should be unified to nccl + # actor_rollout_ref.rollout.checkpoint_engine.backend='hccl' + actor_rollout_ref.rollout.gpu_memory_utilization=0.60 + ) + actor_offload=True + fi + + python3 -m verl.experimental.one_step_off_policy.main_ppo \ + "${common_params[@]}" \ + actor_rollout_ref.actor.fsdp_config.strategy=fsdp2 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=2 \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=2 \ + critic.strategy=fsdp2 \ + actor_rollout_ref.actor.grad_clip=1.0 \ + actor_rollout_ref.model.use_remove_padding=True \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.use_dynamic_bsz=True \ + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True \ + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \ + actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \ + actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \ + actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \ + actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \ + actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \ + actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} $@ + +elif [ "${ACTOR_STRATEGY}" == "megatron" ]; then + echo "Running with Megatron strategy..." + # Megatron specific parameters + gen_tp=2 + train_tp=2 + train_pp=3 + ref_offload=True + actor_offload=False + + if [ "$device_name" ] && [ "$device_name" == "npu" ]; then + common_params+=( + # Todo The checkpoint_engine.backend should be unified to nccl + # actor_rollout_ref.rollout.checkpoint_engine.backend='hccl' + actor_rollout_ref.rollout.gpu_memory_utilization=0.70 + actor_rollout_ref.model.use_remove_padding=True \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.use_dynamic_bsz=True \ + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True \ + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \ + ) + train_tp=2 + actor_offload=True + fi + + python3 -m verl.experimental.one_step_off_policy.main_ppo \ + --config-path=config \ + --config-name='one_step_off_ppo_megatron_trainer.yaml' \ + "${common_params[@]}" \ + actor_rollout_ref.actor.strategy=megatron \ + critic.strategy=megatron \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=2 \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=2 \ + actor_rollout_ref.actor.megatron.param_offload=${actor_offload} \ + actor_rollout_ref.actor.megatron.optimizer_offload=${actor_offload} \ + actor_rollout_ref.actor.megatron.grad_offload=${actor_offload} \ + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${train_pp} \ + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${train_tp} \ + actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \ + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${train_pp} \ + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${train_tp} \ + actor_rollout_ref.ref.megatron.param_offload=${ref_offload} $@ +else + echo "Error: Unknown strategy ${ACTOR_STRATEGY}. Please use 'fsdp2' or 'megatron'" + exit 1 +fi + +echo "One-step-off-policy E2E test completed successfully with ${ACTOR_STRATEGY} strategy" \ No newline at end of file diff --git a/verl/tests/special_e2e/run_ppo_trainer_megatron.sh b/verl/tests/special_e2e/run_ppo_trainer_megatron.sh new file mode 100644 index 0000000000000000000000000000000000000000..c9a237b76c9e7eb779678f9866eea004db941ff1 --- /dev/null +++ b/verl/tests/special_e2e/run_ppo_trainer_megatron.sh @@ -0,0 +1,318 @@ +#!/usr/bin/env bash +set -xeuo pipefail + +export CUDA_DEVICE_MAX_CONNECTIONS=1 # For megatron communication/computation overlapping +export VERL_LOGGING_LEVEL=INFO +export VERL_PPO_LOGGING_LEVEL=INFO + +NUM_GPUS=${NUM_GPUS:-8} + +MODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-0.5B} +MODEL_PATH=${MODEL_PATH:-${HOME}/models/${MODEL_ID}} +RM_MODEL_PATH=${RM_MODEL_PATH:-${HOME}/models/Skywork/Skywork-Reward-V2-Llama-3.2-1B} +#hf download "${MODEL_ID}" --local-dir "${MODEL_PATH}" + +USE_DUMMY_MODEL=${USE_DUMMY_MODEL:-False} +DUMMY_MODEL_PATH=${DUMMY_MODEL_PATH:-${HOME}/dummy_models/${MODEL_ID}} +if [ "$USE_DUMMY_MODEL" = "True" ]; then + if [ -z "${DUMMY_MODEL_CONFIG_PATH}" ]; then + echo "[ERROR] DUMMY_MODEL_CONFIG_PATH not set" + exit 1 + fi + + python scripts/init_random_model.py \ + --hf_model_path "${MODEL_PATH}" \ + --new_config_path "${DUMMY_MODEL_CONFIG_PATH}" \ + --output_path "${DUMMY_MODEL_PATH}" + + MODEL_PATH="${DUMMY_MODEL_PATH}" +fi + +TRAIN_FILES=${TRAIN_FILES:-${HOME}/data/gsm8k/train.parquet} +VAL_FILES=${VAL_FILES:-${HOME}/data/gsm8k/test.parquet} + +ADV_ESTIMATOR=${ADV_ESTIMATOR:-gae} +# Validation +VAL_BEFORE_TRAIN=${VAL_BEFORE_TRAIN:-False} +TEST_FREQ=${TEST_FREQ:--1} +# Save & Resume +RESUME_MODE=${RESUME_MODE:-disable} +SAVE_FREQ=${SAVE_FREQ:--1} +TOTAL_TRAIN_STEPS=${TOTAL_TRAIN_STEPS:-1} + +USE_DYNAMIC_BSZ=${USE_DYNAMIC_BSZ:-True} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN:-2400} +forward_max_token_len_per_gpu=${FWD_MAX_TOKEN_LEN:-4800} +train_traj_micro_bsz_per_gpu=${MICRO_BSZ:-2} # b +n_resp_per_prompt=4 # g + +train_traj_micro_bsz=$((train_traj_micro_bsz_per_gpu * 1)) # b * n +train_traj_mini_bsz=$((train_traj_micro_bsz * 2)) # 2 * b * n +train_prompt_mini_bsz=$((train_traj_mini_bsz * 2)) # 2 * b * n / g +train_prompt_bsz=$((train_prompt_mini_bsz * 2)) # 4 * b * n / g + +LORA_RANK=${LORA_RANK:-0} +CRITIC_LORA_RANK=${CRITIC_LORA_RANK:-$LORA_RANK} +LORA_ALPHA=${LORA_ALPHA:-${LORA_RANK}} +LORA_TARGET_MODULES=${LORA_TARGET_MODULES:-"['linear_qkv','linear_proj','linear_fc1','linear_fc2']"} +LORA_MERGE=${LORA_MERGE:-False} + +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-512} +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-512} +MAX_RM_LENGTH=$((MAX_PROMPT_LENGTH + MAX_RESPONSE_LENGTH)) + +COMMON_PP=${COMMON_PP:-2} +COMMON_VPP=${COMMON_VPP:-2} +COMMON_CP=${COMMON_CP:-2} +COMMON_TP=${COMMON_TP:-2} +COMMON_EP=${COMMON_EP:-1} +COMMON_ETP=${COMMON_ETP:-1} + +TRAIN_TP=${TRAIN_TP:-$COMMON_TP} +INFER_TP=${INFER_TP:-$COMMON_TP} + +ACTOR_PP=${ACTOR_PP:-$COMMON_PP} +ACTOR_VPP=${ACTOR_VPP:-$COMMON_VPP} +ACTOR_CP=${ACTOR_CP:-$COMMON_CP} +ACTOR_TP=${ACTOR_TP:-$TRAIN_TP} +ACTOR_EP=${ACTOR_EP:-$COMMON_EP} +ACTOR_ETP=${ACTOR_ETP:-$COMMON_ETP} +ROLLOUT_TP=${ROLLOUT_TP:-$INFER_TP} +REF_PP=${REF_PP:-$COMMON_PP} +REF_VPP=${REF_VPP:-$COMMON_VPP} +REF_CP=${REF_CP:-$COMMON_CP} +REF_TP=${REF_TP:-$TRAIN_TP} +REF_EP=${REF_EP:-$COMMON_EP} +REF_ETP=${REF_ETP:-$COMMON_ETP} +CRITIC_PP=${CRITIC_PP:-$COMMON_PP} +CRITIC_VPP=${CRITIC_VPP:-$COMMON_VPP} +CRITIC_CP=${CRITIC_CP:-$COMMON_CP} +CRITIC_TP=${CRITIC_TP:-$TRAIN_TP} +CRITIC_EP=${CRITIC_EP:-$COMMON_EP} +CRITIC_ETP=${CRITIC_ETP:-$COMMON_ETP} + +ALL_OFFLOAD=${ALL_OFFLOAD:-True} +COMMON_PARAM_OFFLOAD=${COMMON_PARAM_OFFLOAD:-$ALL_OFFLOAD} +COMMON_GRAD_OFFLOAD=${COMMON_GRAD_OFFLOAD:-$ALL_OFFLOAD} +COMMON_OPTIMIZER_OFFLOAD=${COMMON_OPTIMIZER_OFFLOAD:-$ALL_OFFLOAD} + +ACTOR_PARAM_OFFLOAD=${ACTOR_PARAM_OFFLOAD:-$COMMON_PARAM_OFFLOAD} +ACTOR_GRAD_OFFLOAD=${ACTOR_GRAD_OFFLOAD:-$COMMON_GRAD_OFFLOAD} +ACTOR_OPTIMIZER_OFFLOAD=${ACTOR_OPTIMIZER_OFFLOAD:-$COMMON_OPTIMIZER_OFFLOAD} +REF_PARAM_OFFLOAD=${REF_PARAM_OFFLOAD:-$COMMON_PARAM_OFFLOAD} +CRITIC_PARAM_OFFLOAD=${CRITIC_PARAM_OFFLOAD:-$COMMON_PARAM_OFFLOAD} +CRITIC_GRAD_OFFLOAD=${CRITIC_GRAD_OFFLOAD:-$COMMON_GRAD_OFFLOAD} +CRITIC_OPTIMIZER_OFFLOAD=${CRITIC_OPTIMIZER_OFFLOAD:-$COMMON_OPTIMIZER_OFFLOAD} +RM_PARAM_OFFLOAD=${RM_PARAM_OFFLOAD:-$COMMON_PARAM_OFFLOAD} +USE_MBRIDGE=${USE_MBRIDGE:-True} +VANILLA_MBRIDGE=${VANILLA_MBRIDGE:-True} +VALUE_VANILLA_MBRIDGE=${VALUE_VANILLA_MBRIDGE:-$VANILLA_MBRIDGE} +USE_MEGATRON_FSDP=${USE_MEGATRON_FSDP:-False} +USE_FUSED_KERNELS=${USE_FUSED_KERNELS:-False} + +LR_WARMUP_STEPS=${LR_WARMUP_STEPS:-null} + +CHECKPOINT_CONTENTS=['model','hf_model','optimizer','extra'] +SKIP_SAVE_HF_MODEL=${SKIP_SAVE_HF_MODEL:-0} +if [ $SKIP_SAVE_HF_MODEL -eq 1 ]; then + CHECKPOINT_CONTENTS=['model','optimizer','extra'] +fi + +USE_DIST_CKPT=${USE_DIST_CKPT:-False} +DIST_CKPT_PATH=${DIST_CKPT_PATH:-${HOME}/dist_ckpt/${MODEL_ID}} +if [ "$USE_DIST_CKPT" = "True" ]; then + if [ "$USE_DUMMY_MODEL" = "True" ]; then + DIST_CKPT_PATH=${HOME}/dist_ckpt_dummy/${MODEL_ID} + fi + python scripts/converter_hf_to_mcore.py \ + --hf_model_path "${MODEL_PATH}" \ + --output_path "${DIST_CKPT_PATH}" +fi + +ENGINE=${ENGINE:-"vllm"} +if [ "$ENGINE" = "vllm" ]; then + export VLLM_USE_V1=1 +fi + +exp_name="$(basename "${MODEL_ID,,}")-megatron-gsm8k-minimal" +ROLLOUT_MODE="async" +ROLLOUT_QUANTIZATION=${ROLLOUT_QUANTIZATION:-null} + +RETURN_RAW_CHAT="True" +SKIP_TOKENIZER_INIT="True" + +OPTIM_MEMORY_EFFICIENT=${OPTIM_MEMORY_EFFICIENT:-False} + +PROFILE_ENABLE=${PROFILE_ENABLE:-False} +PROFILE_STEPS=${PROFILE_STEPS:-[1]} +PROFILE_RANKS_ALL=${PROFILE_RANKS_ALL:-True} +PROFILE_RANKS=${PROFILE_RANKS:-[0,1,2,3]} +DISCRETE=${DISCRETE:-True} # or True + +USE_REMOVE_PADDING=${USE_REMOVE_PADDING:-False} +ROUTING_REPLAY_MODE=${ROUTING_REPLAY_MODE:-"disabled"} + +if [ "$ROUTING_REPLAY_MODE" = "R3" ]; then + ENABLE_ROLLOUT_ROUTING_REPLAY=True +else + ENABLE_ROLLOUT_ROUTING_REPLAY=False +fi + +common_params=( + algorithm.adv_estimator="${ADV_ESTIMATOR}" + data.train_files="${TRAIN_FILES}" + data.val_files="${VAL_FILES}" + data.train_batch_size=${train_prompt_bsz} + data.max_prompt_length=${MAX_PROMPT_LENGTH} + data.max_response_length=${MAX_RESPONSE_LENGTH} + data.return_raw_chat=${RETURN_RAW_CHAT} + data.filter_overlong_prompts=True + data.truncation='error' + actor_rollout_ref.model.path="${MODEL_PATH}" + actor_rollout_ref.model.use_fused_kernels=${USE_FUSED_KERNELS} + actor_rollout_ref.model.use_remove_padding=${USE_REMOVE_PADDING} + actor_rollout_ref.model.lora.rank=${LORA_RANK} + actor_rollout_ref.model.lora.alpha=${LORA_ALPHA} + actor_rollout_ref.model.lora.target_modules=${LORA_TARGET_MODULES} + actor_rollout_ref.model.lora.merge=${LORA_MERGE} + +actor_rollout_ref.model.lora.fully_sharded_loras=True + actor_rollout_ref.actor.optim.lr_warmup_steps=$LR_WARMUP_STEPS + actor_rollout_ref.actor.megatron.router_replay.mode=${ROUTING_REPLAY_MODE} + +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=$OPTIM_MEMORY_EFFICIENT + +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=$OPTIM_MEMORY_EFFICIENT + +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=$OPTIM_MEMORY_EFFICIENT + actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${train_traj_micro_bsz_per_gpu} + actor_rollout_ref.actor.use_dynamic_bsz=${USE_DYNAMIC_BSZ} + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.megatron.use_mbridge=${USE_MBRIDGE} + actor_rollout_ref.actor.megatron.vanilla_mbridge=${VANILLA_MBRIDGE} + actor_rollout_ref.actor.megatron.use_megatron_fsdp=${USE_MEGATRON_FSDP} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=$ACTOR_PP + actor_rollout_ref.actor.megatron.virtual_pipeline_model_parallel_size=$ACTOR_VPP + actor_rollout_ref.actor.megatron.context_parallel_size=$ACTOR_CP + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=$ACTOR_TP + actor_rollout_ref.actor.megatron.expert_model_parallel_size=$ACTOR_EP + actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=$ACTOR_ETP + actor_rollout_ref.actor.megatron.param_offload=${ACTOR_PARAM_OFFLOAD} + actor_rollout_ref.actor.megatron.optimizer_offload=${ACTOR_OPTIMIZER_OFFLOAD} + actor_rollout_ref.actor.megatron.grad_offload=${ACTOR_GRAD_OFFLOAD} + actor_rollout_ref.actor.megatron.use_dist_checkpointing=${USE_DIST_CKPT} + actor_rollout_ref.actor.megatron.dist_checkpointing_path=${DIST_CKPT_PATH} + actor_rollout_ref.actor.use_kl_loss=True + actor_rollout_ref.actor.kl_loss_coef=0.001 + actor_rollout_ref.actor.kl_loss_type=low_var_kl + actor_rollout_ref.actor.checkpoint.save_contents=$CHECKPOINT_CONTENTS + actor_rollout_ref.actor.profiler.enable=$PROFILE_ENABLE + actor_rollout_ref.actor.profiler.ranks=$PROFILE_RANKS + actor_rollout_ref.actor.profiler.all_ranks=$PROFILE_RANKS_ALL + actor_rollout_ref.rollout.name="${ENGINE}" + actor_rollout_ref.rollout.mode="${ROLLOUT_MODE}" + actor_rollout_ref.rollout.tensor_model_parallel_size=$ROLLOUT_TP + actor_rollout_ref.rollout.gpu_memory_utilization=0.6 + actor_rollout_ref.rollout.n=${n_resp_per_prompt} + ++actor_rollout_ref.rollout.quantization=${ROLLOUT_QUANTIZATION} + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${train_traj_micro_bsz_per_gpu} + actor_rollout_ref.rollout.enable_rollout_routing_replay=${ENABLE_ROLLOUT_ROUTING_REPLAY} + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${train_traj_micro_bsz_per_gpu} + actor_rollout_ref.ref.megatron.use_mbridge=${USE_MBRIDGE} + actor_rollout_ref.ref.megatron.vanilla_mbridge=${VANILLA_MBRIDGE} + actor_rollout_ref.ref.megatron.use_megatron_fsdp=${USE_MEGATRON_FSDP} + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=$REF_PP + actor_rollout_ref.ref.megatron.virtual_pipeline_model_parallel_size=$REF_VPP + actor_rollout_ref.ref.megatron.context_parallel_size=$REF_CP + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=$REF_TP + actor_rollout_ref.ref.megatron.expert_model_parallel_size=$REF_EP + actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=$REF_ETP + actor_rollout_ref.ref.megatron.param_offload=${REF_PARAM_OFFLOAD} + actor_rollout_ref.ref.megatron.use_dist_checkpointing=${USE_DIST_CKPT} + actor_rollout_ref.ref.megatron.dist_checkpointing_path=${DIST_CKPT_PATH} + critic.optim.lr=2e-5 + critic.optim.lr_warmup_steps=$LR_WARMUP_STEPS + +critic.optim.override_optimizer_config.optimizer_cpu_offload=$OPTIM_MEMORY_EFFICIENT + +critic.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=$OPTIM_MEMORY_EFFICIENT + +critic.optim.override_optimizer_config.use_precision_aware_optimizer=$OPTIM_MEMORY_EFFICIENT + critic.model.path="${MODEL_PATH}" + critic.model.lora.rank=${CRITIC_LORA_RANK} + critic.model.lora.alpha=${LORA_ALPHA} + critic.model.lora.target_modules=${LORA_TARGET_MODULES} + critic.ppo_micro_batch_size_per_gpu=${train_traj_micro_bsz_per_gpu} + critic.ppo_max_token_len_per_gpu=${forward_max_token_len_per_gpu} + critic.megatron.use_mbridge=${USE_MBRIDGE} + critic.megatron.vanilla_mbridge=${VALUE_VANILLA_MBRIDGE} + critic.megatron.use_megatron_fsdp=${USE_MEGATRON_FSDP} + critic.megatron.pipeline_model_parallel_size=$CRITIC_PP + critic.megatron.virtual_pipeline_model_parallel_size=$CRITIC_VPP + critic.megatron.context_parallel_size=$CRITIC_CP + critic.megatron.tensor_model_parallel_size=$CRITIC_TP + critic.megatron.expert_model_parallel_size=$CRITIC_EP + critic.megatron.expert_tensor_parallel_size=$CRITIC_ETP + critic.megatron.param_offload=${CRITIC_PARAM_OFFLOAD} + critic.megatron.optimizer_offload=${CRITIC_OPTIMIZER_OFFLOAD} + critic.megatron.grad_offload=${CRITIC_GRAD_OFFLOAD} + critic.megatron.use_dist_checkpointing=${USE_DIST_CKPT} + critic.megatron.dist_checkpointing_path=${DIST_CKPT_PATH} + critic.checkpoint.save_contents=$CHECKPOINT_CONTENTS + critic.profiler.enable=$PROFILE_ENABLE + critic.profiler.ranks=$PROFILE_RANKS + critic.profiler.all_ranks=$PROFILE_RANKS_ALL + reward.num_workers=8 + reward.reward_model.enable=True + reward.reward_model.model_path="${RM_MODEL_PATH}" + reward.reward_model.rollout.name=${ENGINE} + reward.reward_model.rollout.gpu_memory_utilization=0.6 + reward.reward_model.rollout.tensor_model_parallel_size=${INFER_TP} + reward.reward_model.rollout.prompt_length=${MAX_RM_LENGTH} + reward.reward_model.rollout.response_length=${MAX_RESPONSE_LENGTH} + algorithm.use_kl_in_reward=False + algorithm.kl_penalty=kl + algorithm.kl_ctrl.kl_coef=0.001 + trainer.critic_warmup=0 + trainer.logger=console + trainer.project_name='verl-test' + trainer.experiment_name="${exp_name}" + trainer.nnodes=1 + trainer.n_gpus_per_node=${NUM_GPUS} + trainer.val_before_train="${VAL_BEFORE_TRAIN}" + trainer.test_freq="${TEST_FREQ}" + trainer.save_freq="${SAVE_FREQ}" + trainer.resume_mode="${RESUME_MODE}" + trainer.total_epochs=2 + trainer.total_training_steps="${TOTAL_TRAIN_STEPS}" + global_profiler.profile_continuous_steps=True + global_profiler.tool=nsys + global_profiler.steps=$PROFILE_STEPS + global_profiler.global_tool_config.nsys.discrete=$DISCRETE +) + + # Detect device + device_name=$(python3 - <<'EOF' +from verl.utils.device import get_device_name +print(get_device_name()) +EOF +) + +if [ -n "$device_name" ] && [ "$device_name" == "cuda" ]; then + python3 -m verl.trainer.main_ppo \ + --config-path=config \ + --config-name='ppo_megatron_trainer.yaml' \ + "${common_params[@]}" $@ + +elif [ -n "$device_name" ] && [ "$device_name" == "npu" ]; then + python3 -m verl.trainer.main_ppo \ + --config-path=config \ + --config-name='ppo_megatron_trainer.yaml' \ + "${common_params[@]}" \ + +actor_rollout_ref.actor.megatron.override_transformer_config.context_parallel_size=${ACTOR_CP} \ + +actor_rollout_ref.actor.megatron.override_transformer_config.use_flash_attn=True \ + ++actor_rollout_ref.ref.megatron.override_transformer_config.use_flash_attn=True \ + global_profiler.tool=npu \ + actor_rollout_ref.actor.profiler.tool_config.npu.contents=[npu,cpu,memory,shapes,module] \ + actor_rollout_ref.actor.profiler.tool_config.npu.level='level1' \ + actor_rollout_ref.actor.profiler.tool_config.npu.discrete=False \ + actor_rollout_ref.actor.profiler.tool_config.npu.analysis=False \ + global_profiler.save_path="${HOME}/profiling" $@ +else + echo "Unknown device: $device_name" + exit 1 +fi diff --git a/verl/tests/special_e2e/run_ppo_trainer_veomni.sh b/verl/tests/special_e2e/run_ppo_trainer_veomni.sh new file mode 100644 index 0000000000000000000000000000000000000000..e43915c0e1581b3597531c01efb977010b277c84 --- /dev/null +++ b/verl/tests/special_e2e/run_ppo_trainer_veomni.sh @@ -0,0 +1,122 @@ +#!/usr/bin/env bash +set -xeuo pipefail + + +SAVE_PATH=tests/utils/ci/profiler_data +rm -rf "$SAVE_PATH" + +PROFILE_STEPS=[1] +PROFILE_RANKS_ALL=False +PROFILE_RANKS=[0] +DISCRETE=True + +# Download model if not exists +MODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-0.5B-Instruct} +MODEL_PATH=${MODEL_PATH:-${HOME}/models/${MODEL_ID}} +#huggingface-cli download "${MODEL_ID}" --local-dir "${MODEL_PATH}" + +TRAIN_FILES=${TRAIN_FILES:-${HOME}/data/gsm8k/train.parquet} +VAL_FILES=${VAL_FILES:-${HOME}/data/gsm8k/test.parquet} +VAL_BEFORE_TRAIN=${VAL_BEFORE_TRAIN:-True} +NUM_GPUS=${NUM_GPUS:-8} +FSDP_SIZE=${FSDP_SIZE:-4} +SP_SIZE=${SP_SIZE:-2} +EP_SIZE=${EP_SIZE:-1} +MODEL_NAME_ONLY=${MODEL_ID##*/} +VERL_EXP_NAME=${VERL_EXP_NAME:-${MODEL_NAME_ONLY}-function-reward-minimal-fsdp-size${FSDP_SIZE}} + +device_name=$(python3 - <<'EOF' +from verl.utils.device import get_device_name +print(get_device_name()) +EOF +) + +common_params=( + model_engine=veomni \ + algorithm.adv_estimator=grpo \ + data.train_files="${TRAIN_FILES}" \ + data.val_files="${VAL_FILES}" \ + data.train_batch_size=16 \ + data.max_prompt_length=512 \ + data.max_response_length=128 \ + data.filter_overlong_prompts=True \ + data.truncation='error' \ + actor_rollout_ref.model.path="${MODEL_PATH}" \ + actor_rollout_ref.actor.optim.lr=5e-7 \ + actor_rollout_ref.model.use_remove_padding=True \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.veomni.param_offload=True \ + actor_rollout_ref.actor.veomni.optimizer_offload=True \ + actor_rollout_ref.actor.ppo_mini_batch_size=8 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.actor.use_kl_loss=True \ + actor_rollout_ref.actor.kl_loss_coef=0.001 \ + actor_rollout_ref.actor.kl_loss_type=low_var_kl \ + actor_rollout_ref.actor.entropy_coeff=0 \ + actor_rollout_ref.actor.use_torch_compile=False \ + actor_rollout_ref.actor.veomni.fsdp_size="${FSDP_SIZE}" \ + actor_rollout_ref.actor.veomni.ulysses_parallel_size="${SP_SIZE}" \ + actor_rollout_ref.actor.veomni.expert_parallel_size="${EP_SIZE}" \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.ref.veomni.param_offload=True \ + actor_rollout_ref.ref.use_torch_compile=False \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=2 \ + actor_rollout_ref.rollout.enable_chunked_prefill=False \ + actor_rollout_ref.rollout.name=vllm \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \ + actor_rollout_ref.rollout.free_cache_engine=True \ + actor_rollout_ref.rollout.enforce_eager=True \ + actor_rollout_ref.rollout.n=2 \ + actor_rollout_ref.ref.veomni.optimizer_offload=True \ + algorithm.kl_ctrl.kl_coef=0.001 \ + trainer.critic_warmup=0 \ + trainer.logger=console \ + trainer.project_name='verl_veomni_test' \ + trainer.experiment_name="${VERL_EXP_NAME}" \ + trainer.n_gpus_per_node="${NUM_GPUS}" \ + trainer.val_before_train="${VAL_BEFORE_TRAIN}" \ + trainer.nnodes=1 \ + trainer.save_freq=-1 \ + trainer.test_freq=-1 \ + trainer.total_epochs=1 \ + trainer.total_training_steps=1 \ + actor_rollout_ref.actor.profiler.enable=True \ + actor_rollout_ref.actor.profiler.all_ranks=$PROFILE_RANKS_ALL \ + actor_rollout_ref.actor.profiler.ranks=$PROFILE_RANKS \ + actor_rollout_ref.ref.profiler.enable=True \ + actor_rollout_ref.ref.profiler.all_ranks=$PROFILE_RANKS_ALL \ + actor_rollout_ref.ref.profiler.ranks=$PROFILE_RANKS \ + global_profiler.steps=$PROFILE_STEPS \ + global_profiler.save_path="$SAVE_PATH" \ +) + +if [ -n "$device_name" ] && [ "$device_name" == "cuda" ]; then + CONTENTS=['cuda'] + python3 -m verl.trainer.main_ppo \ + "${common_params[@]}" \ + actor_rollout_ref.actor.profiler.tool_config.torch.discrete=$DISCRETE \ + actor_rollout_ref.actor.profiler.tool_config.torch.contents=$CONTENTS \ + actor_rollout_ref.ref.profiler.tool_config.torch.discrete=$DISCRETE \ + actor_rollout_ref.ref.profiler.tool_config.torch.contents=$CONTENTS \ + global_profiler.tool=torch $@ + + python3 "tests/utils/test_check_profiler_output.py" --profiler_dir="$SAVE_PATH" --device="gpu" + +elif [ -n "$device_name" ] && [ "$device_name" == "npu" ]; then + CONTENTS=['npu','cpu'] + python3 -m verl.trainer.main_ppo \ + "${common_params[@]}" \ + actor_rollout_ref.actor.profiler.tool_config.npu.discrete=$DISCRETE \ + actor_rollout_ref.actor.profiler.tool_config.npu.contents=$CONTENTS \ + actor_rollout_ref.ref.profiler.tool_config.npu.discrete=$DISCRETE \ + actor_rollout_ref.ref.profiler.tool_config.npu.contents=$CONTENTS \ + global_profiler.tool=npu $@ + + python3 "tests/utils/test_check_profiler_output.py" --profiler_dir="$SAVE_PATH" --device="npu" +else + echo "Unknown device: $device_name" + exit 1 +fi + +rm -rf "$SAVE_PATH" diff --git a/verl/tests/special_e2e/sft/compare_sft_engine_results.py b/verl/tests/special_e2e/sft/compare_sft_engine_results.py new file mode 100644 index 0000000000000000000000000000000000000000..322f5353c06e7fd8463b9236a69b8fe078f9adb9 --- /dev/null +++ b/verl/tests/special_e2e/sft/compare_sft_engine_results.py @@ -0,0 +1,58 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import os + +import torch + + +def get_result(file): + file = os.path.expanduser(file) + result = [] + with open(file) as f: + lines = f.readlines() + for line in lines: + result.append(json.loads(line)) + return result + + +def compare_results(golden_results, other_result): + golden_loss = golden_results[0]["data"]["train/loss"] + golden_grad_norm = golden_results[0]["data"]["train/grad_norm"] + + loss = other_result[0]["data"]["train/loss"] + grad_norm = other_result[0]["data"]["train/grad_norm"] + + torch.testing.assert_close(golden_loss, loss, atol=1e-2, rtol=1e-2) + torch.testing.assert_close(golden_grad_norm, grad_norm, atol=1e-4, rtol=3e-2) + + +if __name__ == "__main__": + golden_results = get_result("~/verl/test/log/golden.jsonl") + + # get all other results + other_results = {} + # walk through all files in ~/verl/test/log + for file in os.listdir(os.path.expanduser("~/verl/test/log/verl_sft_test")): + if file.endswith(".jsonl"): + other_results[file] = get_result(os.path.join(os.path.expanduser("~/verl/test/log/verl_sft_test"), file)) + + # # compare results + for file, other_result in other_results.items(): + print(f"compare results {file}") + compare_results(golden_results, other_result) + print(f"compare results {file} done") + + print("All results are close to golden results") diff --git a/verl/tests/special_e2e/sft/run_sft.sh b/verl/tests/special_e2e/sft/run_sft.sh new file mode 100644 index 0000000000000000000000000000000000000000..4078ff260efc781de27be1313c22da3021469544 --- /dev/null +++ b/verl/tests/special_e2e/sft/run_sft.sh @@ -0,0 +1,58 @@ +#!/usr/bin/env bash +set -xeuo pipefail + +ENTRYPOINT=${ENTRYPOINT:-"-m verl.trainer.sft_trainer"} + +NUM_GPUS=${NUM_GPUS:-8} + +MODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-0.5B-Instruct} +MODEL_PATH=${MODEL_PATH:-${HOME}/models/${MODEL_ID}} +#hf download "${MODEL_ID}" --local-dir "${MODEL_PATH}" + +TRAIN_FILES=${TRAIN_FILES:-$HOME/data/gsm8k_sft/train.parquet} +VAL_FILES=${VAL_FILES:-$HOME/data/gsm8k_sft/test.parquet} + +SP_SIZE=${SP_SIZE:-1} +LIGER=${LIGER:-False} +MULTITURN=${MULTITURN:-False} +LORA_RANK=${LORA_RANK:-0} +RM_PAD=${RM_PAD:-True} + +TOTAL_TRAIN_STEP=${TOTAL_TRAIN_STEP:-1} +RESUME_MODE=${RESUME_MODE:-disable} +SAVE_FREQ=${SAVE_FREQ:-1} + +micro_bsz=2 +NUM_GPUS=8 + +project_name="verl-test" +exp_name="$(basename "${MODEL_ID,,}")-sft-minimal" +ckpts_home=${ckpts_home:-$HOME/${project_name}/${exp_name}} + +mkdir -p "${ckpts_home}" + +torchrun --standalone --nnodes=1 --nproc_per_node=${NUM_GPUS} ${ENTRYPOINT} \ + data.train_files="${TRAIN_FILES}" \ + data.val_files="${VAL_FILES}" \ + data.messages_key=messages \ + data.micro_batch_size_per_gpu=${micro_bsz} \ + optim.lr=1e-4 \ + engine=fsdp \ + engine.ulysses_sequence_parallel_size="${SP_SIZE}" \ + model.path="${MODEL_PATH}" \ + model.lora_rank="${LORA_RANK}" \ + model.lora_alpha=16 \ + model.target_modules=all-linear \ + model.use_liger="${LIGER}" \ + model.use_remove_padding="${RM_PAD}" \ + trainer.default_local_dir="${ckpts_home}" \ + trainer.project_name="${project_name}" \ + trainer.experiment_name="${exp_name}" \ + trainer.total_training_steps=${TOTAL_TRAIN_STEP} \ + trainer.save_freq=${SAVE_FREQ} \ + checkpoint.save_contents=[model,optimizer,extra,hf_model] \ + trainer.max_ckpt_to_keep=1 \ + trainer.resume_mode=${RESUME_MODE} \ + trainer.logger=['console'] $@ + +rm -rf "${ckpts_home:?}/*" \ No newline at end of file diff --git a/verl/tests/special_e2e/sft/run_sft_engine.sh b/verl/tests/special_e2e/sft/run_sft_engine.sh new file mode 100644 index 0000000000000000000000000000000000000000..e7350ee99cf747a2c2e66c0d26b880417bd353bc --- /dev/null +++ b/verl/tests/special_e2e/sft/run_sft_engine.sh @@ -0,0 +1,182 @@ +#!/usr/bin/env bash +set -xeuo pipefail + +NUM_GPUS=${NUM_GPUS:-1} + +mode=${mode:-spmd} + +if [ "$mode" = "spmd" ]; then + ENTRYPOINT=${ENTRYPOINT:-"-m verl.trainer.sft_trainer"} + COMMAND="torchrun --standalone --nnodes=${NNODES:-1} --nproc-per-node=${NUM_GPUS:-1} ${ENTRYPOINT}" +else + ENTRYPOINT=${ENTRYPOINT:-"-m verl.trainer.sft_trainer_ray"} + COMMAND="python ${ENTRYPOINT} trainer.nnodes=${NNODES:-1} trainer.n_gpus_per_node=${NUM_GPUS:-1}" +fi + +DATASET_DIR=${DATASET_DIR:-~/data/gsm8k_sft} +TRAIN_FILES=${DATASET_DIR}/train.parquet +VAL_FILES=${DATASET_DIR}/test.parquet + +backend=${BACKEND:-fsdp} + +project_name=verl_sft_test + +RESUME_MODE=disable + +ckpts_home=${ckpts_home:-~/verl/test/gsm8k-sft-${backend}} + +MODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-0.5B} +MODEL_PATH=${MODEL_PATH:-${HOME}/models/${MODEL_ID}} +#hf download "${MODEL_ID}" --local-dir "${MODEL_PATH}" + +SP_SIZE=${SP_SIZE:-1} +FSDP_SIZE=${FSDP_SIZE:-1} +FSDP_STRATEGY=${FSDP_STRATEGY:-"fsdp"} + +TP_SIZE=${TP_SIZE:-1} +PP_SIZE=${PP_SIZE:-1} +VPP_SIZE=${VPP_SIZE:-null} +CP_SIZE=${CP_SIZE:-1} + +PAD_MODE=${PAD_MODE:-no_padding} + +USE_REMOVE_PADDING=${USE_REMOVE_PADDING:-True} + +FSDP_ENGINE_CONFIG="\ + engine=${backend} \ + model=hf_model \ + model.path=$MODEL_PATH \ + optim=${backend} \ + optim.lr=1e-5 \ + optim.lr_warmup_steps_ratio=0.2 \ + optim.weight_decay=0.1 \ + optim.betas="[0.9,0.95]" \ + optim.clip_grad=1.0 \ + optim.min_lr_ratio=0.1 \ + optim.lr_scheduler_type=cosine \ + engine.ulysses_sequence_parallel_size=${SP_SIZE} \ + engine.strategy=${FSDP_STRATEGY} \ + engine.fsdp_size=${FSDP_SIZE}" + +VEOMNI_ENGINE_CONFIG="\ + engine=${backend} \ + model=hf_model \ + model.path=$MODEL_PATH \ + optim=${backend} \ + optim.lr=1e-5 \ + optim.lr_warmup_steps_ratio=0.2 \ + optim.weight_decay=0.1 \ + optim.betas="[0.9,0.95]" \ + optim.clip_grad=1.0 \ + optim.lr_min=1e-6 \ + optim.lr_scheduler_type=cosine \ + engine.ulysses_parallel_size=${SP_SIZE} \ + engine.fsdp_size=${FSDP_SIZE}" + +MEGATRON_ENGINE_CONFIG="\ + engine=${backend} \ + model=hf_model \ + model.path=$MODEL_PATH \ + optim=${backend} \ + optim.lr=1e-5 \ + optim.lr_warmup_steps_ratio=0.2 \ + optim.weight_decay=0.1 \ + optim.betas="[0.9,0.95]" \ + optim.clip_grad=1.0 \ + optim.lr_warmup_init=0 \ + optim.lr_decay_style=cosine \ + optim.min_lr=1e-6 \ + engine.tensor_model_parallel_size=${TP_SIZE} \ + engine.pipeline_model_parallel_size=${PP_SIZE} \ + engine.virtual_pipeline_model_parallel_size=${VPP_SIZE} \ + engine.context_parallel_size=${CP_SIZE} \ + +engine.override_transformer_config.context_parallel_size=${CP_SIZE} \ + engine.use_mbridge=True" + +TORCHTITAN_ENGINE_CONFIG="\ + engine=${backend} \ + model=hf_model \ + model.path=${MODEL_PATH} \ + optim=${backend} \ + optim.lr=1e-5 \ + optim.lr_warmup_steps_ratio=0.2 \ + optim.weight_decay=0.1 \ + optim.betas="[0.9,0.95]" \ + optim.clip_grad=1.0 \ + optim.min_lr_factor=0.1 \ + optim.decay_type=cosine \ + optim.total_training_steps=1000 \ + engine.tensor_parallel_size=${TP_SIZE} \ + engine.pipeline_parallel_size=${PP_SIZE} \ + engine.context_parallel_size=${CP_SIZE} \ + engine.data_parallel_shard_size=${FSDP_SIZE} \ + engine.use_torch_compile=False" + +AUTOMODEL_ENGINE_CONFIG="\ + engine=${backend} \ + model=hf_model \ + model.path=${MODEL_PATH} \ + optim=${backend} \ + optim.lr=1e-5 \ + optim.lr_warmup_steps_ratio=0.2 \ + optim.weight_decay=0.1 \ + optim.betas="[0.9,0.95]" \ + optim.clip_grad=1.0 \ + optim.min_lr_ratio=0.1 \ + optim.lr_scheduler_type=cosine \ + engine.tp_size=${TP_SIZE} \ + engine.cp_size=${CP_SIZE} \ + engine.use_torch_compile=False" + + +if [ "$backend" = "fsdp" ]; then + ENGINE_CONFIG="$FSDP_ENGINE_CONFIG" + echo "Using fsdp engine" + exp_name=gsm8k-${backend}-${FSDP_STRATEGY}-sp${SP_SIZE}-fsdp${FSDP_SIZE}-pad-${PAD_MODE}-use_remove_padding-${USE_REMOVE_PADDING}-mode-${mode} +elif [ "$backend" = "veomni" ]; then + ENGINE_CONFIG="$VEOMNI_ENGINE_CONFIG" + echo "Using veomni engine" + exp_name=gsm8k-${backend}-sp${SP_SIZE}-fsdp${FSDP_SIZE}-pad-${PAD_MODE}-use_remove_padding-${USE_REMOVE_PADDING}-mode-${mode} +elif [ "$backend" = "torchtitan" ]; then + ENGINE_CONFIG="$TORCHTITAN_ENGINE_CONFIG" + echo "Using torchtitan engine" + exp_name=gsm8k-${backend}-tp${TP_SIZE}-pp${PP_SIZE}-cp${CP_SIZE}-dp${FSDP_SIZE}-pad-${PAD_MODE}-use_remove_padding-${USE_REMOVE_PADDING}-mode-${mode} +elif [ "$backend" = "automodel" ]; then + ENGINE_CONFIG="$AUTOMODEL_ENGINE_CONFIG" + echo "Using automodel engine" + exp_name=gsm8k-${backend}-tp${TP_SIZE}-pp${PP_SIZE}-cp${CP_SIZE}-pad-${PAD_MODE}-use_remove_padding-${USE_REMOVE_PADDING}-mode-${mode} +else + ENGINE_CONFIG="$MEGATRON_ENGINE_CONFIG" + echo "Using megatron engine" + exp_name=gsm8k-${backend}-tp${TP_SIZE}-pp${PP_SIZE}-vpp${VPP_SIZE}-cp${CP_SIZE}-pad-${PAD_MODE}-use_remove_padding-${USE_REMOVE_PADDING}-mode-${mode} +fi + +mkdir -p "${ckpts_home}" + +$COMMAND \ + data.train_files="${TRAIN_FILES}" \ + data.val_files="${VAL_FILES}" \ + data.train_batch_size=128 \ + data.pad_mode=${PAD_MODE} \ + data.truncation=error \ + data.use_dynamic_bsz=True \ + data.max_token_len_per_gpu=2048 \ + data.messages_key=messages \ + model.use_remove_padding=${USE_REMOVE_PADDING} \ + data.ignore_input_ids_mismatch=True \ + ${ENGINE_CONFIG} \ + trainer.test_freq=after_each_epoch \ + trainer.save_freq=-1 \ + trainer.logger=['console','file'] \ + trainer.project_name="${project_name}" \ + trainer.experiment_name="${exp_name}" \ + trainer.total_epochs=2 \ + trainer.total_training_steps=2 \ + trainer.default_local_dir="${ckpts_home}" \ + trainer.resume_mode=${RESUME_MODE} \ + + # trainer.total_training_steps=${TOTAL_TRAIN_STEP} \ + # trainer.checkpoint.save_contents=[model,optimizer,extra,hf_model] \ + # trainer.max_ckpt_to_keep=1 \ + +rm -rf "${ckpts_home:?}/*" diff --git a/verl/tests/special_e2e/sft/test_sft_engine_all.sh b/verl/tests/special_e2e/sft/test_sft_engine_all.sh new file mode 100644 index 0000000000000000000000000000000000000000..2cf262e2a8de0a0aca884da248ea2c33e4317bac --- /dev/null +++ b/verl/tests/special_e2e/sft/test_sft_engine_all.sh @@ -0,0 +1,59 @@ +#!/usr/bin/env bash +set -xeuo pipefail + +rm -rf ~/verl/test/log +mkdir -p ~/verl/test/log + +export VERL_FILE_LOGGER_ROOT=~/verl/test/log +VPP_SIZE=${VPP_SIZE:-2} + +# test with single gpu as golden +echo "run with single gpu as golden" +BACKEND=fsdp SP_SIZE=1 FSDP_SIZE=1 NUM_GPUS=1 FSDP_STRATEGY=fsdp VERL_FILE_LOGGER_PATH=~/verl/test/log/golden.jsonl bash tests/special_e2e/sft/run_sft_engine.sh + +# test with fsdp 1 +echo "run with sp2 fsdp_size2 num_gpus8 fsdp_strategy fsdp pad_mode no_padding" +BACKEND=fsdp SP_SIZE=2 FSDP_SIZE=2 NUM_GPUS=8 FSDP_STRATEGY=fsdp PAD_MODE=no_padding bash tests/special_e2e/sft/run_sft_engine.sh + +# test with fsdp 1 use_remove_padding and pad_mode no_padding +echo "run with sp4 fsdp_size4 num_gpus8 fsdp_strategy fsdp pad_mode no_padding use_remove_padding False" +BACKEND=fsdp SP_SIZE=1 FSDP_SIZE=-1 NUM_GPUS=8 FSDP_STRATEGY=fsdp PAD_MODE=no_padding USE_REMOVE_PADDING=False bash tests/special_e2e/sft/run_sft_engine.sh + + +# test with fsdp 2 +echo "run with sp2 fsdp_size2 num_gpus8 fsdp_strategy fsdp2" +BACKEND=fsdp SP_SIZE=2 FSDP_SIZE=2 NUM_GPUS=8 FSDP_STRATEGY=fsdp2 bash tests/special_e2e/sft/run_sft_engine.sh + +# test with veomni +echo "run with sp2 fsdp_size4 num_gpus8 fsdp_strategy fsdp2 backend veomni" +BACKEND=veomni SP_SIZE=2 FSDP_SIZE=4 NUM_GPUS=8 FSDP_STRATEGY=fsdp2 bash tests/special_e2e/sft/run_sft_engine.sh + + +# test with megatron +echo "run with tp2 pp2 vpp2 cp2 num_gpus8" +BACKEND=megatron TP_SIZE=2 PP_SIZE=2 VPP_SIZE=${VPP_SIZE} CP_SIZE=2 NUM_GPUS=8 bash tests/special_e2e/sft/run_sft_engine.sh + +# test with cp in ray +echo "run with tp2 pp2 vpp2 cp2 num_gpus8 mode=ray" +BACKEND=megatron TP_SIZE=2 PP_SIZE=2 VPP_SIZE=${VPP_SIZE} CP_SIZE=2 NUM_GPUS=8 mode=ray bash tests/special_e2e/sft/run_sft_engine.sh + +# TODO: Will add back torchtitan CI once everything is ready +# # test with torchtitan fsdp=2 +# echo "run with tp1 pp1 cp1 fsdp2 num_gpus2" +# BACKEND=torchtitan TP_SIZE=1 PP_SIZE=1 CP_SIZE=1 FSDP_SIZE=2 NUM_GPUS=2 bash tests/special_e2e/sft/run_sft_engine.sh + +# # test with torchtitan tp2 fsdp=2 +# echo "run with tp2 pp1 cp1 fsdp2 num_gpus4" +# BACKEND=torchtitan TP_SIZE=2 PP_SIZE=1 CP_SIZE=1 FSDP_SIZE=2 NUM_GPUS=4 bash tests/special_e2e/sft/run_sft_engine.sh + +# # test with automodel dp=2 +# echo "run with automodel tp1 pp1 cp1 dp2 num_gpus2" +# BACKEND=automodel TP_SIZE=1 PP_SIZE=1 CP_SIZE=1 FSDP_SIZE=2 NUM_GPUS=2 bash tests/special_e2e/sft/run_sft_engine.sh + +# # test with automodel tp2 dp=2 +# echo "run with automodel tp2 pp1 cp1 dp2 num_gpus4" +# BACKEND=automodel TP_SIZE=2 PP_SIZE=1 CP_SIZE=1 FSDP_SIZE=2 NUM_GPUS=4 bash tests/special_e2e/sft/run_sft_engine.sh + +python3 tests/special_e2e/sft/compare_sft_engine_results.py + +rm -rf ~/verl/test/log diff --git a/verl/tests/special_npu/nightly_ci_ascend/run_dapo_moonlight-16b_megatron_npu.sh b/verl/tests/special_npu/nightly_ci_ascend/run_dapo_moonlight-16b_megatron_npu.sh new file mode 100644 index 0000000000000000000000000000000000000000..2055eb3d72e0818c6107e9d77dae154354ba62a5 --- /dev/null +++ b/verl/tests/special_npu/nightly_ci_ascend/run_dapo_moonlight-16b_megatron_npu.sh @@ -0,0 +1,191 @@ +set -x + +project_name='moonlight' +exp_name='exp' +adv_estimator=grpo +use_kl_in_reward=False +kl_penalty="kl" +kl_coef=0.0 +use_kl_loss=True +kl_loss_coef=0.001 + +clip_ratio_low=0.2 +clip_ratio_high=0.28 +max_prompt_length=$((1024 * 1)) +max_response_length=$((1024 * 2)) +enable_overlong_buffer=True +overlong_buffer_len=$((1024 * 1)) +overlong_penalty_factor=1.0 + +loss_agg_mode="token-mean" +train_prompt_bsz=32 +n_resp_per_prompt=16 +train_prompt_mini_bsz=32 +train_ppo_micro_batch_size_per_gpu=2 +infer_ppo_micro_batch_size_per_gpu=2 + +# Paths +MODEL_ID=${MODEL_ID:-moonshotai/Moonlight-16B-A3B-Instruct} +MODEL_PATH=${MODEL_PATH:-${HOME}/.cache/models/${MODEL_ID}} + +TRAIN_FILE=$HOME/.cache/datasets/dapo-math-17k.parquet +TEST_FILE=$HOME/.cache/datasets/dapo-math-17k.parquet + +# Algorithm +temperature=1.0 +top_p=1.0 +top_k=-1 # 0 for HF rollout, -1 for vLLM rollout +val_top_p=0.7 + +# Performance Related Parameter +use_dynamic_bsz=True +actor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 1)) +infer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3)) + +optimizer_offload_fraction=1 + +COMMON_PP=${COMMON_PP:-4} +COMMON_VPP=${COMMON_VPP:-null} +COMMON_CP=${COMMON_CP:-2} +COMMON_TP=${COMMON_TP:-1} +COMMON_EP=${COMMON_EP:-2} +COMMON_ETP=${COMMON_ETP:-1} + +TRAIN_TP=${TRAIN_TP:-$COMMON_TP} +INFER_TP=${INFER_TP:-4} +ACTOR_PP=${ACTOR_PP:-$COMMON_PP} +ACTOR_VPP=${ACTOR_VPP:-$COMMON_VPP} +ACTOR_CP=${ACTOR_CP:-$COMMON_CP} +ACTOR_TP=${ACTOR_TP:-$TRAIN_TP} +ACTOR_EP=${ACTOR_EP:-$COMMON_EP} +ACTOR_ETP=${ACTOR_ETP:-$COMMON_ETP} +ROLLOUT_TP=${ROLLOUT_TP:-$INFER_TP} +REF_PP=${REF_PP:-$COMMON_PP} +REF_VPP=${REF_VPP:-$COMMON_VPP} +REF_CP=${REF_CP:-$COMMON_CP} +REF_TP=${REF_TP:-$TRAIN_TP} +REF_EP=${REF_EP:-$COMMON_EP} +REF_ETP=${REF_ETP:-$COMMON_ETP} +CRITIC_PP=${CRITIC_PP:-$COMMON_PP} +CRITIC_VPP=${CRITIC_VPP:-$COMMON_VPP} +CRITIC_CP=${CRITIC_CP:-$COMMON_CP} +CRITIC_TP=${CRITIC_TP:-$TRAIN_TP} +CRITIC_EP=${CRITIC_EP:-$COMMON_EP} +CRITIC_ETP=${CRITIC_ETP:-$COMMON_ETP} +RM_PP=${RM_PP:-$COMMON_PP} +RM_VPP=${RM_VPP:-$COMMON_VPP} +RM_CP=${RM_CP:-$COMMON_CP} +RM_TP=${RM_TP:-$TRAIN_TP} +RM_EP=${RM_EP:-$COMMON_EP} +RM_ETP=${RM_ETP:-$COMMON_ETP} + +USE_MBRIDGE=True +USE_DIST_CKPT=False + +first_layer=7 +last_layer=6 +python3 -m recipe.dapo.main_dapo \ + --config-path=config \ + --config-name="dapo_megatron_trainer" \ + data.shuffle=False \ + data.validation_shuffle=False \ + data.train_files="${TRAIN_FILE}" \ + data.val_files="${TEST_FILE}" \ + data.prompt_key=prompt \ + data.truncation='left' \ + data.max_prompt_length=${max_prompt_length} \ + data.max_response_length=${max_response_length} \ + data.train_batch_size=${train_prompt_bsz} \ + data.trust_remote_code=True \ + actor_rollout_ref.rollout.n=${n_resp_per_prompt} \ + algorithm.adv_estimator=${adv_estimator} \ + algorithm.use_kl_in_reward=${use_kl_in_reward} \ + algorithm.kl_penalty=${kl_penalty} \ + algorithm.kl_ctrl.kl_coef=${kl_coef} \ + actor_rollout_ref.model.path="${MODEL_PATH}" \ + actor_rollout_ref.model.trust_remote_code=True \ + actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \ + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \ + actor_rollout_ref.actor.policy_loss.loss_mode=vanilla \ + algorithm.filter_groups.enable=False \ + actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \ + actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \ + actor_rollout_ref.actor.clip_ratio_c=10.0 \ + actor_rollout_ref.model.use_fused_kernels=False \ + actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \ + actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${train_ppo_micro_batch_size_per_gpu} \ + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \ + actor_rollout_ref.actor.optim.lr=3e-6 \ + +actor_rollout_ref.actor.megatron.override_transformer_config.num_layers_in_first_pipeline_stage=$first_layer \ + +actor_rollout_ref.actor.megatron.override_transformer_config.num_layers_in_last_pipeline_stage=$last_layer \ + +actor_rollout_ref.actor.megatron.override_transformer_config.tensor_model_parallel_size=${ACTOR_TP} \ + +actor_rollout_ref.actor.megatron.override_transformer_config.multi_latent_attention=True \ + +actor_rollout_ref.actor.megatron.override_transformer_config.use_flash_attn=True \ + actor_rollout_ref.actor.megatron.param_offload=True \ + actor_rollout_ref.actor.megatron.optimizer_offload=True \ + actor_rollout_ref.actor.megatron.grad_offload=True \ + actor_rollout_ref.ref.megatron.param_offload=True \ + ++actor_rollout_ref.actor.megatron.override_transformer_config.attention_backend=fused \ + actor_rollout_ref.actor.megatron.use_mbridge=$USE_MBRIDGE \ + actor_rollout_ref.actor.megatron.use_dist_checkpointing=$USE_DIST_CKPT \ + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${ACTOR_TP} \ + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${ACTOR_PP} \ + actor_rollout_ref.actor.megatron.virtual_pipeline_model_parallel_size=${ACTOR_VPP} \ + actor_rollout_ref.actor.megatron.context_parallel_size=${ACTOR_CP} \ + +actor_rollout_ref.actor.megatron.override_transformer_config.context_parallel_size=${ACTOR_CP} \ + actor_rollout_ref.actor.megatron.expert_model_parallel_size=${ACTOR_EP} \ + actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${ACTOR_ETP} \ + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform \ + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full \ + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 \ + actor_rollout_ref.actor.entropy_coeff=0 \ + actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \ + actor_rollout_ref.rollout.load_format=safetensors \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${infer_ppo_micro_batch_size_per_gpu} \ + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.65 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=${INFER_TP} \ + actor_rollout_ref.rollout.data_parallel_size=1 \ + actor_rollout_ref.rollout.enable_chunked_prefill=False \ + actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \ + actor_rollout_ref.rollout.temperature=${temperature} \ + actor_rollout_ref.rollout.top_p=${top_p} \ + actor_rollout_ref.rollout.top_k=${top_k} \ + actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \ + actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \ + actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \ + actor_rollout_ref.rollout.val_kwargs.do_sample=True \ + actor_rollout_ref.rollout.val_kwargs.n=1 \ + actor_rollout_ref.rollout.name=vllm \ + actor_rollout_ref.rollout.calculate_log_probs=True \ + actor_rollout_ref.rollout.enforce_eager=True \ + actor_rollout_ref.rollout.free_cache_engine=True \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${infer_ppo_micro_batch_size_per_gpu} \ + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \ + actor_rollout_ref.ref.megatron.use_dist_checkpointing=${USE_DIST_CKPT} \ + actor_rollout_ref.ref.megatron.param_offload=True \ + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${REF_TP} \ + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${REF_PP} \ + actor_rollout_ref.ref.megatron.virtual_pipeline_model_parallel_size=${REF_VPP} \ + actor_rollout_ref.ref.megatron.context_parallel_size=${REF_CP} \ + actor_rollout_ref.ref.megatron.expert_model_parallel_size=${REF_EP} \ + actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=${REF_ETP} \ + reward_model.reward_manager=dapo \ + +reward_model.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \ + +reward_model.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \ + +reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \ + +reward_model.reward_kwargs.overlong_buffer_cfg.log=False \ + +reward_model.reward_kwargs.max_resp_len=${max_response_length} \ + trainer.logger=['console'] \ + trainer.project_name="${project_name}" \ + trainer.experiment_name="${exp_name}" \ + trainer.n_gpus_per_node=8 \ + trainer.nnodes=1 \ + trainer.val_before_train=False \ + trainer.balance_batch=False \ + trainer.test_freq=-1 \ + trainer.save_freq=-1 \ + trainer.resume_mode=auto \ + trainer.log_val_generations=10 \ + trainer.total_training_steps=15 2>&1 | tee /root/.cache/nightly_log/moonlight/dapo_moonlight16b_megatron_npu-$(date +%Y%m%d_%H%M).log diff --git a/verl/tests/special_npu/nightly_ci_ascend/run_grpo_qwen25-7b-instruct_fsdp_npu.sh b/verl/tests/special_npu/nightly_ci_ascend/run_grpo_qwen25-7b-instruct_fsdp_npu.sh new file mode 100644 index 0000000000000000000000000000000000000000..c379c77bea996d4499787e3a803be84812841580 --- /dev/null +++ b/verl/tests/special_npu/nightly_ci_ascend/run_grpo_qwen25-7b-instruct_fsdp_npu.sh @@ -0,0 +1,49 @@ +set -x + +# Some models are optimized by vllm ascend. While in some case, e.g. rlhf training, +# the optimized model may not be suitable. In this case, set this value to 0 to disable the optimized model. + +MODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-7B-Instruct} +MODEL_PATH=${MODEL_PATH:-${HOME}/.cache/models/${MODEL_ID}} + +python3 -m verl.trainer.main_ppo \ + algorithm.adv_estimator=grpo \ + data.train_files=$HOME/data/gsm8k/train.parquet \ + data.val_files=$HOME/data/gsm8k/test.parquet \ + data.train_batch_size=32 \ + data.max_prompt_length=1024 \ + data.max_response_length=1024 \ + data.filter_overlong_prompts=True \ + data.truncation='error' \ + actor_rollout_ref.model.path="${MODEL_PATH}" \ + actor_rollout_ref.actor.optim.lr=5e-8 \ + actor_rollout_ref.model.use_remove_padding=False \ + actor_rollout_ref.actor.ppo_mini_batch_size=32 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \ + actor_rollout_ref.actor.use_kl_loss=True \ + actor_rollout_ref.actor.entropy_coeff=0 \ + actor_rollout_ref.actor.kl_loss_coef=0.001 \ + actor_rollout_ref.actor.kl_loss_type=low_var_kl \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.fsdp_config.param_offload=False \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=2 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=4 \ + actor_rollout_ref.rollout.name=vllm \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \ + actor_rollout_ref.rollout.n=5 \ + actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=4096 \ + actor_rollout_ref.rollout.enable_chunked_prefill=False \ + actor_rollout_ref.rollout.calculate_log_probs=True \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=2 \ + actor_rollout_ref.ref.fsdp_config.param_offload=True \ + algorithm.use_kl_in_reward=False \ + trainer.critic_warmup=0 \ + trainer.logger=console \ + trainer.project_name='verl_grpo_example_gsm8k' \ + trainer.experiment_name='qwen2_5_7b_instruct_fsdp' \ + trainer.n_gpus_per_node=8 \ + trainer.nnodes=1 \ + trainer.save_freq=-1 \ + trainer.test_freq=-1 \ + trainer.total_training_steps=15 2>&1 | tee /root/.cache/nightly_log/qwen25-7b/grpo_qwen25-7b-instruct_fsdp_npu-$(date +%Y%m%d_%H%M).log \ No newline at end of file diff --git a/verl/tests/special_npu/nightly_ci_ascend/run_grpo_qwen25-vl-3b-instruct_fsdp_npu.sh b/verl/tests/special_npu/nightly_ci_ascend/run_grpo_qwen25-vl-3b-instruct_fsdp_npu.sh new file mode 100644 index 0000000000000000000000000000000000000000..fc47b67daf2b064bf470e7d0b27eac58cde3abf7 --- /dev/null +++ b/verl/tests/special_npu/nightly_ci_ascend/run_grpo_qwen25-vl-3b-instruct_fsdp_npu.sh @@ -0,0 +1,54 @@ +set -x +ENGINE=${1:-vllm} + +# Some models are optimized by vllm ascend. While in some case, e.g. rlhf training, +# the optimized model may not be suitable. In this case, set this value to 0 to disable the optimized model. +export USE_OPTIMIZED_MODEL=0 + +MODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-VL-3B-Instruct} +MODEL_PATH=${MODEL_PATH:-${HOME}/.cache/models/${MODEL_ID}} + +python3 -m verl.trainer.main_ppo \ + algorithm.adv_estimator=grpo \ + data.train_files=$HOME/data/geo3k/train.parquet \ + data.val_files=$HOME/data/geo3k/test.parquet \ + data.train_batch_size=16 \ + data.max_prompt_length=1024 \ + data.max_response_length=2048 \ + data.filter_overlong_prompts=True \ + data.truncation='error' \ + data.image_key=images \ + actor_rollout_ref.model.path="${MODEL_PATH}" \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.model.use_remove_padding=True \ + actor_rollout_ref.actor.ppo_mini_batch_size=16 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \ + actor_rollout_ref.actor.use_kl_loss=True \ + actor_rollout_ref.actor.kl_loss_coef=0.01 \ + actor_rollout_ref.actor.kl_loss_type=low_var_kl \ + actor_rollout_ref.actor.entropy_coeff=0 \ + actor_rollout_ref.actor.use_torch_compile=False \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.fsdp_config.param_offload=False \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=2 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=2 \ + actor_rollout_ref.rollout.name=$ENGINE \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \ + actor_rollout_ref.rollout.enable_chunked_prefill=False \ + actor_rollout_ref.rollout.enforce_eager=True \ + actor_rollout_ref.rollout.free_cache_engine=True \ + actor_rollout_ref.rollout.n=5 \ + actor_rollout_ref.rollout.calculate_log_probs=True \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=2 \ + actor_rollout_ref.ref.fsdp_config.param_offload=True \ + algorithm.use_kl_in_reward=False \ + trainer.critic_warmup=0 \ + trainer.logger=console \ + trainer.project_name='verl_grpo_example_geo3k' \ + trainer.experiment_name='qwen2_5_vl_3b_function_rm' \ + trainer.n_gpus_per_node=8 \ + trainer.nnodes=1 \ + trainer.save_freq=-1 \ + trainer.test_freq=-1 \ + trainer.total_training_steps=15 2>&1 | tee /root/.cache/nightly_log/qwen25-vl-3b/grpo_qwen25-vl-3b-instruct_fsdp_npu-$(date +%Y%m%d_%H%M).log \ No newline at end of file diff --git a/verl/tests/special_npu/nightly_ci_ascend/run_gspo_qwen3_30b_megatron_npu.sh b/verl/tests/special_npu/nightly_ci_ascend/run_gspo_qwen3_30b_megatron_npu.sh new file mode 100644 index 0000000000000000000000000000000000000000..c0e479890b405d1d01e07abb1cd7dd38d8c4107d --- /dev/null +++ b/verl/tests/special_npu/nightly_ci_ascend/run_gspo_qwen3_30b_megatron_npu.sh @@ -0,0 +1,145 @@ +#!/usr/bin/env bash +set -x + +export HCCL_OP_EXPANSION_MODE="AIV" +export VLLM_USE_V1=1 + +# ---- user-adjustable ---- +MODEL_ID=${MODEL_ID:-Qwen/Qwen3-30B-A3B-Instruct-2507} +MODEL_PATH=${MODEL_PATH:-${HOME}/.cache/models/${MODEL_ID}} + +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-16} + +train_batch_size=${TRAIN_BATCH_SIZE:-32} +ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-16} +max_prompt_length=${MAX_PROMPT_LENGTH:-2048} +max_response_length=${MAX_RESPONSE_LENGTH:-2048} +ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-30720} + +actor_lr=${ACTOR_LR:-1e-6} +entropy_coeff=${ENTROPY_COEFF:-0} + +clip_ratio_low=${CLIP_RATIO_LOW:-3e-4} +clip_ratio_high=${CLIP_RATIO_HIGH:-4e-4} + +actor_tp=${ACTOR_TP:-2} +actor_pp=${ACTOR_PP:-1} +actor_ep=${ACTOR_EP:-8} +actor_etp=${ACTOR_ETP:-1} + +rollout_tp=${ROLLOUT_TP:-8} +rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.7} +rollout_n=${ROLLOUT_N:-4} + +total_epochs=${TOTAL_EPOCHS:-10} +save_freq=${SAVE_FREQ:--1} +test_freq=${TEST_FREQ:--1} + +project_name=${PROJECT_NAME:-verl_gspo_qwen3_moe} +experiment_name=${EXPERIMENT_NAME:-qwen3_30b_a3b_vllm_megatron} + +train_file=$HOME/data/gsm8k/train.parquet +val_file=$HOME/data/gsm8k/test.parquet + +# ---- end user-adjustable ---- +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files="['$train_file']" + data.val_files="['$val_file']" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.filter_overlong_prompts=True + data.truncation='error' +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True +) + +ACTOR=( + actor_rollout_ref.actor.policy_loss.loss_mode=gspo + actor_rollout_ref.actor.loss_agg_mode=seq-mean-token-mean + actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} + actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} + actor_rollout_ref.actor.clip_ratio_c=10.0 + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=False + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${actor_pp} + actor_rollout_ref.actor.megatron.expert_model_parallel_size=${actor_ep} + actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${actor_etp} + actor_rollout_ref.actor.megatron.param_offload=True + actor_rollout_ref.actor.megatron.grad_offload=True + actor_rollout_ref.actor.megatron.optimizer_offload=True + actor_rollout_ref.actor.megatron.use_mbridge=True + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32 + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=True + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 + +actor_rollout_ref.actor.megatron.override_transformer_config.apply_rope_fusion=True +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.rollout.val_kwargs.n=1 + actor_rollout_ref.rollout.val_kwargs.temperature=1.0 + actor_rollout_ref.rollout.val_kwargs.top_p=0.7 + actor_rollout_ref.rollout.calculate_log_probs=True +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${actor_pp} + actor_rollout_ref.ref.megatron.expert_model_parallel_size=${actor_ep} + actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=${actor_etp} + actor_rollout_ref.ref.megatron.param_offload=True + actor_rollout_ref.ref.megatron.use_mbridge=True +) + +TRAINER=( + trainer.balance_batch=True + trainer.critic_warmup=0 + trainer.logger='["console"]' + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.val_before_train=False + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} + trainer.total_training_steps=15 +) + +EXTRA=( + model_engine=megatron +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/verl/tests/special_npu/nightly_ci_ascend/run_ppo_qwen3-8b_fsdp_npu.sh b/verl/tests/special_npu/nightly_ci_ascend/run_ppo_qwen3-8b_fsdp_npu.sh new file mode 100644 index 0000000000000000000000000000000000000000..16df63ecc281e0649518567759dce94d6ff7d9dd --- /dev/null +++ b/verl/tests/special_npu/nightly_ci_ascend/run_ppo_qwen3-8b_fsdp_npu.sh @@ -0,0 +1,60 @@ +set -x + +# Some models are optimized by vllm ascend. While in some case, e.g. rlhf training, +# the optimized model may not be suitable. In this case, set this value to 0 to disable the optimized model. + +MODEL_ID=${MODEL_ID:-Qwen/Qwen3-8B} +MODEL_PATH=${MODEL_PATH:-${HOME}/.cache/models/${MODEL_ID}} + +python3 -m verl.trainer.main_ppo \ + algorithm.adv_estimator=gae \ + data.train_files=$HOME/data/gsm8k/train.parquet \ + data.val_files=$HOME/data/gsm8k/test.parquet \ + data.train_batch_size=32 \ + data.max_prompt_length=2000 \ + data.max_response_length=2000 \ + data.shuffle=False \ + actor_rollout_ref.model.path="${MODEL_PATH}" \ + actor_rollout_ref.model.use_remove_padding=True \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.actor.ppo_mini_batch_size=32 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.actor.fsdp_config.param_offload=True \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \ + actor_rollout_ref.actor.use_kl_loss=False \ + actor_rollout_ref.actor.ulysses_sequence_parallel_size=2 \ + actor_rollout_ref.actor.use_dynamic_bsz=True \ + actor_rollout_ref.actor.use_torch_compile=False \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=1 \ + actor_rollout_ref.rollout.name=vllm \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.9 \ + actor_rollout_ref.rollout.max_num_batched_tokens=4000 \ + actor_rollout_ref.rollout.max_num_seqs=64 \ + actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=4096 \ + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \ + actor_rollout_ref.rollout.enable_chunked_prefill=True \ + actor_rollout_ref.rollout.enforce_eager=False \ + actor_rollout_ref.rollout.calculate_log_probs=True \ + critic.optim.lr=1e-5 \ + critic.model.use_remove_padding=True \ + critic.model.path="${MODEL_PATH}" \ + critic.model.enable_gradient_checkpointing=True \ + critic.ppo_micro_batch_size_per_gpu=1 \ + critic.ulysses_sequence_parallel_size=2 \ + critic.fsdp.param_offload=True \ + critic.fsdp.optimizer_offload=True \ + critic.use_dynamic_bsz=True \ + trainer.critic_warmup=0 \ + trainer.logger=console \ + trainer.project_name='verl_example_ppo_gsm8k' \ + trainer.experiment_name='qwen3_8b_fsdp' \ + trainer.n_gpus_per_node=8 \ + trainer.nnodes=1 \ + trainer.save_freq=-1 \ + trainer.test_freq=-1 \ + trainer.val_before_train=False \ + trainer.max_actor_ckpt_to_keep=1 \ + trainer.max_critic_ckpt_to_keep=1 \ + trainer.total_training_steps=15 2>&1 | tee /root/.cache/nightly_log/qwen3-8b-ppo/ppo_qwen3-8b_fsdp_npu-$(date +%Y%m%d_%H%M).log \ No newline at end of file diff --git a/verl/tests/special_npu/run_qwen2_5_05b_grpo.sh b/verl/tests/special_npu/run_qwen2_5_05b_grpo.sh new file mode 100644 index 0000000000000000000000000000000000000000..c6e86dfdbf486737cce49be60a7b0f22b260fb69 --- /dev/null +++ b/verl/tests/special_npu/run_qwen2_5_05b_grpo.sh @@ -0,0 +1,78 @@ +set -x + +MODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-0.5B-Instruct} +MODEL_PATH=${MODEL_PATH:-${HOME}/.cache/models/${MODEL_ID}} + +SAVE_PATH=tests/utils/ci/profiler_data +rm -rf "$SAVE_PATH" + +LEVEL="level0" +CONTENTS=['npu','cpu'] +ANALYSIS=False +PROFILE_STEPS=[1] +PROFILE_RANKS_ALL=False +PROFILE_RANKS=[0] +DISCRETE=True + +python3 -m verl.trainer.main_ppo \ + algorithm.adv_estimator=grpo \ + data.train_files=$HOME/data/gsm8k/train.parquet \ + data.val_files=$HOME/data/gsm8k/test.parquet \ + data.train_batch_size=16 \ + data.max_prompt_length=512 \ + data.max_response_length=128 \ + data.filter_overlong_prompts=True \ + data.truncation='error' \ + actor_rollout_ref.model.path="${MODEL_PATH}" \ + actor_rollout_ref.actor.optim.lr=5e-7 \ + actor_rollout_ref.model.use_remove_padding=False \ + actor_rollout_ref.actor.ppo_mini_batch_size=8 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.actor.use_kl_loss=True \ + actor_rollout_ref.actor.kl_loss_coef=0.001 \ + actor_rollout_ref.actor.kl_loss_type=low_var_kl \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.fsdp_config.param_offload=False \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \ + actor_rollout_ref.actor.use_torch_compile=False \ + actor_rollout_ref.ref.use_torch_compile=False \ + +actor_rollout_ref.rollout.engine_kwargs.vllm.compilation_config.cudagraph_mode="FULL_AND_PIECEWISE" \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.rollout.enable_chunked_prefill=False \ + actor_rollout_ref.rollout.tensor_model_parallel_size=2 \ + actor_rollout_ref.rollout.name=vllm \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \ + actor_rollout_ref.rollout.n=2 \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.ref.fsdp_config.param_offload=True \ + algorithm.kl_ctrl.kl_coef=0.001 \ + trainer.critic_warmup=0 \ + trainer.logger=console \ + trainer.project_name='verl_grpo_example_gsm8k' \ + trainer.experiment_name='qwen2_7b_function_rm' \ + trainer.n_gpus_per_node=8 \ + trainer.nnodes=1 \ + trainer.save_freq=-1 \ + trainer.test_freq=-1 \ + trainer.total_epochs=1 \ + trainer.total_training_steps=1 \ + actor_rollout_ref.actor.profiler.enable=True \ + actor_rollout_ref.actor.profiler.all_ranks=$PROFILE_RANKS_ALL \ + actor_rollout_ref.actor.profiler.ranks=$PROFILE_RANKS \ + actor_rollout_ref.actor.profiler.tool_config.npu.discrete=$DISCRETE \ + actor_rollout_ref.actor.profiler.tool_config.npu.contents=$CONTENTS \ + actor_rollout_ref.actor.profiler.tool_config.npu.level=$LEVEL \ + actor_rollout_ref.actor.profiler.tool_config.npu.analysis=$ANALYSIS \ + actor_rollout_ref.ref.profiler.enable=True \ + actor_rollout_ref.ref.profiler.all_ranks=$PROFILE_RANKS_ALL \ + actor_rollout_ref.ref.profiler.ranks=$PROFILE_RANKS \ + actor_rollout_ref.ref.profiler.tool_config.npu.discrete=$DISCRETE \ + actor_rollout_ref.ref.profiler.tool_config.npu.contents=$CONTENTS \ + actor_rollout_ref.ref.profiler.tool_config.npu.level=$LEVEL \ + actor_rollout_ref.ref.profiler.tool_config.npu.analysis=$ANALYSIS \ + global_profiler.tool=npu \ + global_profiler.steps=$PROFILE_STEPS \ + global_profiler.save_path="$SAVE_PATH" $@ + +python3 "tests/utils/test_check_profiler_output.py" --profiler_dir="$SAVE_PATH" --device="npu" +rm -rf "$SAVE_PATH" diff --git a/verl/tests/special_npu/run_qwen2_5_05b_grpo_mindspeed.sh b/verl/tests/special_npu/run_qwen2_5_05b_grpo_mindspeed.sh new file mode 100644 index 0000000000000000000000000000000000000000..b57acac1dfabc883ff7b75c2d5d80d7446527584 --- /dev/null +++ b/verl/tests/special_npu/run_qwen2_5_05b_grpo_mindspeed.sh @@ -0,0 +1,68 @@ +set -x + +MODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-0.5B-Instruct} +MODEL_PATH=${MODEL_PATH:-${HOME}/.cache/models/${MODEL_ID}} + +USE_DIST_CKPT=${USE_DIST_CKPT:-False} +DIST_CKPT_PATH=${DIST_CKPT_PATH:-${HOME}/dist_ckpt/qwen2_5_05b_grpo_mindspeed} +if [ "$USE_DIST_CKPT" = "True" ]; then + if [ "$USE_DUMMY_MODEL" = "True" ]; then + DIST_CKPT_PATH=${HOME}/dist_ckpt_dummy/${MODEL_ID} + fi + python scripts/converter_hf_to_mcore.py \ + --hf_model_path "${MODEL_PATH}" \ + --output_path "${DIST_CKPT_PATH}" +fi + + +python3 -m verl.trainer.main_ppo --config-path=config \ + --config-name='ppo_megatron_trainer.yaml' \ + algorithm.adv_estimator=grpo \ + data.train_files=$HOME/data/gsm8k/train.parquet \ + data.val_files=$HOME/data/gsm8k/test.parquet \ + data.train_batch_size=16 \ + data.max_prompt_length=512 \ + data.max_response_length=128 \ + data.filter_overlong_prompts=True \ + data.truncation='error' \ + actor_rollout_ref.model.path=${MODEL_PATH} \ + actor_rollout_ref.actor.optim.lr=5e-7 \ + actor_rollout_ref.actor.ppo_mini_batch_size=8 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.actor.strategy=megatron \ + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=2 \ + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=2 \ + actor_rollout_ref.actor.megatron.expert_model_parallel_size=1 \ + actor_rollout_ref.actor.megatron.use_dist_checkpointing=True \ + actor_rollout_ref.actor.megatron.dist_checkpointing_path=${DIST_CKPT_PATH} \ + actor_rollout_ref.actor.use_kl_loss=True \ + actor_rollout_ref.actor.kl_loss_coef=0.001 \ + actor_rollout_ref.actor.kl_loss_type=low_var_kl \ + actor_rollout_ref.actor.use_torch_compile=False \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.rollout.enable_chunked_prefill=False \ + actor_rollout_ref.rollout.tensor_model_parallel_size=2 \ + actor_rollout_ref.rollout.name=vllm \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \ + +actor_rollout_ref.rollout.engine_kwargs.vllm.compilation_config.cudagraph_mode="FULL_AND_PIECEWISE" \ + actor_rollout_ref.rollout.n=2 \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.ref.strategy=megatron \ + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=2 \ + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=2 \ + actor_rollout_ref.ref.megatron.expert_model_parallel_size=1 \ + actor_rollout_ref.ref.megatron.use_dist_checkpointing=True \ + actor_rollout_ref.ref.megatron.dist_checkpointing_path=${DIST_CKPT_PATH} \ + actor_rollout_ref.ref.use_torch_compile=False \ + algorithm.kl_ctrl.kl_coef=0.001 \ + trainer.critic_warmup=0 \ + trainer.logger=console \ + trainer.project_name='verl_grpo_example_gsm8k' \ + trainer.experiment_name='qwen2_7b_function_rm' \ + trainer.n_gpus_per_node=8 \ + trainer.nnodes=1 \ + trainer.save_freq=-1 \ + trainer.test_freq=-1 \ + trainer.total_epochs=1 \ + trainer.total_training_steps=1 \ + +actor_rollout_ref.actor.megatron.override_transformer_config.use_flash_attn=True $@ diff --git a/verl/tests/special_npu/run_qwen2_5_05b_sft_peft_sp2.sh b/verl/tests/special_npu/run_qwen2_5_05b_sft_peft_sp2.sh new file mode 100644 index 0000000000000000000000000000000000000000..ba61ca10d0fab142fbfb7a6b2f1e0e900607c361 --- /dev/null +++ b/verl/tests/special_npu/run_qwen2_5_05b_sft_peft_sp2.sh @@ -0,0 +1,65 @@ +set -x + +NUM_GPUS=${NUM_GPUS:-4} + +mode=${mode:-spmd} + +if [ "$mode" = "spmd" ]; then + ENTRYPOINT=${ENTRYPOINT:-"-m verl.trainer.sft_trainer"} + COMMAND="torchrun --standalone --nnodes=${NNODES:-1} --nproc-per-node=${NUM_GPUS:-1} ${ENTRYPOINT}" +else + ENTRYPOINT=${ENTRYPOINT:-"-m verl.trainer.sft_trainer_ray"} + COMMAND="python ${ENTRYPOINT} trainer.nnodes=${NNODES:-1} trainer.n_gpus_per_node=${NUM_GPUS:-1}" +fi + +RESUME_MODE=disable + +ckpts_home=${ckpts_home:-~/verl/test/gsm8k-sft-fsdp} + +MODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-0.5B-Instruct} +MODEL_PATH=${MODEL_PATH:-${HOME}/.cache/models/${MODEL_ID}} + +DATASET_DIR=${DATASET_DIR:-$HOME/data/gsm8k_sft} +TRAIN_FILES=${DATASET_DIR}/train.parquet +VAL_FILES=${DATASET_DIR}/test.parquet + +exp_name=gsm8k-sft-qwen-2.5-0.5b-instruct-mode-${mode} + +mkdir -p "${ckpts_home}" + +$COMMAND \ + data.train_files=$TRAIN_FILES \ + data.val_files=$VAL_FILES \ + data.pad_mode=no_padding \ + data.truncation=error \ + data.use_dynamic_bsz=True \ + data.max_token_len_per_gpu=2048 \ + data.messages_key=messages \ + model.path=$MODEL_PATH \ + model.use_remove_padding=True \ + model.lora_rank=32 \ + model.lora_alpha=16 \ + model.target_modules=all-linear \ + engine=fsdp \ + optim=fsdp \ + optim.lr=1e-5 \ + optim.lr_warmup_steps_ratio=0.2 \ + optim.weight_decay=0.1 \ + optim.betas="[0.9,0.95]" \ + optim.clip_grad=1.0 \ + optim.min_lr_ratio=0.1 \ + optim.lr_scheduler_type=cosine \ + engine.ulysses_sequence_parallel_size=2 \ + engine.strategy=fsdp2 \ + engine.fsdp_size=2 \ + trainer.test_freq=after_each_epoch \ + trainer.save_freq=-1 \ + trainer.logger=['console','file'] \ + trainer.project_name=gsm8k-sft \ + trainer.experiment_name=gsm8k-sft-qwen-2.5-0.5b-instruct \ + trainer.total_epochs=2 \ + trainer.total_training_steps=2 \ + trainer.default_local_dir="${ckpts_home}" \ + trainer.resume_mode=${RESUME_MODE} \ + +rm -rf "${ckpts_home:?}/*" diff --git a/verl/tests/special_npu/run_qwen2_5_vl_3b_npu.sh b/verl/tests/special_npu/run_qwen2_5_vl_3b_npu.sh new file mode 100644 index 0000000000000000000000000000000000000000..a66c2f77de4e171f309dadec4910c23e7f1a171d --- /dev/null +++ b/verl/tests/special_npu/run_qwen2_5_vl_3b_npu.sh @@ -0,0 +1,57 @@ +set -x + +ENGINE=${1:-vllm} + +# Some models are optimized by vllm ascend. While in some case, e.g. rlhf training, +# the optimized model may not be suitable. In this case, set this value to 0 to disable the optimized model. +export USE_OPTIMIZED_MODEL=0 + +MODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-VL-3B-Instruct} +MODEL_PATH=${MODEL_PATH:-${HOME}/.cache/models/${MODEL_ID}} + +python3 -m verl.trainer.main_ppo \ + algorithm.adv_estimator=grpo \ + data.train_files=$HOME/data/geo3k/train.parquet \ + data.val_files=$HOME/data/geo3k/test.parquet \ + data.train_batch_size=16 \ + data.max_prompt_length=1024 \ + data.max_response_length=2048 \ + data.filter_overlong_prompts=True \ + data.truncation='error' \ + data.image_key=images \ + actor_rollout_ref.model.path="${MODEL_PATH}" \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.model.use_remove_padding=True \ + actor_rollout_ref.actor.ppo_mini_batch_size=8 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.actor.use_kl_loss=True \ + actor_rollout_ref.actor.kl_loss_coef=0.01 \ + actor_rollout_ref.actor.kl_loss_type=low_var_kl \ + actor_rollout_ref.actor.entropy_coeff=0 \ + actor_rollout_ref.actor.use_torch_compile=False \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.fsdp_config.param_offload=False \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \ + actor_rollout_ref.ref.use_torch_compile=False \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=2 \ + actor_rollout_ref.rollout.name=$ENGINE \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \ + actor_rollout_ref.rollout.enable_chunked_prefill=False \ + actor_rollout_ref.rollout.enforce_eager=True \ + +actor_rollout_ref.rollout.engine_kwargs.vllm.compilation_config.cudagraph_mode="FULL_AND_PIECEWISE" \ + actor_rollout_ref.rollout.free_cache_engine=True \ + actor_rollout_ref.rollout.n=2 \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.ref.fsdp_config.param_offload=True \ + algorithm.use_kl_in_reward=False \ + trainer.critic_warmup=0 \ + trainer.logger=console \ + trainer.project_name='verl_grpo_example_geo3k' \ + trainer.experiment_name='qwen2_5_vl_3b_function_rm' \ + trainer.n_gpus_per_node=8 \ + trainer.nnodes=1 \ + trainer.save_freq=-1 \ + trainer.test_freq=-1 \ + trainer.total_epochs=1 \ + trainer.total_training_steps=1 $@ \ No newline at end of file diff --git a/verl/tests/special_npu/run_qwen3_06b_ppo.sh b/verl/tests/special_npu/run_qwen3_06b_ppo.sh new file mode 100644 index 0000000000000000000000000000000000000000..d3844414db5fc018261dd2b06123936686295d53 --- /dev/null +++ b/verl/tests/special_npu/run_qwen3_06b_ppo.sh @@ -0,0 +1,52 @@ +set -x + +MODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-0.5B-Instruct} # TODO: change to Qwen3-0.6B when CI server is ready +MODEL_PATH=${MODEL_PATH:-${HOME}/.cache/models/${MODEL_ID}} + +python3 -m verl.trainer.main_ppo \ + algorithm.adv_estimator=gae \ + data.train_files=$HOME/data/gsm8k/train.parquet \ + data.val_files=$HOME/data/gsm8k/test.parquet \ + data.train_batch_size=16 \ + data.max_prompt_length=512 \ + data.max_response_length=128 \ + data.shuffle=False \ + actor_rollout_ref.model.path="${MODEL_PATH}" \ + actor_rollout_ref.model.use_remove_padding=True \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.actor.ppo_mini_batch_size=8 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.actor.fsdp_config.param_offload=True \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \ + actor_rollout_ref.actor.use_kl_loss=False \ + actor_rollout_ref.actor.ulysses_sequence_parallel_size=2 \ + actor_rollout_ref.actor.use_dynamic_bsz=True \ + actor_rollout_ref.actor.use_torch_compile=False \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=1 \ + actor_rollout_ref.rollout.name=vllm \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \ + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \ + actor_rollout_ref.rollout.enable_chunked_prefill=True \ + actor_rollout_ref.rollout.enforce_eager=False \ + +actor_rollout_ref.rollout.engine_kwargs.vllm.compilation_config.cudagraph_mode="FULL_AND_PIECEWISE" \ + critic.optim.lr=1e-5 \ + critic.model.use_remove_padding=True \ + critic.model.path="${MODEL_PATH}" \ + critic.model.enable_gradient_checkpointing=True \ + critic.ppo_micro_batch_size_per_gpu=1 \ + critic.ulysses_sequence_parallel_size=2 \ + critic.fsdp.param_offload=True \ + critic.fsdp.optimizer_offload=True \ + critic.use_dynamic_bsz=True \ + trainer.critic_warmup=0 \ + trainer.logger='["console"]' \ + trainer.project_name='verl_ppo_example_gsm8k_qwen3' \ + trainer.experiment_name='qwen3_06b_fsdp' \ + trainer.n_gpus_per_node=8 \ + trainer.nnodes=1 \ + trainer.save_freq=-1 \ + trainer.test_freq=-1 \ + trainer.total_epochs=1 \ + trainer.total_training_steps=1 $@ diff --git a/verl/tests/special_npu/run_qwen3_30b_grpo_mindspeed.sh b/verl/tests/special_npu/run_qwen3_30b_grpo_mindspeed.sh new file mode 100644 index 0000000000000000000000000000000000000000..89933caa844c145d4b8e8fb1abcdb05c952edf60 --- /dev/null +++ b/verl/tests/special_npu/run_qwen3_30b_grpo_mindspeed.sh @@ -0,0 +1,134 @@ +#!/usr/bin/env bash +set -xeuo pipefail + + +MODEL_ID=${MODEL_ID:-Qwen/Qwen3-30B-A3B-Instruct-2507} +MODEL_PATH=${MODEL_PATH:-${HOME}/.cache/models/${MODEL_ID}} +USE_DIST_CKPT=${USE_DIST_CKPT:-False} +DIST_CKPT_PATH=${DIST_CKPT_PATH:-${HOME}/dist_ckpt/qwen3_30b_grpo_mindspeed} + +# use dummy model +if [[ "$USE_DUMMY_MODEL" == "True" ]]; then + DUMMY_MODEL_PATH=${DUMMY_MODEL_PATH:-${HOME}/models_dummy/${MODEL_ID}} + if [ -z "${DUMMY_MODEL_CONFIG_PATH}" ]; then + echo "[ERROR] DUMMY_MODEL_CONFIG_PATH not set" + exit 1 + fi + + # make sure the path is empty + if [[ -d $DUMMY_MODEL_PATH && $DUMMY_MODEL_PATH != "/" ]]; then + rm -rf $DUMMY_MODEL_PATH + fi + + # init model + python scripts/init_random_model.py \ + --hf_model_path "${MODEL_PATH}" \ + --new_config_path "${DUMMY_MODEL_CONFIG_PATH}" \ + --output_path "${DUMMY_MODEL_PATH}" + + # replace model path + MODEL_PATH=$DUMMY_MODEL_PATH +fi + +# convert to megatron +if [[ "$USE_DIST_CKPT" == "True" ]]; then + + if [[ "$USE_DUMMY_MODEL" == "True" ]]; then + DIST_CKPT_PATH=${HOME}/dist_ckpt/qwen3_30b_grpo_mindspeed_dummy + + if [[ -d $DIST_CKPT_PATH && $DIST_CKPT_PATH != "/" ]];then + rm -rf $DIST_CKPT_PATH + fi + fi + + torchrun --nproc_per_node 2 --nnodes 1 scripts/converter_hf_to_mcore.py \ + --hf_model_path "${MODEL_PATH}" \ + --output_path "${DIST_CKPT_PATH}" +fi + +exp_name='Qwen3-30B-A3B-GRPO-MindSpeed' + +max_prompt_length=512 +max_response_length=1024 + +train_prompt_bsz=16 + +actor_ppo_max_token_len=$(((max_prompt_length + max_response_length))) +infer_ppo_max_token_len=$(((max_prompt_length + max_response_length))) + +python3 -m verl.trainer.main_ppo --config-path=config \ + --config-name='ppo_megatron_trainer.yaml' \ + data.train_files=${HOME}/data/gsm8k/train.parquet \ + data.val_files=${HOME}/data/gsm8k/test.parquet \ + data.train_batch_size=${train_prompt_bsz} \ + data.max_prompt_length=${max_prompt_length} \ + data.max_response_length=${max_response_length} \ + data.filter_overlong_prompts=True \ + data.shuffle=False \ + data.truncation='left' \ + algorithm.adv_estimator=grpo \ + algorithm.use_kl_in_reward=False \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.rollout.enable_chunked_prefill=False \ + actor_rollout_ref.rollout.tensor_model_parallel_size=2 \ + actor_rollout_ref.rollout.name=vllm \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \ + actor_rollout_ref.rollout.n=2 \ + actor_rollout_ref.rollout.temperature=1.0 \ + actor_rollout_ref.rollout.top_p=1.0 \ + actor_rollout_ref.rollout.top_k=-1 \ + actor_rollout_ref.rollout.enforce_eager=True \ + actor_rollout_ref.rollout.free_cache_engine=True \ + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \ + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \ + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \ + actor_rollout_ref.actor.strategy=megatron \ + actor_rollout_ref.actor.kl_loss_coef=0.0 \ + actor_rollout_ref.actor.clip_ratio_low=0.2 \ + actor_rollout_ref.actor.clip_ratio_high=0.28 \ + actor_rollout_ref.actor.clip_ratio_c=10.0 \ + actor_rollout_ref.actor.ppo_epochs=1 \ + actor_rollout_ref.actor.use_dynamic_bsz=True \ + actor_rollout_ref.model.path="${MODEL_PATH}" \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.actor.ppo_mini_batch_size=8 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \ + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True \ + +actor_rollout_ref.rollout.engine_kwargs.vllm.compilation_config.cudagraph_mode="FULL_AND_PIECEWISE" \ + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=2 \ + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=2 \ + actor_rollout_ref.actor.megatron.expert_model_parallel_size=2 \ + actor_rollout_ref.actor.megatron.use_dist_checkpointing=${USE_DIST_CKPT} \ + actor_rollout_ref.actor.megatron.dist_checkpointing_path=${DIST_CKPT_PATH} \ + actor_rollout_ref.actor.use_kl_loss=False \ + actor_rollout_ref.actor.loss_agg_mode="token-mean" \ + actor_rollout_ref.ref.strategy=megatron \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=2 \ + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=2 \ + actor_rollout_ref.ref.megatron.expert_model_parallel_size=2 \ + actor_rollout_ref.ref.megatron.use_dist_checkpointing=${USE_DIST_CKPT} \ + actor_rollout_ref.ref.megatron.dist_checkpointing_path=${DIST_CKPT_PATH} \ + reward.reward_manager.name=naive \ + algorithm.kl_ctrl.kl_coef=0.0 \ + trainer.logger=['console'] \ + trainer.project_name='verl_gsm8k_example' \ + trainer.experiment_name='qwen3_30b_a3b_cut_gsm8k_mindspeed' \ + trainer.n_gpus_per_node=8 \ + trainer.nnodes=1 \ + trainer.save_freq=-1 \ + trainer.test_freq=-1 \ + trainer.total_epochs=1 \ + trainer.total_training_steps=1 \ + actor_rollout_ref.actor.use_torch_compile=False \ + actor_rollout_ref.ref.use_torch_compile=False \ + +actor_rollout_ref.actor.megatron.override_transformer_config.use_flash_attn=True $@ + +# clean up +if [[ "$USE_DUMMY_MODEL" == "True" ]]; then + rm -rf $DUMMY_MODEL_PATH + if [[ "$USE_DIST_CKPT" == "True" ]]; then + rm -rf $DIST_CKPT_PATH + fi +fi diff --git a/verl/tests/special_npu/run_qwen3_30b_grpo_mindspeedllm.sh b/verl/tests/special_npu/run_qwen3_30b_grpo_mindspeedllm.sh new file mode 100644 index 0000000000000000000000000000000000000000..90fcc155116351a3ad4cf5ac0543e55596fda580 --- /dev/null +++ b/verl/tests/special_npu/run_qwen3_30b_grpo_mindspeedllm.sh @@ -0,0 +1,267 @@ +set -x + +# Project Configuration +project_name='GRPO-Qwen3-30B-BASE-TEST' +exp_name='GRPO-Qwen3-30B-BASE-MindSpeedLLM-SGLang' + +# Necessary env +export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 +export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 +export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050 +export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050 + +# Node Info +NNODES=${NNODES:-1} +NPUS_PER_NODE=${NPUS_PER_NODE:-8} + +# Model Weights Paths +MODEL_ID=${MODEL_ID:-Qwen/Qwen3-30B-A3B-Instruct-2507} +MODEL_PATH=${MODEL_PATH:-${HOME}/.cache/models/${MODEL_ID}} + +# use dummy model +if [[ "$USE_DUMMY_MODEL" == "True" ]]; then + DUMMY_MODEL_PATH=${DUMMY_MODEL_PATH:-${HOME}/models_dummy/${MODEL_ID}} + if [ -z "${DUMMY_MODEL_CONFIG_PATH}" ]; then + echo "[ERROR] DUMMY_MODEL_CONFIG_PATH not set" + exit 1 + fi + + # make sure the path is empty + if [[ -d $DUMMY_MODEL_PATH && $DUMMY_MODEL_PATH != "/" ]]; then + rm -rf $DUMMY_MODEL_PATH + fi + + # init model + python scripts/init_random_model.py \ + --hf_model_path "${MODEL_PATH}" \ + --new_config_path "${DUMMY_MODEL_CONFIG_PATH}" \ + --output_path "${DUMMY_MODEL_PATH}" + + # replace model path + MODEL_PATH=$DUMMY_MODEL_PATH +fi + +# File System Paths +TRAIN_FILE=$HOME/data/gsm8k/train.parquet +TEST_FILE=$HOME/data/gsm8k/test.parquet +# Data Length Configuration +max_prompt_length=$((512)) +max_response_length=$((128)) + +# Training Batch Configuration +train_prompt_bsz=16 +train_prompt_mini_bsz=16 +n_resp_per_prompt=2 +micro_batch_size=1 + +# Algorithm Configuration +adv_estimator=grpo +use_kl_in_reward=False +kl_coef=0.0 +use_kl_loss=True +kl_loss_coef=0.001 + +# Performance and Memory Management Configuration +all_offload=True +use_dynamic_bsz=False +actor_ppo_max_token_len=$(((max_prompt_length + max_response_length))) +infer_ppo_max_token_len=$(((max_prompt_length + max_response_length))) + +# Megatron Parallelism Configuration +train_tp=4 +train_ep=4 +train_etp=1 +train_pp=2 +train_cp=1 + +# SGLang Generation Configuration +gen_tp=4 +gen_dp=1 +gen_ep=1 +gpu_memory_utilization=0.5 +max_model_len=$((max_prompt_length + max_response_length)) +max_num_batched_tokens=$(((max_prompt_length + max_response_length) * 1)) + +# Data Configuration +DATA_CONFIG=( + # File Paths + data.train_files="${TRAIN_FILE}" + data.val_files="${TEST_FILE}" + # Data Structure + data.prompt_key=prompt + # Batch and Length Configuration + data.train_batch_size=${train_prompt_bsz} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + # Preprocessing + data.filter_overlong_prompts=True + data.truncation='left' +) + +# Model Configuration +MODEL_CONFIG=( + # Model Path + actor_rollout_ref.model.path="${MODEL_PATH}" + # Model Processing + actor_rollout_ref.model.use_remove_padding=True +) + +# Reinforcement Learning Algorithm Configuration +ALGORITHM_CONFIG=( + # Advantage Estimation + algorithm.adv_estimator=${adv_estimator} + # KL Divergence Control + algorithm.use_kl_in_reward=${use_kl_in_reward} + algorithm.kl_ctrl.kl_coef=${kl_coef} +) + +ACTOR_CONFIG=( + # Core Runtime Settings + actor_rollout_ref.actor.use_torch_compile=False + actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} + # Loss Function Configuration + actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.entropy_coeff=0 + # PPO Training Parameters + actor_rollout_ref.actor.ppo_epochs=1 + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${micro_batch_size} + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} + actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} + # Optimizer Settings + actor_rollout_ref.actor.optim.lr=1e-6 + # Megatron Parallelism Strategy + actor_rollout_ref.actor.mindspeed.tensor_model_parallel_size=${train_tp} + actor_rollout_ref.actor.mindspeed.pipeline_model_parallel_size=${train_pp} + actor_rollout_ref.actor.mindspeed.expert_model_parallel_size=${train_ep} + actor_rollout_ref.actor.mindspeed.expert_tensor_parallel_size=${train_etp} + # Memory Optimization + actor_rollout_ref.actor.mindspeed.param_offload=${all_offload} + actor_rollout_ref.actor.mindspeed.optimizer_offload=${all_offload} + actor_rollout_ref.actor.mindspeed.grad_offload=${all_offload} + # Model Weights Management + actor_rollout_ref.actor.mindspeed.use_mbridge=True + actor_rollout_ref.actor.mindspeed.vanilla_mbridge=True + # Transformer Architecture Optimizations + actor_rollout_ref.actor.mindspeed.llm_kwargs.spec='[mindspeed_llm.tasks.models.spec.qwen3_spec, layer_spec]' + actor_rollout_ref.actor.mindspeed.llm_kwargs.seq_length=${max_model_len} + actor_rollout_ref.actor.mindspeed.llm_kwargs.micro_batch_size=${micro_batch_size} + +actor_rollout_ref.actor.mindspeed.llm_kwargs.num_query_groups=4 + +actor_rollout_ref.actor.mindspeed.llm_kwargs.recompute_method=uniform + +actor_rollout_ref.actor.mindspeed.llm_kwargs.recompute_granularity=full + +actor_rollout_ref.actor.mindspeed.llm_kwargs.recompute_num_layers=1 + # MOE + +actor_rollout_ref.actor.mindspeed.llm_kwargs.moe_router_load_balancing_type=aux_loss + +actor_rollout_ref.actor.mindspeed.llm_kwargs.moe_permutation_async_comm=True + +actor_rollout_ref.actor.mindspeed.llm_kwargs.moe_token_dispatcher_type=alltoall + +actor_rollout_ref.actor.mindspeed.llm_kwargs.moe_aux_loss_coeff=0.001 + +actor_rollout_ref.actor.mindspeed.llm_kwargs.moe_grouped_gemm=True + +actor_rollout_ref.actor.mindspeed.llm_kwargs.fix_router=True +) + +REF_CONFIG=( + # Core Runtime Settings + actor_rollout_ref.ref.use_torch_compile=False + # Log Probability Inference + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${micro_batch_size} + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} + # Megatron Parallelism Strategy + actor_rollout_ref.ref.mindspeed.tensor_model_parallel_size=${train_tp} + actor_rollout_ref.ref.mindspeed.pipeline_model_parallel_size=${train_pp} + actor_rollout_ref.ref.mindspeed.expert_model_parallel_size=${train_ep} + actor_rollout_ref.ref.mindspeed.expert_tensor_parallel_size=${train_etp} + # Memory Optimization + actor_rollout_ref.ref.mindspeed.param_offload=${all_offload} + # Model Weights Management + actor_rollout_ref.ref.mindspeed.use_mbridge=True + actor_rollout_ref.ref.mindspeed.vanilla_mbridge=True +) + +ROLLOUT_CONFIG=( + # Rollout Engine + actor_rollout_ref.rollout.name=sglang + +actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend="ascend" + # Generation Parameters + actor_rollout_ref.rollout.n=${n_resp_per_prompt} + actor_rollout_ref.rollout.top_p=1.0 + actor_rollout_ref.rollout.top_k=-1 + actor_rollout_ref.rollout.temperature=1.0 + # Log Probability Inference + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${micro_batch_size} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} + # Memory Management + actor_rollout_ref.rollout.gpu_memory_utilization=${gpu_memory_utilization} + # Parallelism Strategy + actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} + actor_rollout_ref.rollout.data_parallel_size=${gen_dp} + actor_rollout_ref.rollout.expert_parallel_size=${gen_ep} + +actor_rollout_ref.rollout.engine_kwargs.sglang.enable_dp_attention=False + # Performance Optimization + +actor_rollout_ref.rollout.engine_kwargs.sglang.chunked_prefill_size=-1 + actor_rollout_ref.rollout.enforce_eager=False + # Validation Generation + actor_rollout_ref.rollout.val_kwargs.n=1 + actor_rollout_ref.rollout.val_kwargs.do_sample=True + actor_rollout_ref.rollout.val_kwargs.top_p=1.0 + actor_rollout_ref.rollout.val_kwargs.top_k=-1 + actor_rollout_ref.rollout.val_kwargs.temperature=1.0 +) + +TRAINER_CONFIG=( + # Logger Configuration + trainer.logger='["console"]' + # Project Settings + trainer.project_name="${project_name}" + trainer.experiment_name="${exp_name}" + # Hardware Configuration + trainer.nnodes="${NNODES}" + trainer.n_gpus_per_node="${NPUS_PER_NODE}" + trainer.device='npu' + # Training Schedule + trainer.total_epochs=1 + trainer.val_before_train=False + trainer.test_freq=-1 + trainer.save_freq=-1 + trainer.total_training_steps=1 +) + +# profiling configuration +PROF_CONFIG=( + global_profiler.tool=npu + global_profiler.steps=null + global_profiler.save_path=/profpath + actor_rollout_ref.actor.profiler.enable=True + actor_rollout_ref.actor.profiler.ranks="[0]" + actor_rollout_ref.actor.profiler.all_ranks=False + actor_rollout_ref.actor.profiler.tool_config.npu.discrete=True + actor_rollout_ref.actor.profiler.tool_config.npu.contents=['npu','cpu'] + actor_rollout_ref.actor.profiler.tool_config.npu.level=level0 + actor_rollout_ref.actor.profiler.tool_config.npu.analysis=True + actor_rollout_ref.rollout.profiler.enable=True + actor_rollout_ref.rollout.profiler.ranks="[0]" + actor_rollout_ref.rollout.profiler.all_ranks=False +) + +python3 -m verl.trainer.main_ppo \ + --config-path=config \ + --config-name='ppo_trainer.yaml' \ + model_engine=mindspeed \ + "${DATA_CONFIG[@]}" \ + "${MODEL_CONFIG[@]}" \ + "${ACTOR_CONFIG[@]}" \ + "${REF_CONFIG[@]}" \ + "${ROLLOUT_CONFIG[@]}" \ + "${ALGORITHM_CONFIG[@]}" \ + "${TRAINER_CONFIG[@]}" \ + "${PROF_CONFIG[@]}" \ + "$@" + +# clean up +if [[ "$USE_DUMMY_MODEL" == "True" ]]; then + rm -rf $DUMMY_MODEL_PATH + if [[ "$USE_DIST_CKPT" == "True" ]]; then + rm -rf $DIST_CKPT_PATH + fi +fi \ No newline at end of file diff --git a/verl/tests/special_npu/run_qwen3_8b_grpo_mindspeedllm.sh b/verl/tests/special_npu/run_qwen3_8b_grpo_mindspeedllm.sh new file mode 100644 index 0000000000000000000000000000000000000000..6ac8c20fbc307918045dc3b991be7e95644694a5 --- /dev/null +++ b/verl/tests/special_npu/run_qwen3_8b_grpo_mindspeedllm.sh @@ -0,0 +1,222 @@ +set -x + +# Project Configuration +project_name='GRPO-Qwen3-8B-BASE-TEST' +exp_name='GRPO-Qwen3-8B-BASE-MindSpeedLLM-SGLang' + +# Necessary env +export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 +export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 +export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050 +export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050 + +# Node Info +NNODES=${NNODES:-1} +NPUS_PER_NODE=${NPUS_PER_NODE:-8} + +# Model Weights Paths +MODEL_ID=${MODEL_ID:-Qwen/Qwen3-8B} +MODEL_PATH=${MODEL_PATH:-${HOME}/.cache/models/${MODEL_ID}} + +# File System Paths +TRAIN_FILE=$HOME/data/gsm8k/train.parquet +TEST_FILE=$HOME/data/gsm8k/test.parquet +# Data Length Configuration +max_prompt_length=$((512)) +max_response_length=$((128)) + +# Training Batch Configuration +train_prompt_bsz=16 +train_prompt_mini_bsz=16 +n_resp_per_prompt=2 +micro_batch_size=1 + +# Algorithm Configuration +adv_estimator=grpo +use_kl_in_reward=False +kl_coef=0.0 +use_kl_loss=True +kl_loss_coef=0.001 + +# Performance and Memory Management Configuration +all_offload=True +use_dynamic_bsz=False +actor_ppo_max_token_len=$(((max_prompt_length + max_response_length))) +infer_ppo_max_token_len=$(((max_prompt_length + max_response_length))) + +# Megatron Parallelism Configuration +train_tp=4 +train_pp=2 + +# SGLang Generation Configuration +gen_tp=4 +gen_dp=1 +gpu_memory_utilization=0.5 +max_model_len=$((max_prompt_length + max_response_length)) +max_num_batched_tokens=$(((max_prompt_length + max_response_length) * 1)) + +# Data Configuration +DATA_CONFIG=( + # File Paths + data.train_files="${TRAIN_FILE}" + data.val_files="${TEST_FILE}" + # Data Structure + data.prompt_key=prompt + # Batch and Length Configuration + data.train_batch_size=${train_prompt_bsz} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + # Preprocessing + data.filter_overlong_prompts=True + data.truncation='left' +) + +# Model Configuration +MODEL_CONFIG=( + # Model Path + actor_rollout_ref.model.path="${MODEL_PATH}" + # Model Processing + actor_rollout_ref.model.use_remove_padding=True +) + +# Reinforcement Learning Algorithm Configuration +ALGORITHM_CONFIG=( + # Advantage Estimation + algorithm.adv_estimator=${adv_estimator} + # KL Divergence Control + algorithm.use_kl_in_reward=${use_kl_in_reward} + algorithm.kl_ctrl.kl_coef=${kl_coef} +) + +ACTOR_CONFIG=( + # Core Runtime Settings + actor_rollout_ref.actor.use_torch_compile=False + actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} + # Loss Function Configuration + actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} + actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} + actor_rollout_ref.actor.entropy_coeff=0 + # PPO Training Parameters + actor_rollout_ref.actor.ppo_epochs=1 + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${micro_batch_size} + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} + actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} + # Optimizer Settings + actor_rollout_ref.actor.optim.lr=1e-6 + # Megatron Parallelism Strategy + actor_rollout_ref.actor.mindspeed.tensor_model_parallel_size=${train_tp} + actor_rollout_ref.actor.mindspeed.pipeline_model_parallel_size=${train_pp} + # Memory Optimization + actor_rollout_ref.actor.mindspeed.param_offload=${all_offload} + actor_rollout_ref.actor.mindspeed.optimizer_offload=${all_offload} + actor_rollout_ref.actor.mindspeed.grad_offload=${all_offload} + # Model Weights Management + actor_rollout_ref.actor.mindspeed.use_mbridge=True + actor_rollout_ref.actor.mindspeed.vanilla_mbridge=True + # Transformer Architecture Optimizations + actor_rollout_ref.actor.mindspeed.llm_kwargs.spec='[mindspeed_llm.tasks.models.spec.qwen3_spec, layer_spec]' + actor_rollout_ref.actor.mindspeed.llm_kwargs.seq_length=${max_model_len} + actor_rollout_ref.actor.mindspeed.llm_kwargs.micro_batch_size=${micro_batch_size} + +actor_rollout_ref.actor.mindspeed.llm_kwargs.num_query_groups=8 + +actor_rollout_ref.actor.mindspeed.llm_kwargs.recompute_method=uniform + +actor_rollout_ref.actor.mindspeed.llm_kwargs.recompute_granularity=full + +actor_rollout_ref.actor.mindspeed.llm_kwargs.recompute_num_layers=1 + +actor_rollout_ref.actor.mindspeed.llm_kwargs.overlap_grad_reduce=True + +actor_rollout_ref.actor.mindspeed.llm_kwargs.overlap_param_gather=True +) + +REF_CONFIG=( + # Core Runtime Settings + actor_rollout_ref.ref.use_torch_compile=False + # Log Probability Inference + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${micro_batch_size} + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} + # Megatron Parallelism Strategy + actor_rollout_ref.ref.mindspeed.tensor_model_parallel_size=${train_tp} + actor_rollout_ref.ref.mindspeed.pipeline_model_parallel_size=${train_pp} + # Memory Optimization + actor_rollout_ref.ref.mindspeed.param_offload=${all_offload} + # Model Weights Management + actor_rollout_ref.ref.mindspeed.use_mbridge=True + actor_rollout_ref.ref.mindspeed.vanilla_mbridge=True +) + +ROLLOUT_CONFIG=( + # Rollout Engine + actor_rollout_ref.rollout.name=sglang + +actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend="ascend" + # Generation Parameters + actor_rollout_ref.rollout.n=${n_resp_per_prompt} + actor_rollout_ref.rollout.top_p=1.0 + actor_rollout_ref.rollout.top_k=-1 + actor_rollout_ref.rollout.temperature=1.0 + # Log Probability Inference + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${micro_batch_size} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} + # Memory Management + actor_rollout_ref.rollout.gpu_memory_utilization=${gpu_memory_utilization} + # Parallelism Strategy + actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} + actor_rollout_ref.rollout.data_parallel_size=${gen_dp} + +actor_rollout_ref.rollout.engine_kwargs.sglang.enable_dp_attention=False + # Performance Optimization + +actor_rollout_ref.rollout.engine_kwargs.sglang.chunked_prefill_size=-1 + actor_rollout_ref.rollout.enforce_eager=False + # Validation Generation + actor_rollout_ref.rollout.val_kwargs.n=1 + actor_rollout_ref.rollout.val_kwargs.do_sample=True + actor_rollout_ref.rollout.val_kwargs.top_p=1.0 + actor_rollout_ref.rollout.val_kwargs.top_k=-1 + actor_rollout_ref.rollout.val_kwargs.temperature=1.0 +) + +TRAINER_CONFIG=( + # Logger Configuration + trainer.logger='["console"]' + # Project Settings + trainer.project_name="${project_name}" + trainer.experiment_name="${exp_name}" + # Hardware Configuration + trainer.nnodes="${NNODES}" + trainer.n_gpus_per_node="${NPUS_PER_NODE}" + trainer.device='npu' + # Training Schedule + trainer.total_epochs=1 + trainer.val_before_train=False + trainer.test_freq=-1 + trainer.save_freq=-1 + trainer.total_training_steps=1 +) + +# profiling configuration +PROF_CONFIG=( + global_profiler.tool=npu + global_profiler.steps=null + global_profiler.save_path=/profpath + actor_rollout_ref.actor.profiler.enable=True + actor_rollout_ref.actor.profiler.ranks="[0]" + actor_rollout_ref.actor.profiler.all_ranks=False + actor_rollout_ref.actor.profiler.tool_config.npu.discrete=True + actor_rollout_ref.actor.profiler.tool_config.npu.contents=['npu','cpu'] + actor_rollout_ref.actor.profiler.tool_config.npu.level=level0 + actor_rollout_ref.actor.profiler.tool_config.npu.analysis=True + actor_rollout_ref.rollout.profiler.enable=True + actor_rollout_ref.rollout.profiler.ranks="[0]" + actor_rollout_ref.rollout.profiler.all_ranks=False +) + +python3 -m verl.trainer.main_ppo \ + --config-path=config \ + --config-name='ppo_trainer.yaml' \ + model_engine=mindspeed \ + "${DATA_CONFIG[@]}" \ + "${MODEL_CONFIG[@]}" \ + "${ACTOR_CONFIG[@]}" \ + "${REF_CONFIG[@]}" \ + "${ROLLOUT_CONFIG[@]}" \ + "${ALGORITHM_CONFIG[@]}" \ + "${TRAINER_CONFIG[@]}" \ + "${PROF_CONFIG[@]}" \ + "$@" diff --git a/verl/tests/special_sanity/check_api_docs.py b/verl/tests/special_sanity/check_api_docs.py new file mode 100644 index 0000000000000000000000000000000000000000..03756f2d284ddcb58b41e068b4abd560b2d074f7 --- /dev/null +++ b/verl/tests/special_sanity/check_api_docs.py @@ -0,0 +1,142 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Fail CI if any function or class that is publicly exported via +``__all__`` lacks a docstring. + +Usage +----- + # Check specific modules or packages + python check_docstrings.py mypkg.core mypkg.utils + + # Check an entire source tree (all top-level packages under cwd) + python check_docstrings.py +""" + +from __future__ import annotations + +import argparse +import importlib +import inspect +import pkgutil +import sys +from pathlib import Path +from types import ModuleType +from typing import Iterable + +_ALLOW_LIST = [ + "verl.third_party.vllm.LLM", + "verl.third_party.vllm.parallel_state", + "verl.utils.profiler.WorkerProfiler", + "verl.utils.profiler.WorkerProfilerExtension", + "verl.utils.profiler.log_gpu_memory_usage", + "verl.utils.profiler.log_print", + "verl.utils.profiler.mark_annotate", + "verl.utils.profiler.mark_end_range", + "verl.utils.profiler.mark_start_range", + "verl.models.mcore.qwen2_5_vl.get_vision_model_config", + "verl.models.mcore.qwen2_5_vl.get_vision_projection_config", + "verl.models.mcore.mbridge.freeze_moe_router", + "verl.models.mcore.mbridge.make_value_model", + "verl.utils.transformers_compat.flash_attn_supports_top_left_mask", +] + + +def iter_submodules(root: ModuleType) -> Iterable[ModuleType]: + """Yield *root* and every sub-module inside it.""" + yield root + + def print_pkg_error(pkg_name): + print(f"[warn] Skipping {pkg_name!r}", file=sys.stderr) + + if getattr(root, "__path__", None): # only packages have __path__ + for mod_info in pkgutil.walk_packages(root.__path__, prefix=f"{root.__name__}.", onerror=print_pkg_error): + try: + yield importlib.import_module(mod_info.name) + except Exception as exc: + print(f"[warn] Skipping {mod_info.name!r}: {exc}", file=sys.stderr) + + +def names_missing_doc(mod: ModuleType) -> list[str]: + """Return fully-qualified names that need docstrings.""" + missing: list[str] = [] + public = getattr(mod, "__all__", []) + for name in public: + obj = getattr(mod, name, None) + if f"{mod.__name__}.{name}" in _ALLOW_LIST: + continue + if obj is None: + # Exported but not found in the module: flag it anyway. + missing.append(f"{mod.__name__}.{name} (not found)") + continue + + if inspect.isfunction(obj) or inspect.isclass(obj): + doc = inspect.getdoc(obj) + if not doc or not doc.strip(): + missing.append(f"{mod.__name__}.{name}") + return missing + + +def check_module(qualname: str) -> list[str]: + """Import *qualname* and check it (and sub-modules).""" + try: + module = importlib.import_module(qualname) + except ModuleNotFoundError as exc: + print(f"[error] Cannot import '{qualname}': {exc}", file=sys.stderr) + return [qualname] + + missing: list[str] = [] + for submod in iter_submodules(module): + missing.extend(names_missing_doc(submod)) + return missing + + +def autodiscover_packages() -> list[str]: + """Detect top-level packages under CWD when no argument is given.""" + pkgs: list[str] = [] + for p in Path.cwd().iterdir(): + if p.is_dir() and (p / "__init__.py").exists(): + pkgs.append(p.name) + return pkgs + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "modules", + nargs="*", + help="Fully-qualified module or package names (defaults to every top-level package found in CWD).", + ) + args = parser.parse_args() + + targets = args.modules or autodiscover_packages() + if not targets: + raise ValueError("[error] No modules specified and none detected automatically.") + + all_missing: list[str] = [] + for modname in targets: + all_missing.extend(check_module(modname)) + + if all_missing: + print("\nMissing docstrings:") + for name in sorted(all_missing): + print(f" - {name}") + raise ValueError("Missing docstrings detected. Please enhance them with docs accordingly.") + + print("✅ All exported functions/classes have docstrings.") + + +if __name__ == "__main__": + main() diff --git a/verl/tests/special_sanity/check_dataproto_usage.py b/verl/tests/special_sanity/check_dataproto_usage.py new file mode 100644 index 0000000000000000000000000000000000000000..7c8521ab12e0fc2f39dd965d3aefbb4f303c12c9 --- /dev/null +++ b/verl/tests/special_sanity/check_dataproto_usage.py @@ -0,0 +1,69 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +This CI test is used for checking whether DataProto is used in the code of some directory +""" + +import os +from argparse import ArgumentParser +from pathlib import Path + +SEARCH_WHITELIST = [] + +SEARCH_KEYWORDS = ["DataProto"] + + +if __name__ == "__main__": + parser = ArgumentParser() + parser.add_argument("--directory", "-d", required=True, type=str) + args = parser.parse_args() + directory_in_str = args.directory + + pathlist = Path(directory_in_str).glob("**/*.py") + for path in pathlist: + path_in_str = str(path.absolute()) + + # judge whether current path is in pre-defined search whitelist or not. + path_in_whitelist = False + + for sw in SEARCH_WHITELIST: + # for easy debugging in non-linux system + sw = sw.replace("/", os.sep) + if sw in path_in_str: + print(f"[SKIP] File {path_in_str} is in device api usage check whitelist, checking is skipped.") + path_in_whitelist = True + break + + if path_in_whitelist: + continue + + with open(path_in_str, encoding="utf-8") as f: + file_content = f.read() + + find_invalid_device_management = False + + for sk in SEARCH_KEYWORDS: + if sk in file_content: + find_invalid_device_management = True + break + + print( + f"[CHECK] File {path_in_str} is detected for DataProto usage check, check result: " + f"{'success' if not find_invalid_device_management else f'failed, because detect {sk}'}." + ) + + assert not find_invalid_device_management, ( + f"file {path_in_str} contains DataProto usage, please use TensorDict directly!" + ) diff --git a/verl/tests/special_sanity/check_device_api_usage.py b/verl/tests/special_sanity/check_device_api_usage.py new file mode 100644 index 0000000000000000000000000000000000000000..afbc012f6dcb7439238de9e8906ce50d7cc3667d --- /dev/null +++ b/verl/tests/special_sanity/check_device_api_usage.py @@ -0,0 +1,110 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +This CI test is used for checking whether device api usage is irregular, suggest using api in `verl/utils/device.py`. +Search targets include .py files in verl/recipe and verl/verl. +Some files that must contain ".cuda", "cuda" or "nccl" keyword is pre-defined in whitelist below. +""" + +import os +from argparse import ArgumentParser +from pathlib import Path + +# directory or file path must contain keyword ".cuda" or "cuda" +CUDA_KEYWORD_CHECK_WHITELIST = [ + "verl/utils/device.py", + "verl/utils/torch_functional.py", # import flash_attn only on cuda + "verl/utils/profiler/nvtx_profile.py", # appear in NsightSystemsProfiler + "verl/utils/profiler/torch_profile.py", # appear in TorchProfiler + "verl/utils/profiler/config.py", # appear in TorchProfilerToolConfig + "verl/utils/kernel/linear_cross_entropy.py", # appear in nvidia nvtx + "verl/utils/rendezvous/ray_backend.py", # appear in cupy importance + "verl/single_controller/ray/base.py", # appear in default device_name + "verl/trainer/ppo/ray_trainer.py", # appear in default device_name + "verl/experimental/transfer_queue/ray_trainer.py", # appear in docstring as default device_name + "verl/experimental/one_step_off_policy/ray_trainer.py", # appear in docstring as default device_name + "verl/utils/reward_score/sandbox_fusion/utils.py", # appear in sandbox language type + "verl/third_party/torch/distributed/_state_dict_utils.py", # torch monkey patch fixes + "verl/third_party/torch/distributed/checkpoint/state_dict.py", # torch monkey patch fixes + "verl/workers/engine/base.py", # appear in default device_name + "verl/workers/engine/utils.py", # appear in enable_full_determinism + "verl/workers/engine/fsdp/transformer_impl.py", # appear in default device_name + "verl/workers/engine/veomni/transformer_impl.py", # appear in default device_name + "verl/workers/engine/torchtitan/transformer_impl.py", # appear in default device_name + "verl/workers/engine/torchtitan/utils.py", # appear in torch.cuda.empty_cache() + "verl/workers/engine/automodel/transformer_impl.py", # appear in default device_name + "verl/workers/rollout/vllm_rollout/vllm_async_server.py", # appear in config.cudagraph_capture_sizes + "verl/workers/rollout/sglang_rollout/async_sglang_server.py", # manually set CUDA_VISIBLE_DEVICES + "verl/workers/rollout/trtllm_rollout/trtllm_async_server.py", # appear in config.cudagraph_capture_sizes + "verl/workers/rollout/replica.py", # appear in default device_name + "verl/checkpoint_engine", # checkpoint engine backend are device specific + "verl/utils/modelopt/megatron_qat_patch.py", # appear in torch.cuda.empty_cache() + "verl/models/mcore/patch.py", # checkpoint patch only on cuda +] + +# directory or file path must contain keyword "nccl" +NCCL_KEYWORD_CHECK_WHITELIST = [ + "verl/utils/device.py", + "verl/third_party/sglang/parallel_state.py", # appear in default backend +] + +SEARCH_WHITELIST = CUDA_KEYWORD_CHECK_WHITELIST + NCCL_KEYWORD_CHECK_WHITELIST + +SEARCH_KEYWORDS = [".cuda", '"cuda"', '"nccl"'] + + +if __name__ == "__main__": + parser = ArgumentParser() + parser.add_argument("--directory", "-d", required=True, type=str) + args = parser.parse_args() + directory_in_str = args.directory + + pathlist = Path(directory_in_str).glob("**/*.py") + for path in pathlist: + path_in_str = str(path.absolute()) + + # judge whether current path is in pre-defined search whitelist or not. + path_in_whitelist = False + + for sw in SEARCH_WHITELIST: + # for easy debugging in non-linux system + sw = sw.replace("/", os.sep) + if sw in path_in_str: + print(f"[SKIP] File {path_in_str} is in device api usage check whitelist, checking is skipped.") + path_in_whitelist = True + break + + if path_in_whitelist: + continue + + with open(path_in_str, encoding="utf-8") as f: + file_content = f.read() + + find_invalid_device_management = False + + for sk in SEARCH_KEYWORDS: + if sk in file_content: + find_invalid_device_management = True + break + + print( + f"[CHECK] File {path_in_str} is detected for device api usage check, check result: " + f"{'success' if not find_invalid_device_management else f'failed, because detect {sk}'}." + ) + + assert not find_invalid_device_management, ( + f'file {path_in_str} contains .cuda/"cuda"/"nccl" usage, please use api in ' + f"verl/utils/device.py directly." + ) diff --git a/verl/tests/special_sanity/check_docs_time_info.py b/verl/tests/special_sanity/check_docs_time_info.py new file mode 100644 index 0000000000000000000000000000000000000000..a54d1d50a7e9d21202387e2c9c8e3c6c73a5d807 --- /dev/null +++ b/verl/tests/special_sanity/check_docs_time_info.py @@ -0,0 +1,84 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Check that every .md and .rst file under docs/ contains the substring "Last updated", +with an allow-list for exceptions. +""" + +import sys +from pathlib import Path + +# === CONFIGURATION === + +# Relative paths (to docs/) or glob patterns to skip checking +ALLOW_LIST = { + "docs/README.md", # you can list individual files + "docs/legacy/*.rst", # or glob patterns + "docs/index.rst", + "docs/start/install.rst", + "docs/start/quickstart.rst", + "docs/README_vllm0.7.md", +} + +# The folder to scan +DOCS_DIR = Path("docs") + +# === SCRIPT === + + +def is_allowed(path: Path) -> bool: + """ + Return True if `path` matches any entry in ALLOW_LIST. + """ + rel = str(path) + for pattern in ALLOW_LIST: + if Path(rel).match(pattern): + return True + return False + + +def main(): + if not DOCS_DIR.exists(): + print(f"Error: Documentation directory '{DOCS_DIR}' does not exist.", file=sys.stderr) + sys.exit(1) + + missing = [] + + # Gather all .md and .rst files under docs/ + for ext in ("*.md", "*.rst"): + for path in DOCS_DIR.rglob(ext): + if is_allowed(path): + continue + + text = path.read_text(encoding="utf-8", errors="ignore") + if "Last updated" not in text: + missing.append(path) + + # Report + if missing: + print("\nThe following files are missing the 'Last updated' string:\n") + for p in missing: + print(f" - {p}") + print(f"\nTotal missing: {len(missing)}\n", file=sys.stderr) + raise AssertionError( + "Some documentation files lack a 'Last updated' line. Please include info such as " + "'Last updated: mm/dd/yyyy' to indicate the last update time of the document." + ) + else: + print("✅ All checked files contain 'Last updated'.") + + +if __name__ == "__main__": + main() diff --git a/verl/tests/special_sanity/check_docstrings.py b/verl/tests/special_sanity/check_docstrings.py new file mode 100644 index 0000000000000000000000000000000000000000..829de9d2b76c636b99da35e68360f8f6e1b7693e --- /dev/null +++ b/verl/tests/special_sanity/check_docstrings.py @@ -0,0 +1,154 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Python script to check docstrings for functions and classes in specified files. +Checks that every public function and class has proper docstring documentation. +""" + +import ast +import os +import sys + + +class DocstringChecker(ast.NodeVisitor): + """AST visitor to check for missing docstrings in functions and classes.""" + + def __init__(self, filename: str): + self.filename = filename + self.missing_docstrings: list[tuple[str, str, int]] = [] + self.current_class = None + self.function_nesting_level = 0 + + def visit_FunctionDef(self, node: ast.FunctionDef): + """Visit function definitions and check for docstrings.""" + if not node.name.startswith("_") and self.function_nesting_level == 0: + if not self._has_docstring(node): + func_name = f"{self.current_class}.{node.name}" if self.current_class else node.name + self.missing_docstrings.append((func_name, self.filename, node.lineno)) + + self.function_nesting_level += 1 + self.generic_visit(node) + self.function_nesting_level -= 1 + + def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef): + """Visit async function definitions and check for docstrings.""" + if not node.name.startswith("_") and self.function_nesting_level == 0: + if not self._has_docstring(node): + func_name = f"{self.current_class}.{node.name}" if self.current_class else node.name + self.missing_docstrings.append((func_name, self.filename, node.lineno)) + + self.function_nesting_level += 1 + self.generic_visit(node) + self.function_nesting_level -= 1 + + def visit_ClassDef(self, node: ast.ClassDef): + """Visit class definitions and check for docstrings.""" + if not node.name.startswith("_"): + if not self._has_docstring(node): + self.missing_docstrings.append((node.name, self.filename, node.lineno)) + + old_class = self.current_class + self.current_class = node.name + self.generic_visit(node) + self.current_class = old_class + + def _has_docstring(self, node) -> bool: + """Check if a node has a docstring.""" + return ast.get_docstring(node) is not None + + +def check_file_docstrings(filepath: str) -> list[tuple[str, str, int]]: + """Check docstrings in a single file.""" + try: + with open(filepath, encoding="utf-8") as f: + content = f.read() + + tree = ast.parse(content, filename=filepath) + checker = DocstringChecker(filepath) + checker.visit(tree) + return checker.missing_docstrings + + except Exception as e: + print(f"Error processing {filepath}: {e}") + return [] + + +def main(): + """Main function to check docstrings in specified files.""" + + files_to_check = [ + "verl/trainer/ppo/ray_trainer.py", + "verl/trainer/main_ppo.py", + "verl/trainer/ppo/reward.py", + "verl/utils/reward_score/__init__.py", + "verl/trainer/ppo/core_algos.py", + "verl/experimental/agent_loop/agent_loop.py", + ] + + script_dir = os.path.dirname(os.path.abspath(__file__)) + repo_path = os.path.dirname(os.path.dirname(script_dir)) + + if not os.path.exists(repo_path): + print(f"Repository path {repo_path} does not exist!") + sys.exit(1) + + os.chdir(repo_path) + + all_missing_docstrings = [] + + print("Checking docstrings in specified files...") + print("=" * 60) + + for file_path in files_to_check: + if not os.path.exists(file_path): + print(f"Warning: File {file_path} does not exist!") + continue + + print(f"Checking {file_path}...") + missing = check_file_docstrings(file_path) + all_missing_docstrings.extend(missing) + + if missing: + print(f" Found {len(missing)} missing docstrings") + else: + print(" All functions and classes have docstrings [OK]") + + print("=" * 60) + + if all_missing_docstrings: + print(f"\nSUMMARY: Found {len(all_missing_docstrings)} functions/classes missing docstrings:") + print("-" * 60) + + by_file = {} + for name, filepath, lineno in all_missing_docstrings: + if filepath not in by_file: + by_file[filepath] = [] + by_file[filepath].append((name, lineno)) + + for filepath in sorted(by_file.keys()): + print(f"\n{filepath}:") + for name, lineno in sorted(by_file[filepath], key=lambda x: x[1]): + print(f" - {name} (line {lineno})") + + print(f"\nTotal missing docstrings: {len(all_missing_docstrings)}") + + raise Exception(f"Found {len(all_missing_docstrings)} functions/classes without proper docstrings!") + + else: + print("\n[OK] All functions and classes have proper docstrings!") + + +if __name__ == "__main__": + main() diff --git a/verl/tests/special_sanity/check_example_naming.py b/verl/tests/special_sanity/check_example_naming.py new file mode 100644 index 0000000000000000000000000000000000000000..8f1f92798a0863d1b8c8fcda9e71da44271d2a87 --- /dev/null +++ b/verl/tests/special_sanity/check_example_naming.py @@ -0,0 +1,247 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Enforce the canonical naming convention for example run scripts. + +The current convention (see ``examples/README.md``) is:: + + run__.sh + +Where ```` is one of ``fsdp``, ``fsdp2``, ``megatron``, +``mindspeed``, ``automodel`` or ``veomni``, and **must be the last +underscore-separated token before** ``.sh`` — nothing follows it. The legacy +convention used to embed the inference backend (``vllm``/``sglang``/``trtllm``), +platform tokens (``_npu``/``_amd``), machine-type tokens (``_gb200``, +``_blackwell``), quantization variants (``_fp8``), and ad-hoc trailing +suffixes into the filename. All of those are now merged into a single +canonical script and selected at runtime via env-var toggles +(``INFER_BACKEND``, ``DEVICE``, ``MACHINE``, ``QUANT``, ...). + +Usage:: + + python3 tests/special_sanity/check_example_naming.py + python3 tests/special_sanity/check_example_naming.py --root examples \ + --ignore-dirs examples/data_preprocess examples/profile + +Exits with status 1 (and prints offending paths) if any script under +``--root`` violates the convention. +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +# Tokens that may NOT appear as a word in the filename. They used to be +# embedded as ``_vllm_fsdp``/``_sglang_fsdp``/``_trtllm_fsdp``/``_npu``/ +# ``_gb200``/``_fp8`` etc.; the active convention exposes them via env-var +# toggles instead. +FORBIDDEN_TOKENS = ( + # Inference backends. + "vllm", + "sglang", + "trtllm", + # Platform / accelerator vendors. + "npu", + "amd", + # Specific GPU machine types — selected at runtime via ``MACHINE`` env + # var, never embedded in filenames. + "gb200", + "b200", + "blackwell", + # Quantization variants. + "fp8", +) + +# Recognised train-backend / engine markers. The filename must end with +# ``_.sh`` (i.e. the train-backend is the LAST underscore- +# separated token). Generation-only scripts that do not run a trainer are +# listed in ``DEFAULT_IGNORE_FILES`` instead. +ALLOWED_BACKENDS = ( + "fsdp", + "fsdp2", + "megatron", + "mindspeed", + "automodel", + "veomni", +) + +# Directories whose scripts have their own conventions and are exempt from +# the train-backend rule. Paths are interpreted relative to the repo root. +DEFAULT_IGNORE_DIRS = ( + "examples/data_preprocess", + "examples/profile", + "examples/tutorial", + "examples/vllm_omni", + # MTP scripts encode async/multinode dispatch in the filename; they pre- + # date this convention. Migrate separately rather than block this hook. + "examples/mtp_trainer", +) + +# Individual files that are exempt. Use sparingly and document why. +DEFAULT_IGNORE_FILES = ( + # Inference-only entry point — no train-backend in the name. + "examples/generation/run_deepseek_llm_7b.sh", + # Diffusion (flow-grpo) example using the vllm_omni rollout; the ``omni`` + # token in the name refers to that rollout variant rather than a train + # backend, and this is the only example in flowgrpo_trainer/. + "examples/flowgrpo_trainer/run_qwen_image_omni_lora.sh", + # Sibling of run_qwen2_5_7b_fsdp.sh that adds a multi-rollout-server + # dispatch flavour the canonical script does not yet expose. Migrate by + # folding ``_multi_rs`` into a ROLLOUT_SERVER env-var toggle rather than + # in this PR. + "examples/rollout_correction/run_qwen2_5_7b_fsdp_multi_rs.sh", +) + + +def _split_tokens(stem: str) -> list[str]: + """Split a script stem on underscores, dropping the leading ``run`` token.""" + parts = stem.split("_") + if parts and parts[0] == "run": + parts = parts[1:] + return parts + + +def _is_ignored(path: Path, repo_root: Path, ignore_dirs: tuple[str, ...], ignore_files: tuple[str, ...]) -> bool: + rel = path.relative_to(repo_root).as_posix() + if rel in ignore_files: + return True + return any(rel == d or rel.startswith(d.rstrip("/") + "/") for d in ignore_dirs) + + +def check_filename(path: Path, display: str | None = None) -> list[str]: + """Return a list of human-readable violations for one script path. + + The path's *basename* is what we validate; the directory layout is + enforced separately by the caller via the ``--ignore-dirs`` mechanism. + ``display`` controls how the path is shown in error messages (defaults + to ``str(path)``). + """ + errors: list[str] = [] + name = path.name + shown = display if display is not None else str(path) + + if not name.startswith("run_"): + errors.append(f"{shown}: example script must be named 'run__.sh'") + return errors + + if not name.endswith(".sh"): + errors.append(f"{shown}: example script must end with '.sh'") + return errors + + tokens = _split_tokens(path.stem) + + found_forbidden = [t for t in tokens if t in FORBIDDEN_TOKENS] + if found_forbidden: + errors.append( + f"{shown}: filename contains forbidden token(s) {found_forbidden}. " + f"Inference-backend ({{vllm,sglang,trtllm}}), platform ({{npu,amd}}), " + f"machine-type ({{gb200,b200,blackwell}}), and quantization " + f"({{fp8}}) selections must be exposed as env-var toggles inside " + f"the script, not embedded in the filename." + ) + + if not tokens or tokens[-1] not in ALLOWED_BACKENDS: + errors.append( + f"{shown}: filename must end with '_.sh' where " + f"train-backend ∈ {list(ALLOWED_BACKENDS)} " + f"(e.g. 'run_qwen3_8b_fsdp.sh'). Anything after the train-backend " + f"(e.g. '_fp8', '_gb200', '_multi_rs') must be folded into an " + f"env-var toggle. If this is an intentional exception, add it to " + f"DEFAULT_IGNORE_FILES or pass '--ignore-files {shown}'." + ) + + return errors + + +def collect_scripts( + root: Path, + repo_root: Path, + ignore_dirs: tuple[str, ...], + ignore_files: tuple[str, ...], +) -> list[Path]: + return sorted( + p for p in root.rglob("*.sh") if p.is_file() and not _is_ignored(p, repo_root, ignore_dirs, ignore_files) + ) + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument( + "--root", + type=Path, + default=Path("examples"), + help="Directory to scan (default: examples)", + ) + parser.add_argument( + "--repo-root", + type=Path, + default=Path("."), + help="Repository root used for relative-path matching against --ignore-dirs/--ignore-files", + ) + parser.add_argument( + "--ignore-dirs", + nargs="*", + default=list(DEFAULT_IGNORE_DIRS), + help="Directories (relative to --repo-root) whose scripts are exempt", + ) + parser.add_argument( + "--ignore-files", + nargs="*", + default=list(DEFAULT_IGNORE_FILES), + help="Individual files (relative to --repo-root) that are exempt", + ) + args = parser.parse_args(argv) + + repo_root = args.repo_root.resolve() + root = args.root.resolve() + if not root.is_dir(): + print(f"❌ --root '{args.root}' does not exist or is not a directory.", file=sys.stderr) + return 2 + + scripts = collect_scripts( + root, + repo_root, + tuple(args.ignore_dirs), + tuple(args.ignore_files), + ) + + all_errors: list[str] = [] + for script in scripts: + try: + display = script.relative_to(repo_root).as_posix() + except ValueError: + display = str(script) + all_errors.extend(check_filename(script, display=display)) + + if all_errors: + print("❌ Example script naming violations:\n", file=sys.stderr) + for err in all_errors: + print(" - " + err, file=sys.stderr) + print( + "\nNaming convention (see examples/README.md):\n" + " run__.sh (train-backend is the LAST token)\n" + f" train-backend ∈ {list(ALLOWED_BACKENDS)}\n" + f" must NOT contain any of {list(FORBIDDEN_TOKENS)} as a word.\n", + file=sys.stderr, + ) + return 1 + + print(f"✅ {len(scripts)} example scripts under '{args.root}' follow the naming convention.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/verl/tests/special_sanity/check_license.py b/verl/tests/special_sanity/check_license.py new file mode 100644 index 0000000000000000000000000000000000000000..5e1e581d34b7863245dceb2abb9477dc091cab20 --- /dev/null +++ b/verl/tests/special_sanity/check_license.py @@ -0,0 +1,107 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import subprocess +from argparse import ArgumentParser +from pathlib import Path +from typing import Iterable + +license_head_bytedance = "Copyright 2024 Bytedance Ltd. and/or its affiliates" +license_head_bytedance_25 = "Copyright 2025 Bytedance Ltd. and/or its affiliates" +license_head_bytedance_26 = "Copyright 2026 Bytedance Ltd. and/or its affiliates" +# Add custom license headers below +license_head_prime = "Copyright 2024 PRIME team and/or its affiliates" +license_head_individual = "Copyright 2025 Individual Contributor:" +license_head_sglang = "Copyright 2023-2024 SGLang Team" +license_head_modelbest = "Copyright 2025 ModelBest Inc. and/or its affiliates" +license_head_amazon = "Copyright 2025 Amazon.com Inc and/or its affiliates" +license_head_amazon_26 = "Copyright 2026 Amazon.com Inc and/or its affiliates" +license_head_facebook = "Copyright (c) 2016- Facebook, Inc" +license_head_meituan = "Copyright 2025 Meituan Ltd. and/or its affiliates" +license_head_huawei = "Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved." +license_head_huawei_26 = "Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved." +license_head_nvidia = "Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved." +license_headers = [ + license_head_bytedance, + license_head_bytedance_25, + license_head_bytedance_26, + license_head_prime, + license_head_individual, + license_head_sglang, + license_head_modelbest, + license_head_amazon, + license_head_amazon_26, + license_head_facebook, + license_head_meituan, + license_head_huawei, + license_head_huawei_26, + license_head_nvidia, +] + + +def _git_tracked_py_files() -> set[Path]: + """Return the set of .py files tracked by git (respects .gitignore).""" + result = subprocess.run(["git", "ls-files", "*.py", "**/*.py"], capture_output=True, text=True, check=True) + return {Path(line) for line in result.stdout.splitlines() if line} + + +def get_py_files(path_arg: Path, tracked: set[Path]) -> Iterable[Path]: + """get py files under a dir. if already py file return it + + Args: + path_arg (Path): path to scan for py files + tracked (set[Path]): set of git-tracked py files + + Returns: + Iterable[Path]: list of py files + """ + if path_arg.is_dir(): + return (p for p in tracked if p == path_arg or path_arg in p.parents) + elif path_arg.is_file() and path_arg.suffix == ".py": + return [path_arg] + return [] + + +if __name__ == "__main__": + parser = ArgumentParser() + parser.add_argument( + "--directories", + "-d", + required=True, + type=Path, + nargs="+", + help="List of directories to check for license headers", + ) + args = parser.parse_args() + + # Collect all Python files from specified directories (only git-tracked files) + tracked = _git_tracked_py_files() + pathlist = set(path for path_arg in args.directories for path in get_py_files(path_arg, tracked)) + + for path in pathlist: + # because path is object not string + path_in_str = str(path.absolute()) + print(path_in_str) + # ``git ls-files`` may return paths that were deleted in the working tree + # but not yet staged. Skip them so pre-commit keeps working for devs mid-refactor. + if not path.exists(): + continue + with open(path_in_str, encoding="utf-8") as f: + file_content = f.read() + + has_license = False + for lh in license_headers: + if lh in file_content: + has_license = True + break + assert has_license, f"file {path_in_str} does not contain license" diff --git a/verl/tests/special_sanity/check_pr_description.py b/verl/tests/special_sanity/check_pr_description.py new file mode 100644 index 0000000000000000000000000000000000000000..4ed4563db6088e8562273cebd08116e375bc8bb2 --- /dev/null +++ b/verl/tests/special_sanity/check_pr_description.py @@ -0,0 +1,97 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +#!/usr/bin/env python3 +import json +import os + +# Number of lines to check +NUM_LINES = 5 + + +# Custom exception types for clear error handling +class TemplateFileError(Exception): + pass + + +class PRBodyLoadError(Exception): + pass + + +class PRDescriptionError(Exception): + pass + + +# Path to the PR template file +template_file = os.path.join(os.getenv("GITHUB_WORKSPACE", "."), ".github", "PULL_REQUEST_TEMPLATE.md") + + +def load_template(path): + """ + Load only the first NUM_LINES of the PR template file as a list of lines, + without stripping any characters. + """ + lines = [] + try: + with open(path, encoding="utf-8") as f: + for _ in range(NUM_LINES): + line = f.readline() + if not line: + break + lines.append(line.strip()) + return lines + except Exception as e: + raise TemplateFileError(f"Failed to read PR template (first {NUM_LINES} lines) at {path}: {e}") from e + + +def load_pr_body(event_path): + try: + with open(event_path, encoding="utf-8") as f: + payload = json.load(f) + return payload.get("pull_request", {}).get("body", "") or "" + except Exception as e: + raise PRBodyLoadError(f"Failed to read PR body from {event_path}: {e}") from e + + +def check_pr_description(body, template_lines): + """ + Compare the first NUM_LINES lines of the PR body to the template lines. + If they match exactly, the placeholder was not modified. + """ + pr_lines = body.splitlines(keepends=True) + pr_first = [x.strip() for x in pr_lines[:NUM_LINES]] + if pr_first == template_lines: + raise PRDescriptionError( + "It looks like you haven't updated the '### What does this PR do?' section. Please replace " + "the placeholder text with a concise description of what your PR does." + ) + else: + print(pr_first) + print(template_lines) + + +def main(): + event_path = os.getenv("GITHUB_EVENT_PATH") + if not event_path: + raise OSError("GITHUB_EVENT_PATH is not set.") + + template_lines = load_template(template_file) + pr_body = load_pr_body(event_path) + check_pr_description(pr_body, template_lines) + + print("✅ '### What does this PR do?' section has been filled out.") + + +if __name__ == "__main__": + main() diff --git a/verl/tests/special_sanity/check_pr_title.py b/verl/tests/special_sanity/check_pr_title.py new file mode 100644 index 0000000000000000000000000000000000000000..df316d3d080d4b40f77d84d5922ec05e32eb27fa --- /dev/null +++ b/verl/tests/special_sanity/check_pr_title.py @@ -0,0 +1,73 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import re + +# Get PR title from environment +pr_title = os.environ.get("PR_TITLE", "").strip() + +# Define rules +allowed_modules = ["fsdp", "megatron", "veomni", "sglang", "vllm", "trtllm", "rollout", "trainer"] +allowed_modules += ["tests", "training_utils", "recipe", "hardware", "deployment"] +allowed_modules += ["ray", "worker", "single_controller", "misc", "docker", "ci"] +allowed_modules += ["perf", "model", "algo", "env", "tool", "ckpt", "doc", "data", "cfg", "reward"] +allowed_modules += ["fully_async", "one_step_off"] +allowed_types = ["feat", "fix", "refactor", "chore", "test"] + +# Check for [1/N] prefix and extract the rest of the title +progress_match = re.match(r"^\[\d/[\dNn]\]\s*(.+)$", pr_title, re.IGNORECASE) +if progress_match: + pr_title = progress_match.group(1).strip() + +# Check for [BREAKING] prefix and extract the rest of the title +breaking_match = re.match(r"^\[BREAKING\]\s*(.+)$", pr_title, re.IGNORECASE) +if breaking_match: + core_pr_title = breaking_match.group(1).strip() + is_breaking = True +else: + core_pr_title = pr_title + is_breaking = False + +# Build dynamic regex pattern for modules (now working on core_pr_title) +re_modules_pattern = re.compile(r"^\[([a-z_,\s]+)\]", re.IGNORECASE) +re_modules = re_modules_pattern.match(core_pr_title) +if not re_modules: + print(f"❌ Invalid PR title: '{pr_title}'") + print("Expected format: [BREAKING][module] type: description") + print(f"Allowed modules: {', '.join(allowed_modules)}") + raise Exception("Invalid PR title") +else: + modules = re.findall(r"[a-z_]+", re_modules.group(1).lower()) + if not all(module in allowed_modules for module in modules): + invalid_modules = [module for module in modules if module not in allowed_modules] + print(f"❌ Invalid modules: {', '.join(invalid_modules)}") + print(f"Allowed modules: {', '.join(allowed_modules)}") + raise Exception("Invalid PR title") + +types_pattern = "|".join(re.escape(t) for t in allowed_types) +re_types_pattern = re.compile(rf"^\[[a-z_,\s]+\]\s+({types_pattern}):\s+.+$", re.IGNORECASE) +match = re_types_pattern.match(core_pr_title) + +if not match: + print(f"❌ Invalid PR title: '{pr_title}'") + print("Expected format: [BREAKING][module] type: description") + print(f"Allowed types: {', '.join(allowed_types)}") + raise Exception("Invalid PR title") + +change_type = match.group(1).lower() + +# Build the success message +breaking_info = " (BREAKING CHANGE)" if is_breaking else "" +print(f"✅ PR title is valid: {pr_title}, modules: {modules}, type: {change_type}{breaking_info}") diff --git a/verl/tests/special_sanity/test_check_example_naming.py b/verl/tests/special_sanity/test_check_example_naming.py new file mode 100644 index 0000000000000000000000000000000000000000..41b6814df28ce7ed49a94b958ae2519df4d3e8ca --- /dev/null +++ b/verl/tests/special_sanity/test_check_example_naming.py @@ -0,0 +1,119 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit-style smoke tests for ``check_example_naming``.""" + +from __future__ import annotations + +from pathlib import Path + +from tests.special_sanity.check_example_naming import ( + ALLOWED_BACKENDS, + FORBIDDEN_TOKENS, + check_filename, + main, +) + + +def _violations(name: str) -> list[str]: + return check_filename(Path(f"examples/grpo_trainer/{name}")) + + +def test_canonical_name_passes(): + assert _violations("run_qwen3_8b_fsdp.sh") == [] + + +def test_pre_backend_suffix_passes(): + # Tokens before the train-backend are part of ; only the trailing + # token is constrained. + assert _violations("run_qwen3_8b_from_adapter_fsdp.sh") == [] + assert _violations("run_qwen3_8b_merge_fsdp.sh") == [] + + +def test_all_train_backends_accepted(): + for backend in ALLOWED_BACKENDS: + assert _violations(f"run_qwen3_8b_{backend}.sh") == [], backend + + +def test_forbidden_engine_token_rejected(): + errs = _violations("run_qwen3_8b_vllm_fsdp.sh") + # `vllm` is both a forbidden token AND occupies the last-token slot + # (`fsdp` is no longer last). Ensure at least one error mentions vllm. + assert errs and any("vllm" in e for e in errs) + + +def test_forbidden_platform_token_rejected(): + errs = _violations("run_qwen3_8b_fsdp_npu.sh") + assert errs and any("npu" in e for e in errs) + + +def test_forbidden_quantization_token_rejected(): + errs = _violations("run_qwen3_30b_a3b_megatron_fp8.sh") + assert errs and any("fp8" in e for e in errs) + + +def test_forbidden_machine_token_rejected(): + errs = _violations("run_qwen3_8b_fsdp_gb200.sh") + assert errs and any("gb200" in e for e in errs) + + +def test_trailing_suffix_after_backend_rejected(): + # Even if no token is in FORBIDDEN_TOKENS, anything after the train- + # backend is a violation: the train-backend MUST be the last token. + errs = _violations("run_qwen3_8b_fsdp_extra.sh") + assert errs + assert any("must end with '_.sh'" in e for e in errs) + + +def test_missing_train_backend_rejected(): + errs = _violations("run_qwen3_8b.sh") + assert errs and any(b in errs[0] for b in ALLOWED_BACKENDS) + + +def test_non_run_prefix_rejected(): + errs = _violations("badname.sh") + assert errs and "run_" in errs[0] + + +def test_forbidden_tokens_kept_in_sync_with_legacy_pattern(): + # If we ever forget to forbid one of the deprecated tokens, this test + # would also start failing because at least one filename in the pre- + # refactor era used each of these. + for tok in ("vllm", "sglang", "trtllm", "npu", "gb200", "fp8"): + assert tok in FORBIDDEN_TOKENS, tok + + +def test_repo_tree_passes(tmp_path): + # Run the entry point against the actual ``examples/`` tree to mirror + # what pre-commit will do. ``main`` returns 0 on success. + assert main(["--root", "examples", "--repo-root", "."]) == 0 + + +def test_synthetic_violation_fails(tmp_path): + fake = tmp_path / "examples" / "grpo_trainer" + fake.mkdir(parents=True) + (fake / "run_qwen3_8b_vllm_fsdp.sh").write_text("#!/bin/bash\n") + (fake / "run_ok_fsdp.sh").write_text("#!/bin/bash\n") + + rc = main( + [ + "--root", + str(tmp_path / "examples"), + "--repo-root", + str(tmp_path), + "--ignore-dirs", + "--ignore-files", + ] + ) + assert rc == 1 diff --git a/verl/tests/special_sanity/test_config_docs.py b/verl/tests/special_sanity/test_config_docs.py new file mode 100644 index 0000000000000000000000000000000000000000..b8dc74762450fe41a42db6ca09972851e8dcbdc2 --- /dev/null +++ b/verl/tests/special_sanity/test_config_docs.py @@ -0,0 +1,88 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import re +from pathlib import Path + + +def validate_yaml_format(yaml_lines): + errors = [] + i = 0 + + while i < len(yaml_lines): + line = yaml_lines[i] + stripped = line.strip() + + # Skip empty lines + if stripped == "": + i += 1 + continue + + # Match YAML keys like "field:" or "field: value" + key_match = re.match(r"^(\s*)([a-zA-Z0-9_]+):", line) + if key_match: + # Check if there's a comment above + if i == 0 or not yaml_lines[i - 1].strip().startswith("#"): + errors.append(f"Missing comment above line {i + 1}: {line.strip()}") + + # Check for inline comment + if "#" in line and not stripped.startswith("#"): + comment_index = line.index("#") + colon_index = line.index(":") + if comment_index > colon_index: + errors.append(f"Inline comment found on line {i + 1}: {line.strip()}") + + # Check for blank line after this key line (unless next is a deeper indent) + if i + 1 < len(yaml_lines): + next_line = yaml_lines[i + 1] + next_stripped = next_line.strip() + + # If next is not empty and not a deeper nested line, enforce blank line + if next_stripped != "": + errors.append(f"Missing blank line after line {i + 1}: {line.strip()}") + + i += 1 + + return errors + + +def test_trainer_config_doc(): + yamls_to_inspect = [ + "verl/trainer/config/ppo_trainer.yaml", + "verl/trainer/config/actor/actor.yaml", + "verl/trainer/config/actor/dp_actor.yaml", + "verl/trainer/config/critic/critic.yaml", + "verl/trainer/config/critic/dp_critic.yaml", + "verl/trainer/config/ref/ref.yaml", + "verl/trainer/config/ref/dp_ref.yaml", + "verl/trainer/config/rollout/rollout.yaml", + ] + success = True + for yaml_to_inspect in yamls_to_inspect: + yaml_path = Path(yaml_to_inspect) # path to your YAML file + with open(yaml_path) as f: + lines = f.readlines() + + validation_errors = validate_yaml_format(lines) + if validation_errors: + success = False + print("YAML documentation format check failed:") + print(f"Please read the top block of {yaml_to_inspect} to see format rules:\n") + for err in validation_errors: + print(" -", err) + + if not success: + raise Exception("Please fix documentation format.") + else: + print("YAML format check passed ✅") diff --git a/verl/tests/special_sanity/test_import.py b/verl/tests/special_sanity/test_import.py new file mode 100644 index 0000000000000000000000000000000000000000..4f8a918fe65679c353e8055c2c2b0a428fdf8f7a --- /dev/null +++ b/verl/tests/special_sanity/test_import.py @@ -0,0 +1,25 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def test_import(): + import verl + + print(verl.__version__) + + +def test_single_controller_import(): + import verl.single_controller + + print(verl.single_controller.__version__) diff --git a/verl/tests/special_sanity/type_coverage_check.py b/verl/tests/special_sanity/type_coverage_check.py new file mode 100644 index 0000000000000000000000000000000000000000..c35abaeb2d6d4fb23f0d2c58b4dde56986932a37 --- /dev/null +++ b/verl/tests/special_sanity/type_coverage_check.py @@ -0,0 +1,182 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Custom type annotation check tool. +To inspect the type annotation for functions in the entire codebase, please run: +find verl -type f -name "*.py" | xargs -n 1 python3 tests/special_sanity/type_coverage_check.py --all-lines +--debug --target-file +""" + +import argparse +import ast +import linecache +import subprocess +from pathlib import Path + + +def get_changed_files() -> list[Path]: + result = subprocess.run( + ["git", "diff", "--name-only", "--diff-filter=AM", "origin/main...HEAD"], stdout=subprocess.PIPE, text=True + ) + return [Path(f) for f in result.stdout.splitlines() if f.endswith(".py")] + + +def get_changed_lines(file_path: Path) -> set[int]: + result = subprocess.run( + ["git", "diff", "-U0", "origin/main...HEAD", "--", str(file_path)], + stdout=subprocess.PIPE, + text=True, + ) + lines: set[int] = set() + for line in result.stdout.splitlines(): + if line.startswith("@@"): + for part in line.split(): + try: + if part.startswith("+") and "," in part: + start, count = map(int, part[1:].split(",")) + lines.update(range(start, start + count)) + elif part.startswith("+") and "," not in part: + lines.add(int(part[1:])) + except Exception: + # (vermouth1992) There are many edge cases here because + can be in the changed program + pass + return lines + + +CHECK_SUCCESS = 0 +CHECK_WARNING = 1 +CHECK_FAILURE = -1 + + +def should_check_type(arg_name: str) -> bool: + if arg_name in ("self", "cls"): + return False + if arg_name.startswith("*"): + return False + return True + + +def has_type_annotations(node: ast.AST, debug: bool = False) -> int: + if isinstance(node, ast.FunctionDef): + is_private = node.name.startswith("_") + if node.args.vararg is not None or node.args.kwarg is not None: + return CHECK_SUCCESS + has_ann = ( + all(arg.annotation is not None for arg in node.args.args if should_check_type(arg.arg)) + and node.returns is not None + ) + if has_ann or is_private: + return CHECK_SUCCESS + else: + if debug: + print(node, [(arg.annotation, arg.arg) for arg in node.args.args if should_check_type(arg.arg)]) + return CHECK_FAILURE + return CHECK_SUCCESS + + +def check_file( + file_path: Path, changed_lines: set[int], debug: bool = False +) -> tuple[int, int, list[tuple[Path, int, str]], list[tuple[Path, int, str]]]: + with open(file_path) as f: + source: str = f.read() + tree = ast.parse(source, filename=str(file_path)) + annotated = 0 + total = 0 + warning_lines: list[tuple[Path, int, str]] = [] + failure_lines: list[tuple[Path, int, str]] = [] + + for node in ast.walk(tree): + if hasattr(node, "lineno") and node.lineno in changed_lines: + if isinstance(node, ast.FunctionDef | ast.Assign | ast.AnnAssign): + total += 1 + result = has_type_annotations(node, debug) + if result == CHECK_SUCCESS or result == CHECK_WARNING: + annotated += 1 + if result == CHECK_WARNING: + warning_lines.append( + (file_path, node.lineno, linecache.getline(str(file_path), node.lineno).strip()) + ) + else: + source_line = linecache.getline(str(file_path), node.lineno).strip() + failure_lines.append((file_path, node.lineno, source_line)) + + return annotated, total, warning_lines, failure_lines + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--threshold", type=float, default=0.3, help="Minimum ratio of annotated lines required (0.0 - 1.0)" + ) + parser.add_argument("--target-file", type=str, default=None, help="Path to the Python source file to analyse") + parser.add_argument( + "--all-lines", + action="store_true", + help="Check all lines in the file instead of only changed lines based on git", + ) + parser.add_argument("--debug", action="store_true", help="Add debugging logs") + args = parser.parse_args() + + total_changed = 0 + total_annotated = 0 + all_warnings: list[tuple[Path, int, str]] = [] + all_failures: list[tuple[Path, int, str]] = [] + + target_files = [args.target_file] if args.target_file is not None else get_changed_files() + for fpath in target_files: + if "tests/" in str(fpath): + continue + if args.all_lines: + changed_lines = [i + 1 for i in range(len(open(fpath).readlines()))] + else: + changed_lines = get_changed_lines(fpath) + annotated, total, warning_lines, failure_lines = check_file(fpath, changed_lines, args.debug) + total_annotated += annotated + total_changed += total + all_warnings.extend(warning_lines) + all_failures.extend(failure_lines) + + ratio = (total_annotated / total_changed) if total_changed else 1.0 + + print( + f"🔍 Type coverage on {'all' if args.all_lines else 'changed'} lines: " + f"{total_annotated}/{total_changed} = {ratio:.2%}. Files inspected: {target_files}" + ) + + if all_warnings: + print("\n⚠️ Suggest Improve: Lines missing type annotations for inputs and outputs:\n") + for fname, lineno, line in all_warnings: + print(f"{fname}:{lineno}: {line}") + + if all_failures: + print("⚠️ [ERROR] Lines missing type annotations for inputs and outputs:\n") + for fname, lineno, line in all_failures: + print(f"{fname}:{lineno}: {line}") + + if ratio < args.threshold: + print( + f"Please add type annotations for inputs and outputs to meet threshold {args.threshold}. " + f"Cases exempt from checking:" + ) + print("1. Private methods.") + print("2. Args with name in ('self', 'cls'), or *args / **kwargs") + print("3. Files under tests/") + raise Exception(f"\n❌ Type coverage below threshold ({args.threshold:.0%}).") + else: + if all_warnings or all_failures: + print("") + print("✅ Type annotation coverage acceptable.\n") + + +if __name__ == "__main__": + main() diff --git a/verl/tests/special_sanity/validate_imported_docs.py b/verl/tests/special_sanity/validate_imported_docs.py new file mode 100644 index 0000000000000000000000000000000000000000..b36a407be77a777cd72a4abf8ce4727d375eb548 --- /dev/null +++ b/verl/tests/special_sanity/validate_imported_docs.py @@ -0,0 +1,130 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +verify_imported_docs.py + +Assert that every function or class *explicitly imported* (via +`from import `) in a given Python file has a docstring. +""" + +from __future__ import annotations + +import argparse +import ast +import importlib +import inspect +import pathlib +import sys + + +def _parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description="Verify that imported functions/classes have docstrings.") + p.add_argument( + "--target-file", + default="verl/trainer/ppo/ray_trainer.py", + help="Path to the Python source file to analyse (e.g. verl/trainer/ppo/ray_trainer.py)", + ) + p.add_argument( + "--allow-list", + default=["omegaconf.open_dict"], + help="a list of third_party dependencies that do not have proper docs :(", + ) + p.add_argument( + "--project-root", + default=".", + help="Directory to prepend to PYTHONPATH so local packages resolve (default: .)", + ) + p.add_argument( + "--quiet", + action="store_true", + help="Suppress success message (still prints errors).", + ) + return p.parse_args() + + +def _import_attr(module_name: str, attr_name: str): + """Import `module_name` then return `getattr(module, attr_name)`.""" + module = importlib.import_module(module_name) + return getattr(module, attr_name) + + +def _check_file(py_file: pathlib.Path, project_root: pathlib.Path, allow_list: list[str]) -> list[str]: + """Return a list of error strings (empty == success).""" + # Ensure local packages resolve + sys.path.insert(0, str(project_root.resolve())) + + tree = ast.parse(py_file.read_text(), filename=str(py_file)) + problems: list[str] = [] + + for node in ast.walk(tree): + if not isinstance(node, ast.ImportFrom): + continue + + # Relative imports (level > 0) get the leading dots stripped + module_name = "." * node.level + (node.module or "") + for alias in node.names: + if alias.name == "*": + problems.append( + f"{py_file}:{node.lineno} - wildcard import `from {module_name} import *` cannot be verified." + ) + continue + + imported_name = alias.name + + try: + obj = _import_attr(module_name, imported_name) + except Exception: # pragma: no cover – wide net for import quirks + pass + # For some reason the module cannot be imported, skip for now + # problems.append( + # f"{py_file}:{node.lineno} - could not resolve " + # f"`{imported_name}` from `{module_name}` ({exc})" + # ) + continue + + if f"{module_name}.{imported_name}" in allow_list: + continue + if inspect.isfunction(obj) or inspect.isclass(obj): + doc = inspect.getdoc(obj) + if not (doc and doc.strip()): + kind = "class" if inspect.isclass(obj) else "function" + problems.append( + f"{py_file}:{node.lineno} - {kind} `{module_name}.{imported_name}` is missing a docstring." + ) + + return problems + + +def main() -> None: + args = _parse_args() + target_path = pathlib.Path(args.target_file).resolve() + project_root = pathlib.Path(args.project_root).resolve() + + if not target_path.is_file(): + raise Exception(f"❌ Target file not found: {target_path}") + + errors = _check_file(target_path, project_root, args.allow_list) + + if errors: + print("Docstring verification failed:\n") + print("\n".join(f" • {e}" for e in errors)) + raise Exception("❌ Docstring verification failed.") + + if not args.quiet: + print(f"✅ All explicitly imported functions/classes in {target_path} have docstrings.") + + +if __name__ == "__main__": + main() diff --git a/verl/tests/special_sanity/validate_structure.py b/verl/tests/special_sanity/validate_structure.py new file mode 100644 index 0000000000000000000000000000000000000000..56136b206374ceff9c566aa1cd88d5be30f8c73b --- /dev/null +++ b/verl/tests/special_sanity/validate_structure.py @@ -0,0 +1,122 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +#!/usr/bin/env python3 +""" +Validate that test file subfolders mirror the top-level package layout. + +Usage examples +-------------- + +# Typical run (defaults: impl_root=my_project, tests_root=tests) +python check_tests_structure.py + +# Custom layout and extra allowed folders +python check_tests_structure.py \ + --impl-root verl \ + --tests-root tests \ + --allow-dirs special_e2e special_sanity special_standalone special_distributed +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + + +def discover_allowed_modules(impl_root: Path, extra: list[str]) -> set[str]: + """Return the set of first-level directories that tests may live under.""" + allowed = {p.name for p in impl_root.iterdir() if p.is_dir()} + allowed.update(extra) + return allowed + + +def find_violations(tests_root: Path, allowed: set[str], allowed_files: list[str]) -> list[str]: + """Return a list of error strings for test files in the wrong place.""" + errors: list[str] = [] + for test_file in tests_root.rglob("test*.py"): + if str(test_file) in allowed_files: + continue + rel_parts = test_file.relative_to(tests_root).parts + if len(rel_parts) < 2: + errors.append(f"{test_file}: must be inside one of {sorted(allowed)} (not at tests root)") + continue + + first_folder = rel_parts[0] + if first_folder not in allowed: + errors.append( + f"{test_file}: subfolder '{first_folder}' under tests/ is not an allowed module. " + f"The valid ones are: {sorted(allowed)}" + ) + return errors + + +def main() -> None: + parser = argparse.ArgumentParser(description="Check that test files follow tests//… layout.") + parser.add_argument( + "--impl-root", + type=Path, + default="verl", + help="Implementation root (default: my_project)", + ) + parser.add_argument( + "--tests-root", + type=Path, + default="tests", + help="Root of test tree (default: tests)", + ) + parser.add_argument( + "--allow-dirs", + nargs="*", + default=["special_e2e", "special_sanity", "special_standalone", "special_distributed"], + help="Extra top-level test folders that are exempt from the rule", + ) + parser.add_argument( + "--allow-files", + nargs="*", + default=[ + "tests/test_protocol_on_cpu.py", + "tests/test_base_config_on_cpu.py", + "tests/test_protocol_v2_on_cpu.py", + ], + help="Extra top-level test folders that are exempt from the rule", + ) + args = parser.parse_args() + + if not args.impl_root.is_dir(): + raise Exception(f"Implementation root '{args.impl_root}' does not exist.") + if not args.tests_root.is_dir(): + raise Exception(f"Tests root '{args.tests_root}' does not exist.") + + allowed = discover_allowed_modules(args.impl_root, args.allow_dirs) + violations = find_violations(args.tests_root, allowed, args.allow_files) + + if violations: + print("❌ Test layout violations found:\n", file=sys.stderr) + for err in violations: + print(" -", err, file=sys.stderr) + + print( + f"\nGuideline:\n Place each test file under tests//…\n where is " + f"one of the top-level packages inside '{args.impl_root}', or is explicitly listed via --allow-dirs.\n", + file=sys.stderr, + ) + raise Exception("❌ Test layout violations found.") + + print("✅ Tests folder structure looks good.") + + +if __name__ == "__main__": + main() diff --git a/verl/tests/special_standalone/README.md b/verl/tests/special_standalone/README.md new file mode 100644 index 0000000000000000000000000000000000000000..0e3596e1afa9a507c67b6949479d1c254b30aec3 --- /dev/null +++ b/verl/tests/special_standalone/README.md @@ -0,0 +1 @@ +The standalone test folder is reserved for tests that require dedicated environment (e.g. memory stress tests) diff --git a/verl/tests/special_standalone/test_memory_buffers.py b/verl/tests/special_standalone/test_memory_buffers.py new file mode 100644 index 0000000000000000000000000000000000000000..77851534782c7d0f5b9ec93231fde8d4d5e60bb6 --- /dev/null +++ b/verl/tests/special_standalone/test_memory_buffers.py @@ -0,0 +1,66 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Test memory buffers +- We start with two models with the same weights +- We use Memory buffer to make one of the models and then compare the parameters +""" + +import gc + +import torch +from transformers import LlamaConfig, LlamaModel + + +def test_memory_buffers(): + llama_config = LlamaConfig( + vocab_size=256, + hidden_size=4096, + intermediate_size=11008, + num_hidden_layers=2, + num_attention_heads=16, + num_key_value_heads=16, + ) + + model = LlamaModel(config=llama_config).cuda() + model_copy = LlamaModel(config=llama_config).cuda() + model_copy.load_state_dict(model.state_dict()) + + norm_factor = 1024**3 + + t_before = torch.cuda.get_device_properties(0).total_memory / norm_factor + r_before = torch.cuda.memory_reserved(0) / norm_factor + a_before = torch.cuda.memory_allocated(0) / norm_factor + + print(f"Before Total memory: {t_before} GB, reserved: {r_before} GB, allocated: {a_before} GB") + + t = torch.cuda.get_device_properties(0).total_memory / norm_factor + r = torch.cuda.memory_reserved(0) / norm_factor + a = torch.cuda.memory_allocated(0) / norm_factor + + gc.collect() + torch.cuda.empty_cache() + + print(f"After Total memory: {t} GB, reserved: {r} GB, allocated: {a} GB") + + change_ratio = (a - a_before) / a_before + assert change_ratio < 0.01, f"make sure the allocated change is less than 1%, Got {change_ratio}" + + for (name1, param1), (name2, param2) in zip(model.named_parameters(), model_copy.named_parameters(), strict=True): + assert name1 == name2 + assert torch.eq(param1.data, param2.data).all(), f"{param1.data}, {param2.data}, {name1}" + + +if __name__ == "__main__": + test_memory_buffers() diff --git a/verl/tests/test_base_config_on_cpu.py b/verl/tests/test_base_config_on_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..9a50235c8ffa736551781d50cf5c937ce21afce0 --- /dev/null +++ b/verl/tests/test_base_config_on_cpu.py @@ -0,0 +1,42 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from verl.base_config import BaseConfig + + +@pytest.fixture +def base_config_mock(): + """Fixture to create a mock BaseConfig instance with test attributes.""" + mock_config = BaseConfig() + mock_config.test_attr = "test_value" + return mock_config + + +def test_getitem_success(base_config_mock): + """Test __getitem__ with existing attribute (happy path).""" + assert base_config_mock["test_attr"] == "test_value" + + +def test_getitem_nonexistent_attribute(base_config_mock): + """Test __getitem__ with non-existent attribute (exception path 1).""" + with pytest.raises(AttributeError): + _ = base_config_mock["nonexistent_attr"] + + +def test_getitem_invalid_key_type(base_config_mock): + """Test __getitem__ with invalid key type (exception path 2).""" + with pytest.raises(TypeError): + _ = base_config_mock[123] # type: ignore diff --git a/verl/tests/test_protocol_on_cpu.py b/verl/tests/test_protocol_on_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..800d428639239a1b2fc4de3125371d657213ce8b --- /dev/null +++ b/verl/tests/test_protocol_on_cpu.py @@ -0,0 +1,1234 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import random + +import numpy as np +import pytest +import tensordict +import torch +from packaging.version import parse as parse_version +from tensordict import TensorDict + +from verl import DataProto +from verl.protocol import ( + deserialize_single_tensor, + deserialize_tensordict, + serialize_single_tensor, + serialize_tensordict, + union_numpy_dict, + union_tensor_dict, +) +from verl.utils import tensordict_utils as tu + + +def test_union_tensor_dict(): + obs = torch.randn(100, 10) + + data1 = TensorDict({"obs": obs, "act": torch.randn(100, 3)}, batch_size=[100]) + data2 = TensorDict({"obs": obs, "next_obs": torch.randn(100, 10), "rew": torch.randn(100)}, batch_size=[100]) + + data_with_copied_obs = TensorDict( + {"obs": obs.clone(), "next_obs": torch.randn(100, 10), "rew": torch.randn(100)}, batch_size=[100] + ) + + union_tensor_dict(data1, data2) + with pytest.raises(AssertionError): + union_tensor_dict(data1, data_with_copied_obs) + + +def test_union_numpy_dict(): + """ + A comprehensive test suite for union_numpy_dict, covering standard use + cases, N-dimensional arrays, object-dtype arrays, and NaN value handling. + """ + arr_3d = np.arange(8).reshape((2, 2, 2)) + union_numpy_dict({"a": arr_3d}, {"a": arr_3d}) + arr1 = np.array([1, "hello", np.array([2, 3])], dtype=object) + arr2 = np.array([1, "hello", np.array([2, 3])], dtype=object) + union_numpy_dict({"a": arr1}, {"a": arr2}) + # --- Test Case 1: The original test with mixed object/float types --- + # This test case from the original test file is preserved. + data = np.random.random(100) + # This array intentionally mixes float('nan') and the string 'nan' + nan_data = [float("nan") for _ in range(99)] + nan_data.append("nan") + nan_data_arr = np.array(nan_data, dtype=object) + + dict1 = {"a": data, "b": nan_data_arr} + dict2_same = {"a": data.copy(), "b": nan_data_arr.copy()} + dict3_different = {"a": np.random.random(100)} + + union_numpy_dict(dict1, dict2_same) # Should pass + with pytest.raises(AssertionError): + union_numpy_dict(dict1, dict3_different) + + # --- Test Case 2: Standard 3D arrays (fixes the core bug) --- + arr_3d = np.arange(24, dtype=np.int32).reshape((2, 3, 4)) + dict_3d_1 = {"nd_array": arr_3d} + dict_3d_2_same = {"nd_array": arr_3d.copy()} + dict_3d_3_different = {"nd_array": arr_3d + 1} + + union_numpy_dict(dict_3d_1, dict_3d_2_same) # Should pass + with pytest.raises(AssertionError, match="`nd_array` in tensor_dict1 and tensor_dict2 are not the same object."): + union_numpy_dict(dict_3d_1, dict_3d_3_different) + + # --- Test Case 3: Nested 2D and 4D object-dtype arrays --- + sub_arr1 = np.array([1, 2]) + sub_arr2 = np.array([3.0, 4.0]) + # 2D object array + arr_2d_obj = np.array([[sub_arr1, "text"], [sub_arr2, None]], dtype=object) + arr_2d_obj_diff = np.array([[sub_arr1, "text"], [sub_arr2, "other"]], dtype=object) + + union_numpy_dict({"data": arr_2d_obj}, {"data": arr_2d_obj.copy()}) # Should pass + with pytest.raises(AssertionError): + union_numpy_dict({"data": arr_2d_obj}, {"data": arr_2d_obj_diff}) + + # 4D object array to ensure deep recursion is robust + arr_4d_obj = np.array([[[[sub_arr1]]], [[[sub_arr2]]]], dtype=object) + arr_4d_obj_diff = np.array([[[[sub_arr1]]], [[[np.array([9, 9])]]]], dtype=object) + + union_numpy_dict({"data": arr_4d_obj}, {"data": arr_4d_obj.copy()}) # Should pass + with pytest.raises(AssertionError): + union_numpy_dict({"data": arr_4d_obj}, {"data": arr_4d_obj_diff}) + + # --- Test Case 4: Explicit NaN value comparison --- + # This verifies that our new _deep_equal logic correctly handles NaNs. + nan_arr = np.array([1.0, np.nan, 3.0]) + dict_nan_1 = {"data": nan_arr} + dict_nan_2_same = {"data": np.array([1.0, np.nan, 3.0])} # A new array with same values + dict_nan_3_different_val = {"data": np.array([1.0, 2.0, 3.0])} + dict_nan_4_different_pos = {"data": np.array([np.nan, 1.0, 3.0])} + + # NaNs in the same position should be considered equal for merging. + union_numpy_dict(dict_nan_1, dict_nan_2_same) # Should pass + + with pytest.raises(AssertionError): + union_numpy_dict(dict_nan_1, dict_nan_3_different_val) + with pytest.raises(AssertionError): + union_numpy_dict(dict_nan_1, dict_nan_4_different_pos) + + # --- Test Case 5: Circular reference handling --- + # Create two separate, but structurally identical, circular references. + # This should pass without a RecursionError. + circ_arr_1 = np.array([None], dtype=object) + circ_arr_1[0] = circ_arr_1 + + circ_arr_2 = np.array([None], dtype=object) + circ_arr_2[0] = circ_arr_2 + + union_numpy_dict({"data": circ_arr_1}, {"data": circ_arr_2}) # Should pass + + # Create a circular reference and a non-circular one. + # This should fail with an AssertionError because they are different. + non_circ_arr = np.array([None], dtype=object) + + with pytest.raises(AssertionError): + union_numpy_dict({"data": circ_arr_1}, {"data": non_circ_arr}) + + +def test_tensor_dict_constructor(): + obs = torch.randn(100, 10) + act = torch.randn(100, 10, 3) + data = DataProto.from_dict(tensors={"obs": obs, "act": act}) + + assert data.batch.batch_size == torch.Size([100]) + + with pytest.raises(AssertionError): + data = DataProto.from_dict(tensors={"obs": obs, "act": act}, num_batch_dims=2) + + with pytest.raises(AssertionError): + data = DataProto.from_dict(tensors={"obs": obs, "act": act}, num_batch_dims=3) + + +def test_tensor_dict_make_iterator(): + obs = torch.randn(100, 10) + labels = [random.choice(["abc", "cde"]) for _ in range(100)] + dataset = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"labels": labels}) + + data_iter_1 = dataset.make_iterator(mini_batch_size=10, epochs=2, seed=1) + data_list_1 = [] + for data in data_iter_1: + data_list_1.append(data) + + data_iter_2 = dataset.make_iterator(mini_batch_size=10, epochs=2, seed=1) + data_list_2 = [] + for data in data_iter_2: + data_list_2.append(data) + + for data1, data2 in zip(data_list_1, data_list_2, strict=True): + assert isinstance(data1, DataProto) + assert isinstance(data2, DataProto) + result = torch.all(torch.eq(data1.batch["obs"], data2.batch["obs"])) + if not result.item(): + print(data1.batch["obs"]) + print(data2.batch["obs"]) + raise AssertionError() + non_tensor_result = np.all(np.equal(data1.non_tensor_batch["labels"], data2.non_tensor_batch["labels"])) + if not non_tensor_result.item(): + print(data1.non_tensor_batch["labels"]) + print(data2.non_tensor_batch["labels"]) + + +def test_reorder(): + obs = torch.tensor([1, 2, 3, 4, 5, 6]) + labels = ["a", "b", "c", "d", "e", "f"] + data = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"labels": labels}, meta_info={"name": "abdce"}) + data.reorder(torch.tensor([3, 4, 2, 0, 1, 5])) + + assert torch.all(torch.eq(data.batch["obs"], torch.tensor([4, 5, 3, 1, 2, 6]))) + assert np.all(data.non_tensor_batch["labels"] == np.array(["d", "e", "c", "a", "b", "f"])) + assert data.meta_info == {"name": "abdce"} + + +def test_chunk_concat(): + obs = torch.tensor([1, 2, 3, 4, 5, 6]) + labels = ["a", "b", "c", "d", "e", "f"] + data = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"labels": labels}, meta_info={"name": "abdce"}) + + with pytest.raises(AssertionError): + data.chunk(5) + + data_split = data.chunk(2) + assert len(data_split) == 2 + assert torch.all(torch.eq(data_split[0].batch["obs"], torch.tensor([1, 2, 3]))) + assert np.all(data_split[0].non_tensor_batch["labels"] == np.array(["a", "b", "c"])) + assert data_split[0].meta_info == {"name": "abdce"} + + assert torch.all(torch.eq(data_split[1].batch["obs"], torch.tensor([4, 5, 6]))) + assert np.all(data_split[1].non_tensor_batch["labels"] == np.array(["d", "e", "f"])) + assert data_split[1].meta_info == {"name": "abdce"} + + concat_data = DataProto.concat(data_split) + assert torch.all(torch.eq(concat_data.batch["obs"], data.batch["obs"])) + assert np.all(concat_data.non_tensor_batch["labels"] == data.non_tensor_batch["labels"]) + assert concat_data.meta_info == data.meta_info + + +def test_concat_metrics_from_multiple_workers(): + """Test that concat() properly merges metrics from all workers in distributed training.""" + # Simulate 3 workers each with their own metrics + obs1 = torch.tensor([1, 2]) + obs2 = torch.tensor([3, 4]) + obs3 = torch.tensor([5, 6]) + + # Each worker has different metrics (as list of dict format) + worker1_metrics = [{"loss": 0.5, "accuracy": 0.9}] + worker2_metrics = [{"loss": 0.6, "accuracy": 0.85}] + worker3_metrics = [{"loss": 0.55, "accuracy": 0.88}] + + data1 = DataProto.from_dict(tensors={"obs": obs1}, meta_info={"metrics": worker1_metrics, "config_flag": True}) + data2 = DataProto.from_dict(tensors={"obs": obs2}, meta_info={"metrics": worker2_metrics, "config_flag": True}) + data3 = DataProto.from_dict(tensors={"obs": obs3}, meta_info={"metrics": worker3_metrics, "config_flag": True}) + + # Concat all workers' data + concat_data = DataProto.concat([data1, data2, data3]) + + # Verify tensors are concatenated + assert torch.all(torch.eq(concat_data.batch["obs"], torch.tensor([1, 2, 3, 4, 5, 6]))) + + # Verify ALL workers' metrics are flattened to dict of lists + expected_metrics = {"loss": [0.5, 0.6, 0.55], "accuracy": [0.9, 0.85, 0.88]} + assert concat_data.meta_info["metrics"] == expected_metrics + + # Verify config flags are preserved from first worker + assert concat_data.meta_info["config_flag"] is True + + +def test_concat_with_empty_and_non_list_meta_info(): + """Test concat() handles edge cases: empty meta_info, non-list values, and None.""" + obs1 = torch.tensor([1, 2]) + obs2 = torch.tensor([3, 4]) + + # Worker 1 has metrics, worker 2 doesn't + data1 = DataProto.from_dict(tensors={"obs": obs1}, meta_info={"metrics": [{"loss": 0.5}], "flag": True}) + data2 = DataProto.from_dict(tensors={"obs": obs2}, meta_info={"flag": True}) + + concat_data = DataProto.concat([data1, data2]) + + # Should flatten worker1's metrics to dict of lists + assert concat_data.meta_info["metrics"] == {"loss": [0.5]} + assert concat_data.meta_info["flag"] is True + + # Test with non-list meta_info value + data3 = DataProto.from_dict(tensors={"obs": obs1}, meta_info={"single_value": 42}) + data4 = DataProto.from_dict(tensors={"obs": obs2}, meta_info={"single_value": 42}) + + concat_data2 = DataProto.concat([data3, data4]) + assert concat_data2.meta_info["single_value"] == 42 + + +def test_concat_first_worker_missing_metrics(): + """Test that metrics from other workers are preserved even when first worker has no metrics. + + This is a critical edge case - the old buggy implementation only checked data[0].meta_info + and would lose all metrics if the first worker didn't have any. + """ + obs1 = torch.tensor([1, 2]) + obs2 = torch.tensor([3, 4]) + obs3 = torch.tensor([5, 6]) + + # First worker has NO metrics, but workers 2 and 3 do + data1 = DataProto.from_dict(tensors={"obs": obs1}, meta_info={"config_flag": True}) + data2 = DataProto.from_dict(tensors={"obs": obs2}, meta_info={"metrics": {"loss": 0.6}, "config_flag": True}) + data3 = DataProto.from_dict(tensors={"obs": obs3}, meta_info={"metrics": {"loss": 0.55}, "config_flag": True}) + + concat_data = DataProto.concat([data1, data2, data3]) + + # Should flatten metrics from workers 2 and 3 into dict of lists + expected_metrics = {"loss": [0.6, 0.55]} + assert concat_data.meta_info["metrics"] == expected_metrics + assert concat_data.meta_info["config_flag"] is True + + +def test_concat_non_list_metrics(): + """Test that concat() handles non-list metrics (single dict) correctly. + + In some cases, metrics might be a single dict instead of a list. + The implementation should flatten them into a dict of lists. + """ + obs1 = torch.tensor([1, 2]) + obs2 = torch.tensor([3, 4]) + + # Metrics as single dict (not wrapped in list) + data1 = DataProto.from_dict(tensors={"obs": obs1}, meta_info={"metrics": {"loss": 0.5, "accuracy": 0.9}}) + data2 = DataProto.from_dict(tensors={"obs": obs2}, meta_info={"metrics": {"loss": 0.6, "accuracy": 0.85}}) + + concat_data = DataProto.concat([data1, data2]) + + # Should flatten to dict of lists + expected_metrics = {"loss": [0.5, 0.6], "accuracy": [0.9, 0.85]} + assert concat_data.meta_info["metrics"] == expected_metrics + + +def test_concat_merge_different_non_metric_keys(): + """Test that concat() merges non-metric meta_info keys from all workers. + + When different workers have different non-metric keys, all keys should be preserved. + This prevents silent data loss and aligns with the docstring stating meta_info is "merged". + """ + obs1 = torch.tensor([1, 2]) + obs2 = torch.tensor([3, 4]) + obs3 = torch.tensor([5, 6]) + + # Each worker has some unique non-metric keys + data1 = DataProto.from_dict(tensors={"obs": obs1}, meta_info={"config": "A", "shared_key": "X"}) + data2 = DataProto.from_dict(tensors={"obs": obs2}, meta_info={"extra_key": "B", "shared_key": "X"}) + data3 = DataProto.from_dict(tensors={"obs": obs3}, meta_info={"another_key": "C", "shared_key": "X"}) + + concat_data = DataProto.concat([data1, data2, data3]) + + # All unique keys should be preserved + assert concat_data.meta_info["config"] == "A" + assert concat_data.meta_info["extra_key"] == "B" + assert concat_data.meta_info["another_key"] == "C" + assert concat_data.meta_info["shared_key"] == "X" + + +def test_concat_conflicting_non_metric_keys(): + """Test that concat() raises an assertion error when non-metric keys have conflicting values. + + This ensures data integrity by catching cases where workers have different values + for what should be the same configuration parameter. + """ + obs1 = torch.tensor([1, 2]) + obs2 = torch.tensor([3, 4]) + + # Same key "config" but different values + data1 = DataProto.from_dict(tensors={"obs": obs1}, meta_info={"config": "A"}) + data2 = DataProto.from_dict(tensors={"obs": obs2}, meta_info={"config": "B"}) + + # Should raise an assertion error due to conflicting values + with pytest.raises(AssertionError, match="Conflicting values for meta_info key 'config'"): + DataProto.concat([data1, data2]) + + +def test_pop(): + obs = torch.randn(100, 10) + act = torch.randn(100, 3) + dataset = DataProto.from_dict({"obs": obs, "act": act}, meta_info={"2": 2, "1": 1}) + poped_dataset = dataset.pop(batch_keys=["obs"], meta_info_keys=["2"]) + + assert poped_dataset.batch.keys() == {"obs"} + assert poped_dataset.meta_info.keys() == {"2"} + + assert dataset.batch.keys() == {"act"} + assert dataset.meta_info.keys() == {"1"} + + +def test_repeat(): + # Create a DataProto object with some batch and non-tensor data + obs = torch.tensor([[1, 2], [3, 4], [5, 6]]) + labels = ["a", "b", "c"] + data = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"labels": labels}, meta_info={"info": "test_info"}) + + # Test interleave=True + repeated_data_interleave = data.repeat(repeat_times=2, interleave=True) + expected_obs_interleave = torch.tensor([[1, 2], [1, 2], [3, 4], [3, 4], [5, 6], [5, 6]]) + expected_labels_interleave = ["a", "a", "b", "b", "c", "c"] + + assert torch.all(torch.eq(repeated_data_interleave.batch["obs"], expected_obs_interleave)) + assert (repeated_data_interleave.non_tensor_batch["labels"] == expected_labels_interleave).all() + assert repeated_data_interleave.meta_info == {"info": "test_info"} + + # Test interleave=False + repeated_data_no_interleave = data.repeat(repeat_times=2, interleave=False) + expected_obs_no_interleave = torch.tensor([[1, 2], [3, 4], [5, 6], [1, 2], [3, 4], [5, 6]]) + expected_labels_no_interleave = ["a", "b", "c", "a", "b", "c"] + + assert torch.all(torch.eq(repeated_data_no_interleave.batch["obs"], expected_obs_no_interleave)) + assert (repeated_data_no_interleave.non_tensor_batch["labels"] == expected_labels_no_interleave).all() + assert repeated_data_no_interleave.meta_info == {"info": "test_info"} + + +def test_dataproto_pad_unpad(): + obs = torch.tensor([[1, 2], [3, 4], [5, 6]]) + labels = ["a", "b", "c"] + data = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"labels": labels}, meta_info={"info": "test_info"}) + + from verl.protocol import pad_dataproto_to_divisor, unpad_dataproto + + padded_data, pad_size = pad_dataproto_to_divisor(data, size_divisor=2) + assert pad_size == 1 + + expected_obs = torch.tensor([[1, 2], [3, 4], [5, 6], [1, 2]]) + expected_labels = ["a", "b", "c", "a"] + + assert torch.all(torch.eq(padded_data.batch["obs"], expected_obs)) + assert (padded_data.non_tensor_batch["labels"] == expected_labels).all() + assert padded_data.meta_info == {"info": "test_info"} + + unpadd_data = unpad_dataproto(padded_data, pad_size=pad_size) + assert torch.all(torch.eq(unpadd_data.batch["obs"], obs)) + assert (unpadd_data.non_tensor_batch["labels"] == labels).all() + assert unpadd_data.meta_info == {"info": "test_info"} + + padded_data, pad_size = pad_dataproto_to_divisor(data, size_divisor=3) + assert pad_size == 0 + + expected_obs = torch.tensor([[1, 2], [3, 4], [5, 6]]) + expected_labels = ["a", "b", "c"] + + assert torch.all(torch.eq(padded_data.batch["obs"], expected_obs)) + assert (padded_data.non_tensor_batch["labels"] == expected_labels).all() + assert padded_data.meta_info == {"info": "test_info"} + + unpadd_data = unpad_dataproto(padded_data, pad_size=pad_size) + assert torch.all(torch.eq(unpadd_data.batch["obs"], obs)) + assert (unpadd_data.non_tensor_batch["labels"] == labels).all() + assert unpadd_data.meta_info == {"info": "test_info"} + + padded_data, pad_size = pad_dataproto_to_divisor(data, size_divisor=7) + assert pad_size == 4 + + expected_obs = torch.tensor([[1, 2], [3, 4], [5, 6], [1, 2], [3, 4], [5, 6], [1, 2]]) + expected_labels = ["a", "b", "c", "a", "b", "c", "a"] + assert torch.all(torch.eq(padded_data.batch["obs"], expected_obs)) + assert (padded_data.non_tensor_batch["labels"] == expected_labels).all() + assert padded_data.meta_info == {"info": "test_info"} + + unpadd_data = unpad_dataproto(padded_data, pad_size=pad_size) + assert torch.all(torch.eq(unpadd_data.batch["obs"], obs)) + assert (unpadd_data.non_tensor_batch["labels"] == labels).all() + assert unpadd_data.meta_info == {"info": "test_info"} + + +def test_dataproto_fold_unfold(): + from verl.protocol import DataProto, fold_batch_dim, unfold_batch_dim + + obs = torch.tensor([[1, 2], [3, 4], [5, 6]]) + labels = ["a", "b", "c"] + data = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"labels": labels}, meta_info={"info": "test_info"}) + + data1 = data.repeat(repeat_times=2, interleave=True) + + data2 = fold_batch_dim(data1, new_batch_size=3) + + torch.testing.assert_close(data2.batch["obs"], torch.tensor([[[1, 2], [1, 2]], [[3, 4], [3, 4]], [[5, 6], [5, 6]]])) + assert (data2.non_tensor_batch["labels"] == [["a", "a"], ["b", "b"], ["c", "c"]]).all() + + data2.reorder(indices=torch.tensor([1, 2, 0])) + + data3 = unfold_batch_dim(data2, batch_dims=2) + + torch.testing.assert_close(data3.batch["obs"], torch.tensor([[3, 4], [3, 4], [5, 6], [5, 6], [1, 2], [1, 2]])) + assert (data3.non_tensor_batch["labels"] == ["b", "b", "c", "c", "a", "a"]).all() + assert data3.meta_info == {"info": "test_info"} + + +def test_torch_save_data_proto(): + obs = torch.tensor([[1, 2], [3, 4], [5, 6]]) + labels = ["a", "b", "c"] + data = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"labels": labels}, meta_info={"info": "test_info"}) + data.save_to_disk("test_data.pt") + loaded_data = DataProto.load_from_disk("test_data.pt") + + assert torch.all(torch.eq(loaded_data.batch["obs"], data.batch["obs"])) + assert (loaded_data.non_tensor_batch["labels"] == data.non_tensor_batch["labels"]).all() + assert loaded_data.meta_info == data.meta_info + + import os + + os.remove("test_data.pt") + + +def test_len(): + obs = torch.tensor([[1, 2], [3, 4], [5, 6]]) + labels = np.array(["a", "b", "c"], dtype=object) + data = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"labels": labels}, meta_info={"info": "test_info"}) + + assert len(data) == 3 + + data = DataProto(batch=None, non_tensor_batch={"labels": labels}, meta_info={"info": "test_info"}) + + assert len(data) == 3 + + data = DataProto(batch=None, non_tensor_batch={}, meta_info={"info": "test_info"}) + + assert len(data) == 0 + + data = DataProto(batch=None, non_tensor_batch=None, meta_info={"info": "test_info"}) + + assert len(data) == 0 + + +def test_dataproto_index(): + data_len = 100 + idx_num = 10 + + obs = torch.randn(data_len, 10) + labels = [random.choice(["abc", "cde"]) for _ in range(data_len)] + data = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"labels": labels}) + labels_np = np.array(labels) + + idx_np_int = np.random.randint(0, data_len, size=(idx_num,)) + result_np_int = data[idx_np_int] + assert result_np_int.batch.keys() == data.batch.keys() + assert result_np_int.non_tensor_batch.keys() == data.non_tensor_batch.keys() + assert result_np_int.batch["obs"].shape[0] == idx_num + assert result_np_int.non_tensor_batch["labels"].shape[0] == idx_num + assert np.array_equal(result_np_int.batch["obs"].cpu().numpy(), obs[idx_np_int].numpy()) + assert np.array_equal(result_np_int.non_tensor_batch["labels"], labels_np[idx_np_int]) + + idx_torch_int = torch.randint(0, data_len, size=(idx_num,)) + result_torch_int = data[idx_torch_int] + assert result_torch_int.batch.keys() == data.batch.keys() + assert result_torch_int.non_tensor_batch.keys() == data.non_tensor_batch.keys() + assert result_torch_int.batch["obs"].shape[0] == idx_num + assert result_torch_int.non_tensor_batch["labels"].shape[0] == idx_num + assert np.array_equal(result_torch_int.batch["obs"].cpu().numpy(), obs[idx_torch_int].cpu().numpy()) + assert np.array_equal(result_torch_int.non_tensor_batch["labels"], labels_np[idx_torch_int.cpu().numpy()]) + + idx_list_int = [np.random.randint(0, data_len) for _ in range(idx_num)] + result_list_int = data[idx_list_int] + assert result_list_int.batch.keys() == data.batch.keys() + assert result_list_int.non_tensor_batch.keys() == data.non_tensor_batch.keys() + assert result_list_int.batch["obs"].shape[0] == idx_num + assert result_list_int.non_tensor_batch["labels"].shape[0] == idx_num + assert np.array_equal(result_list_int.batch["obs"].cpu().numpy(), obs[idx_list_int].cpu().numpy()) + assert np.array_equal(result_list_int.non_tensor_batch["labels"], labels_np[idx_list_int]) + + idx_np_bool = np.random.randint(0, 2, size=(data_len,), dtype=bool) + result_np_bool = data[idx_np_bool] + assert result_np_bool.batch.keys() == data.batch.keys() + assert result_np_bool.non_tensor_batch.keys() == data.non_tensor_batch.keys() + assert result_np_bool.batch["obs"].shape[0] == idx_np_bool.sum() + assert result_np_bool.non_tensor_batch["labels"].shape[0] == idx_np_bool.sum() + assert np.array_equal(result_np_bool.batch["obs"].cpu().numpy(), obs[idx_np_bool].cpu().numpy()) + assert np.array_equal(result_np_bool.non_tensor_batch["labels"], labels_np[idx_np_bool]) + + idx_torch_bool = torch.randint(0, 2, size=(data_len,), dtype=torch.bool) + result_torch_bool = data[idx_torch_bool] + assert result_torch_bool.batch.keys() == data.batch.keys() + assert result_torch_bool.non_tensor_batch.keys() == data.non_tensor_batch.keys() + assert result_torch_bool.batch["obs"].shape[0] == idx_torch_bool.sum().item() + assert result_torch_bool.non_tensor_batch["labels"].shape[0] == idx_torch_bool.sum().item() + assert np.array_equal(result_torch_bool.batch["obs"].cpu().numpy(), obs[idx_torch_bool].cpu().numpy()) + assert np.array_equal(result_torch_bool.non_tensor_batch["labels"], labels_np[idx_torch_bool]) + + idx_list_bool = [np.random.randint(0, 2, dtype=bool) for _ in range(data_len)] + result_list_bool = data[idx_list_bool] + assert result_list_bool.batch.keys() == data.batch.keys() + assert result_list_bool.non_tensor_batch.keys() == data.non_tensor_batch.keys() + assert result_list_bool.batch["obs"].shape[0] == sum(idx_list_bool) + assert result_list_bool.non_tensor_batch["labels"].shape[0] == sum(idx_list_bool) + assert np.array_equal(result_list_bool.batch["obs"].cpu().numpy(), obs[idx_list_bool].cpu().numpy()) + assert np.array_equal(result_list_bool.non_tensor_batch["labels"], labels_np[idx_list_bool]) + + +def test_old_vs_new_from_single_dict(): + class CustomProto(DataProto): + """Uses the new, fixed from_single_dict.""" + + pass + + class OriginProto(DataProto): + """Mimics the *old* from_single_dict (always returns a DataProto).""" + + @classmethod + def from_single_dict(cls, data, meta_info=None, auto_padding=False): + tensors, non_tensors = {}, {} + for k, v in data.items(): + if torch.is_tensor(v): + tensors[k] = v + else: + non_tensors[k] = v + # always calls DataProto.from_dict, ignoring `cls` + return DataProto.from_dict( + tensors=tensors, + non_tensors=non_tensors, + meta_info=meta_info, + auto_padding=auto_padding, + ) + + sample = {"x": torch.tensor([0])} + + orig = OriginProto.from_single_dict(sample) + # old behavior: always DataProto, not a CustomOriginProto + assert type(orig) is DataProto + assert type(orig) is not OriginProto + + cust = CustomProto.from_single_dict(sample) + # new behavior: respects subclass + assert type(cust) is CustomProto + + +def test_dataproto_no_batch(): + labels = ["a", "b", "c"] + data = DataProto.from_dict(non_tensors={"labels": labels}, meta_info={"info": "test_info"}) + selected = data.select(non_tensor_batch_keys=["labels"]) + assert (selected.non_tensor_batch["labels"] == labels).all() + pop_data = data.pop(non_tensor_batch_keys=["labels"]) + assert (pop_data.non_tensor_batch["labels"] == labels).all() + assert data.non_tensor_batch == {} + + +def test_sample_level_repeat(): + # Create a DataProto object with some batch and non-tensor data + obs = torch.tensor([[1, 2], [3, 4], [5, 6]]) + labels = ["a", "b", "c"] + data = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"labels": labels}, meta_info={"info": "test_info"}) + + # list + repeated_data_interleave = data.sample_level_repeat(repeat_times=[3, 1, 2]) + expected_obs_interleave = torch.tensor([[1, 2], [1, 2], [1, 2], [3, 4], [5, 6], [5, 6]]) + expected_labels_interleave = ["a", "a", "a", "b", "c", "c"] + + assert torch.all(torch.eq(repeated_data_interleave.batch["obs"], expected_obs_interleave)) + assert (repeated_data_interleave.non_tensor_batch["labels"] == expected_labels_interleave).all() + assert repeated_data_interleave.meta_info == {"info": "test_info"} + + # torch.tensor + repeated_data_no_interleave = data.sample_level_repeat(repeat_times=torch.tensor([1, 2, 3])) + expected_obs_no_interleave = torch.tensor([[1, 2], [3, 4], [3, 4], [5, 6], [5, 6], [5, 6]]) + expected_labels_no_interleave = ["a", "b", "b", "c", "c", "c"] + + assert torch.all(torch.eq(repeated_data_no_interleave.batch["obs"], expected_obs_no_interleave)) + assert (repeated_data_no_interleave.non_tensor_batch["labels"] == expected_labels_no_interleave).all() + assert repeated_data_no_interleave.meta_info == {"info": "test_info"} + + +def test_dataproto_unfold_column_chunks(): + obs1 = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]) + obs2 = torch.tensor([[1, 2], [5, 6], [9, 10]]) + + labels = ["a", "b", "c"] + data = DataProto.from_dict( + tensors={"obs1": obs1, "obs2": obs2}, non_tensors={"labels": labels}, meta_info={"name": "abc"} + ) + ret = data.unfold_column_chunks(2, split_keys=["obs1"]) + + expect_obs1 = torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12]]) + expect_obs2 = torch.tensor([[1, 2], [1, 2], [5, 6], [5, 6], [9, 10], [9, 10]]) + expect_labels = ["a", "a", "b", "b", "c", "c"] + assert torch.all(torch.eq(ret.batch["obs1"], expect_obs1)) + assert torch.all(torch.eq(ret.batch["obs2"], expect_obs2)) + assert (ret.non_tensor_batch["labels"] == expect_labels).all() + assert ret.meta_info == {"name": "abc"} + + obs1 = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]) + obs2 = torch.tensor([[1, 2], [5, 6], [9, 10]]) + + labels = [["a1", "a2"], ["b1", "b2"], ["c1", "c2"]] + data = DataProto.from_dict( + tensors={"obs1": obs1, "obs2": obs2}, non_tensors={"labels": labels}, meta_info={"name": "abc"} + ) + ret = data.unfold_column_chunks(2, split_keys=["obs1", "labels"]) + + expect_obs1 = torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12]]) + expect_obs2 = torch.tensor([[1, 2], [1, 2], [5, 6], [5, 6], [9, 10], [9, 10]]) + expect_labels = [["a1"], ["a2"], ["b1"], ["b2"], ["c1"], ["c2"]] + assert torch.all(torch.eq(ret.batch["obs1"], expect_obs1)) + assert torch.all(torch.eq(ret.batch["obs2"], expect_obs2)) + assert (ret.non_tensor_batch["labels"] == expect_labels).all() + assert ret.meta_info == {"name": "abc"} + + obs1 = torch.tensor( + [[[1, 1], [2, 2], [3, 3], [4, 4]], [[5, 5], [6, 6], [7, 7], [8, 8]], [[9, 9], [10, 10], [11, 11], [12, 12]]] + ) + obs2 = torch.tensor([[[1, 1], [2, 2]], [[5, 5], [6, 6]], [[9, 9], [10, 10]]]) + + labels = ["a", "b", "c"] + data = DataProto.from_dict( + tensors={"obs1": obs1, "obs2": obs2}, non_tensors={"labels": labels}, meta_info={"name": "abc"} + ) + ret = data.unfold_column_chunks(2, split_keys=["obs1"]) + + expect_obs1 = torch.tensor( + [ + [[1, 1], [2, 2]], + [[3, 3], [4, 4]], + [[5, 5], [6, 6]], + [[7, 7], [8, 8]], + [[9, 9], [10, 10]], + [[11, 11], [12, 12]], + ] + ) + expect_obs2 = torch.tensor( + [[[1, 1], [2, 2]], [[1, 1], [2, 2]], [[5, 5], [6, 6]], [[5, 5], [6, 6]], [[9, 9], [10, 10]], [[9, 9], [10, 10]]] + ) + expect_labels = ["a", "a", "b", "b", "c", "c"] + assert torch.all(torch.eq(ret.batch["obs1"], expect_obs1)) + assert torch.all(torch.eq(ret.batch["obs2"], expect_obs2)) + assert (ret.non_tensor_batch["labels"] == expect_labels).all() + assert ret.meta_info == {"name": "abc"} + + +def test_dataproto_chunk_after_index(): + data_len = 4 + obs = torch.randn(data_len, 4) + labels = [f"label_{i}" for i in range(data_len)] + data = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"labels": labels}, meta_info={"name": "abc"}) + + # Test with boolean numpy array + bool_mask = np.array([True, False, True, False]) + selected = data[bool_mask] + assert isinstance(selected.batch.batch_size, torch.Size) + assert all(isinstance(d, int) for d in selected.batch.batch_size) # int or List[int] + + # Test with integer numpy array + int_mask = np.array([0, 2]) + selected = data[int_mask] + assert isinstance(selected.batch.batch_size, torch.Size) + assert all(isinstance(d, int) for d in selected.batch.batch_size) + + # Test with boolean list + list_mask = [True, False, True, False] + selected = data[list_mask] + assert isinstance(selected.batch.batch_size, torch.Size) + assert all(isinstance(d, int) for d in selected.batch.batch_size) + + # Test with list + list_mask = [0, 2] + selected = data[list_mask] + assert isinstance(selected.batch.batch_size, torch.Size) + assert all(isinstance(d, int) for d in selected.batch.batch_size) + + # Test with torch tensor (bool) + torch_bool_mask = torch.tensor([True, False, True, False]) + selected = data[torch_bool_mask] + assert isinstance(selected.batch.batch_size, torch.Size) + assert all(isinstance(d, int) for d in selected.batch.batch_size) + + # Test with torch tensor (int) + torch_int_mask = torch.tensor([0, 2]) + selected = data[torch_int_mask] + assert isinstance(selected.batch.batch_size, torch.Size) + assert all(isinstance(d, int) for d in selected.batch.batch_size) + + +@pytest.mark.skipif( + parse_version(tensordict.__version__) < parse_version("0.10"), reason="requires at least tensordict 0.10" +) +def test_to_tensordict(): + obs = torch.tensor([1, 2, 3, 4, 5, 6]) + labels = ["a", "b", "c", "d", "e", "f"] + data = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"labels": labels}, meta_info={"name": "abdce"}) + output = data.to_tensordict() + + assert torch.all(torch.eq(output["obs"], obs)).item() + assert output["labels"] == labels + assert output["name"] == "abdce" + + +@pytest.mark.skipif( + parse_version(tensordict.__version__) < parse_version("0.10"), reason="requires at least tensordict 0.10" +) +def test_from_tensordict(): + tensor_dict = { + "obs": torch.tensor([1, 2, 3, 4, 5, 6]), + "labels": ["a", "b", "c", "d", "e", "f"], + } + non_tensor_dict = {"name": "abdce"} + tensordict = tu.get_tensordict(tensor_dict, non_tensor_dict) + data = DataProto.from_tensordict(tensordict) + + assert data.non_tensor_batch["labels"].tolist() == tensor_dict["labels"] + assert torch.all(torch.eq(data.batch["obs"], tensor_dict["obs"])).item() + assert data.meta_info["name"] == "abdce" + + +@pytest.mark.skipif( + parse_version(tensordict.__version__) < parse_version("0.10"), reason="requires at least tensordict 0.10" +) +def test_to_tensordict_with_nested_lists(): + """Test converting DataProto with nested lists to TensorDict (lists of lists).""" + obs = torch.tensor([1, 2, 3]) + # Simulate turn_scores or tool_rewards: array of lists with varying lengths + turn_scores = [[], [0.5, 0.8], [0.9]] + + data = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"turn_scores": turn_scores}) + + # This should not raise an error + tensordict_output = data.to_tensordict() + + # Verify the data is preserved + assert torch.all(torch.eq(tensordict_output["obs"], obs)).item() + # Verify nested structure is accessible (TensorDict wraps NonTensorStack as LinkedList) + retrieved_scores = tensordict_output["turn_scores"] + assert len(retrieved_scores) == len(turn_scores) + # Verify content matches + assert list(retrieved_scores[0]) == [] + assert list(retrieved_scores[1]) == [0.5, 0.8] + assert list(retrieved_scores[2]) == [0.9] + + +@pytest.mark.skipif( + parse_version(tensordict.__version__) < parse_version("0.10"), reason="requires at least tensordict 0.10" +) +def test_to_tensordict_with_nested_dicts(): + """Test converting DataProto with lists of dicts to TensorDict.""" + obs = torch.tensor([1, 2, 3]) + # Simulate reward_extra_info: array of dicts + reward_extra_info = [{"acc": 1.0}, {"acc": 0.0}, {"acc": 1.0}] + + data = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"reward_extra_info": reward_extra_info}) + + # This should not raise an error - this was the original bug + tensordict_output = data.to_tensordict() + + # Verify the data is preserved + assert torch.all(torch.eq(tensordict_output["obs"], obs)).item() + # Verify nested dicts are accessible + retrieved_info = tensordict_output["reward_extra_info"] + assert len(retrieved_info) == len(reward_extra_info) + # Verify content matches + for i, expected_dict in enumerate(reward_extra_info): + assert dict(retrieved_info[i]) == expected_dict + + +@pytest.mark.skipif( + parse_version(tensordict.__version__) < parse_version("0.10"), reason="requires at least tensordict 0.10" +) +def test_to_tensordict_with_complex_nested_structures(): + """Test converting DataProto with complex nested structures (lists of lists of dicts).""" + obs = torch.tensor([1, 2, 3]) + # Simulate raw_prompt: array of lists containing dicts + raw_prompt = [ + [{"content": "Question 1", "role": "user"}], + [{"content": "Question 2", "role": "user"}, {"content": "Answer 2", "role": "assistant"}], + [{"content": "Question 3", "role": "user"}], + ] + + data = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"raw_prompt": raw_prompt}) + + # This should not raise an error + tensordict_output = data.to_tensordict() + + # Verify the data is preserved + assert torch.all(torch.eq(tensordict_output["obs"], obs)).item() + # Verify complex nested structure is accessible + retrieved_prompt = tensordict_output["raw_prompt"] + assert len(retrieved_prompt) == len(raw_prompt) + # Spot check: verify first prompt has correct structure + assert len(retrieved_prompt[0]) == 1 + assert dict(retrieved_prompt[0][0]) == {"content": "Question 1", "role": "user"} + + +@pytest.mark.skipif( + parse_version(tensordict.__version__) < parse_version("0.10"), reason="requires at least tensordict 0.10" +) +def test_to_tensordict_and_back_with_nested_data(): + """Test round-trip conversion: DataProto → TensorDict → DataProto with nested structures.""" + obs = torch.tensor([1, 2, 3, 4]) + labels = ["a", "b", "c", "d"] + + # Multiple types of nested structures + turn_scores = [[], [0.5], [0.8, 0.9], [0.7]] + reward_extra_info = [ + {"acc": 1.0, "loss": 0.1}, + {"acc": 0.5, "loss": 0.3}, + {"acc": 1.0, "loss": 0.05}, + {"acc": 0.0, "loss": 0.9}, + ] + raw_prompt = [ + [{"content": "Q1", "role": "user"}], + [{"content": "Q2", "role": "user"}], + [{"content": "Q3", "role": "user"}, {"content": "A3", "role": "assistant"}], + [{"content": "Q4", "role": "user"}], + ] + + # Create original DataProto + original_data = DataProto.from_dict( + tensors={"obs": obs}, + non_tensors={ + "labels": labels, + "turn_scores": turn_scores, + "reward_extra_info": reward_extra_info, + "raw_prompt": raw_prompt, + }, + meta_info={"experiment": "test_nested"}, + ) + + # Convert to TensorDict + tensordict_output = original_data.to_tensordict() + + # Convert back to DataProto + reconstructed_data = DataProto.from_tensordict(tensordict_output) + + # Verify tensors are preserved + assert torch.all(torch.eq(reconstructed_data.batch["obs"], obs)).item() + + # Verify non-tensor data is preserved + assert reconstructed_data.non_tensor_batch["labels"].tolist() == labels + + # Verify nested structures are preserved + assert len(reconstructed_data.non_tensor_batch["turn_scores"]) == len(turn_scores) + for orig, recon in zip(turn_scores, reconstructed_data.non_tensor_batch["turn_scores"], strict=True): + assert list(orig) == list(recon) + + assert len(reconstructed_data.non_tensor_batch["reward_extra_info"]) == len(reward_extra_info) + for orig, recon in zip(reward_extra_info, reconstructed_data.non_tensor_batch["reward_extra_info"], strict=True): + assert orig == recon + + assert len(reconstructed_data.non_tensor_batch["raw_prompt"]) == len(raw_prompt) + for orig, recon in zip(raw_prompt, reconstructed_data.non_tensor_batch["raw_prompt"], strict=True): + assert orig == list(recon) + + # Verify meta_info is preserved + assert reconstructed_data.meta_info["experiment"] == "test_nested" + + +@pytest.mark.skipif( + parse_version(tensordict.__version__) < parse_version("0.10"), reason="requires at least tensordict 0.10" +) +def test_to_tensordict_agent_loop_scenario(): + """Test the exact scenario from agent loop: DataProto with tool rewards, acc, etc. + + This test reproduces the exact error from the agent loop where nested structures + (lists of lists, lists of dicts) failed to convert to TensorDict. + """ + # Simulate real agent loop data structure + prompts = torch.tensor([[1, 2, 3], [4, 5, 6]]) + responses = torch.tensor([[7, 8], [9, 10]]) + + # Non-tensor data with nested structures from agent loop + data_source = ["lighteval/MATH", "lighteval/MATH"] + uid = ["uuid-1", "uuid-2"] + num_turns = np.array([2, 4], dtype=np.int32) + acc = np.array([1.0, 0.0]) + turn_scores = [[], [0.5, 0.8]] # Lists of varying lengths + reward_extra_info = [{"acc": 1.0}, {"acc": 0.0}] # List of dicts + raw_prompt = [ + [{"content": "Compute 4 @ 2", "role": "user"}], + [{"content": "Compute 8 @ 7", "role": "user"}], + ] + tool_rewards = [[0.0], []] # List of lists + + data = DataProto.from_dict( + tensors={"prompts": prompts, "responses": responses}, + non_tensors={ + "data_source": data_source, + "uid": uid, + "num_turns": num_turns, + "acc": acc, + "turn_scores": turn_scores, + "reward_extra_info": reward_extra_info, + "raw_prompt": raw_prompt, + "tool_rewards": tool_rewards, + }, + meta_info={"global_steps": 42}, + ) + + # THE KEY TEST: This should not raise ValueError about TensorDict conversion + tensordict_output = data.to_tensordict() + + # Verify tensors are accessible + assert torch.all(torch.eq(tensordict_output["prompts"], prompts)).item() + assert torch.all(torch.eq(tensordict_output["responses"], responses)).item() + + # Verify all nested structures are accessible (content check, not type check) + assert len(tensordict_output["turn_scores"]) == 2 + assert list(tensordict_output["turn_scores"][0]) == [] + assert list(tensordict_output["turn_scores"][1]) == [0.5, 0.8] + + assert len(tensordict_output["reward_extra_info"]) == 2 + assert dict(tensordict_output["reward_extra_info"][0]) == {"acc": 1.0} + + assert len(tensordict_output["raw_prompt"]) == 2 + assert dict(tensordict_output["raw_prompt"][0][0]) == {"content": "Compute 4 @ 2", "role": "user"} + + assert len(tensordict_output["tool_rewards"]) == 2 + assert list(tensordict_output["tool_rewards"][0]) == [0.0] + assert list(tensordict_output["tool_rewards"][1]) == [] + + # Verify round-trip conversion works perfectly + reconstructed = DataProto.from_tensordict(tensordict_output) + assert len(reconstructed) == 2 + assert reconstructed.meta_info["global_steps"] == 42 + assert torch.all(torch.eq(reconstructed.batch["prompts"], prompts)).item() + + +def test_serialize_deserialize_single_tensor(): + """Test serialization and deserialization of a single tensor""" + # Create test tensor + original_tensor = torch.randn(3, 4, 5) + + # Serialize + dtype, shape, data = serialize_single_tensor(original_tensor) + + # Deserialize + reconstructed_tensor = deserialize_single_tensor((dtype, shape, data)) + + # Verify results + assert torch.allclose(original_tensor, reconstructed_tensor) + assert original_tensor.shape == reconstructed_tensor.shape + assert original_tensor.dtype == reconstructed_tensor.dtype + + +def test_serialize_deserialize_tensordict_regular_tensors(): + """Test serialization and deserialization of TensorDict with regular tensors""" + # Create test data + batch_size = (5, 3) + tensor1 = torch.randn(*batch_size, 4) + tensor2 = torch.randint(0, 10, (*batch_size, 2)) + + # Create TensorDict + original_tensordict = TensorDict({"tensor1": tensor1, "tensor2": tensor2}, batch_size=batch_size) + + # Serialize + batch_size_serialized, device, encoded_items = serialize_tensordict(original_tensordict) + + # Deserialize + reconstructed_tensordict = deserialize_tensordict((batch_size_serialized, device, encoded_items)) + + # Verify results + assert original_tensordict.batch_size == reconstructed_tensordict.batch_size + assert set(original_tensordict.keys()) == set(reconstructed_tensordict.keys()) + + for key in original_tensordict.keys(): + original_tensor = original_tensordict[key] + reconstructed_tensor = reconstructed_tensordict[key] + + assert torch.allclose(original_tensor, reconstructed_tensor) + assert original_tensor.shape == reconstructed_tensor.shape + assert original_tensor.dtype == reconstructed_tensor.dtype + + +def test_serialize_deserialize_tensordict_nested_tensors(): + """Test serialization and deserialization of TensorDict with nested tensors""" + # Create nested tensor + tensor_list = [torch.randn(2, 3), torch.randn(3, 4), torch.randn(1, 5)] + nested_tensor = torch.nested.as_nested_tensor(tensor_list) + + # Create regular tensor for comparison + regular_tensor = torch.randn(3, 4, 5) + + # Create TensorDict + original_tensordict = TensorDict({"nested": nested_tensor, "regular": regular_tensor}, batch_size=(3,)) + + # Serialize + batch_size_serialized, device, encoded_items = serialize_tensordict(original_tensordict) + + # Deserialize + reconstructed_tensordict = deserialize_tensordict((batch_size_serialized, device, encoded_items)) + + # Verify results + assert original_tensordict.batch_size == reconstructed_tensordict.batch_size + assert set(original_tensordict.keys()) == set(reconstructed_tensordict.keys()) + + # Verify regular tensor + original_regular = original_tensordict["regular"] + reconstructed_regular = reconstructed_tensordict["regular"] + + assert torch.allclose(original_regular, reconstructed_regular) + assert original_regular.shape == reconstructed_regular.shape + assert original_regular.dtype == reconstructed_regular.dtype + + # Verify nested tensor + original_nested = original_tensordict["nested"] + reconstructed_nested = reconstructed_tensordict["nested"] + + # Check if it's a nested tensor + assert original_nested.is_nested + assert reconstructed_nested.is_nested + + # Check layout + assert original_nested.layout == reconstructed_nested.layout + + # Check each tensor after unbinding + original_unbind = original_nested.unbind() + reconstructed_unbind = reconstructed_nested.unbind() + + assert len(original_unbind) == len(reconstructed_unbind) + + for orig, recon in zip(original_unbind, reconstructed_unbind, strict=False): + assert torch.allclose(orig, recon) + assert orig.shape == recon.shape + assert orig.dtype == recon.dtype + + +def test_serialize_deserialize_tensordict_mixed_types(): + """Test serialization and deserialization of TensorDict with mixed tensor types""" + # Create tensors with different data types + float_tensor = torch.randn(2, 3).float() + double_tensor = torch.randn(2, 3).double() + int_tensor = torch.randint(0, 10, (2, 3)).int() + long_tensor = torch.randint(0, 10, (2, 3)).long() + bool_tensor = torch.tensor([[True, False], [False, True]]) + bfloat16_tensor = torch.randn(2, 3).bfloat16() + + # Add fp8 tensor (if available) + # Note: FP8 is not natively supported in all PyTorch versions + # We'll check if it's available and conditionally include it + has_fp8 = hasattr(torch, "float8_e5m2") or hasattr(torch, "float8_e4m3fn") + if has_fp8: + try: + # Try to create an FP8 tensor (implementation may vary) + # This is a placeholder - actual FP8 support might require specific hardware + fp8_tensor = torch.randn(2, 3) + if hasattr(torch, "float8_e5m2"): + fp8_tensor = fp8_tensor.to(torch.float8_e5m2) + elif hasattr(torch, "float8_e4m3fn"): + fp8_tensor = fp8_tensor.to(torch.float8_e4m3fn) + except Exception: + has_fp8 = False + + # Create nested tensor + tensor_list = [ + torch.randn(2, 3), + torch.randn(3, 4), + ] + nested_tensor = torch.nested.as_nested_tensor(tensor_list) + + # Create TensorDict with all available types + tensordict_data = { + "float": float_tensor, + "double": double_tensor, + "int": int_tensor, + "long": long_tensor, + "bool": bool_tensor, + "bfloat16": bfloat16_tensor, + "nested": nested_tensor, + } + + # Conditionally add fp8 tensor if available + if has_fp8: + tensordict_data["fp8"] = fp8_tensor + + original_tensordict = TensorDict( + tensordict_data, + batch_size=(2,), + ) + + # Serialize + batch_size_serialized, device, encoded_items = serialize_tensordict(original_tensordict) + + # Deserialize + reconstructed_tensordict = deserialize_tensordict((batch_size_serialized, device, encoded_items)) + + # Verify results + assert original_tensordict.batch_size == reconstructed_tensordict.batch_size + assert set(original_tensordict.keys()) == set(reconstructed_tensordict.keys()) + + for key in original_tensordict.keys(): + original_tensor = original_tensordict[key] + reconstructed_tensor = reconstructed_tensordict[key] + + if original_tensor.is_nested: + # For nested tensors, check each tensor after unbinding + original_unbind = original_tensor.unbind() + reconstructed_unbind = reconstructed_tensor.unbind() + + assert len(original_unbind) == len(reconstructed_unbind) + + for orig, recon in zip(original_unbind, reconstructed_unbind, strict=False): + assert torch.allclose(orig, recon, equal_nan=True) + assert orig.shape == recon.shape + assert orig.dtype == recon.dtype + else: + # For regular tensors, compare directly + assert torch.all(original_tensor == reconstructed_tensor) + assert original_tensor.shape == reconstructed_tensor.shape + assert original_tensor.dtype == reconstructed_tensor.dtype + + +def test_serialize_deserialize_tensordict_with_device(): + """Test serialization and deserialization of TensorDict with device information""" + # Create test data + batch_size = (2, 3) + tensor1 = torch.randn(*batch_size, 4) + tensor2 = torch.randint(0, 10, (*batch_size, 2)) + + # Create TensorDict with device information + device = "cpu" + original_tensordict = TensorDict({"tensor1": tensor1, "tensor2": tensor2}, batch_size=batch_size, device=device) + + # Serialize + batch_size_serialized, device_serialized, encoded_items = serialize_tensordict(original_tensordict) + + # Deserialize + reconstructed_tensordict = deserialize_tensordict((batch_size_serialized, device_serialized, encoded_items)) + + # Verify results + assert original_tensordict.batch_size == reconstructed_tensordict.batch_size + assert str(original_tensordict.device) == str(reconstructed_tensordict.device) + assert set(original_tensordict.keys()) == set(reconstructed_tensordict.keys()) + + for key in original_tensordict.keys(): + original_tensor = original_tensordict[key] + reconstructed_tensor = reconstructed_tensordict[key] + + assert torch.allclose(original_tensor.cpu(), reconstructed_tensor.cpu()) + assert original_tensor.shape == reconstructed_tensor.shape + assert original_tensor.dtype == reconstructed_tensor.dtype + + +def test_serialize_dataproto_with_empty_tensordict(): + """Tests that serializing a DataProto with an empty TensorDict does not crash. + + This test verifies the fix for the torch.cat error that occurs when calling + consolidate() on an empty TensorDict during serialization. + """ + import pickle + + # This test requires tensordict >= 0.5.0 to trigger the code path + if parse_version(tensordict.__version__) < parse_version("0.5.0"): + pytest.skip("Test requires tensordict>=0.5.0") + + # Create a DataProto with an empty TensorDict but with a batch size + empty_td = TensorDict({}, batch_size=[10]) + data = DataProto(batch=empty_td) + + # This would crash before the fix with: + # RuntimeError: torch.cat(): expected a non-empty list of Tensors + try: + serialized_data = pickle.dumps(data) + except Exception as e: + pytest.fail(f"Serializing DataProto with empty TensorDict failed with: {e}") + + # Verify deserialization works as expected + deserialized_data = pickle.loads(serialized_data) + assert len(deserialized_data.batch.keys()) == 0 + assert deserialized_data.batch.batch_size == torch.Size([10]) diff --git a/verl/tests/test_protocol_v2_on_cpu.py b/verl/tests/test_protocol_v2_on_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..63d0ee478077c6f69d82984a3955776902036c43 --- /dev/null +++ b/verl/tests/test_protocol_v2_on_cpu.py @@ -0,0 +1,1151 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Replace DataProto with raw TensorDict +""" + +import copy +import random + +import numpy as np +import pytest +import torch +from tensordict.tensorclass import NonTensorData, NonTensorStack + +from verl.utils import tensordict_utils as tu + + +def test_union_tensor_dict(): + obs = torch.randn(100, 10) + + meta_info1 = {"top_p": 0.8} + meta_info2 = {"top_p": 0.9} + data1 = {"obs": obs, "act": torch.randn(100, 3), "data_sources": ["gsm8k"] * 100} + data2 = {"obs": obs, "next_obs": torch.randn(100, 10), "rew": torch.randn(100), "data_sources": ["gsm8k"] * 100} + + data_with_copied_obs = {"obs": obs.clone(), "next_obs": torch.randn(100, 10), "rew": torch.randn(100)} + + data1 = tu.get_tensordict(tensor_dict=data1) + data2 = tu.get_tensordict(tensor_dict=data2) + data_with_copied_obs = tu.get_tensordict(data_with_copied_obs) + + tu.union_tensor_dict(data1, data2) + with pytest.raises(AssertionError): + # conflict in tensor values + tu.union_tensor_dict(data1, data_with_copied_obs) + + data1 = tu.assign_non_tensor(data1, **meta_info1) + tu.union_tensor_dict(data1, data2) # works ok + + data2 = tu.assign_non_tensor(data2, **meta_info2) + + with pytest.raises(AssertionError): + # conflict in NonTensorData + tu.union_tensor_dict(data1, data2) + + data1.pop("top_p") + data2.pop("top_p") + + data2["data_sources"][0] = "math" + with pytest.raises(AssertionError): + # conflict in NonTensorData + tu.union_tensor_dict(data1, data2) + + +def test_tensor_dict_constructor(): + obs = torch.ones(100, 10) + act = torch.zeros(100, 10, 3) + data_source = ["gsm8k"] * 100 + non_tensor_dict = {"name": "abdce"} + + data = tu.get_tensordict( + tensor_dict={"obs": obs, "act": act, "data_source": data_source}, non_tensor_dict=non_tensor_dict + ) + + assert data.batch_size == torch.Size([100]) + + # test slicing + assert torch.all(torch.eq(data[0]["obs"], torch.ones(10))).item() + assert torch.all(torch.eq(data[0]["act"], torch.zeros(10, 3))).item() + assert data[0]["data_source"] == "gsm8k" + + assert torch.all(torch.eq(data[0:2]["obs"], torch.ones(2, 10))).item() + assert torch.all(torch.eq(data[0:2]["act"], torch.zeros(2, 10, 3))).item() + assert data[0:2]["data_source"] == ["gsm8k"] * 2 + + # test non tensor data + assert data["name"] == "abdce" + + +def test_index_select_tensor_dict(): + vocab_size = 128 + a = torch.randint(low=0, high=vocab_size, size=(11,)) + b = torch.randint(low=0, high=vocab_size, size=(13,)) + c = torch.randint(low=0, high=vocab_size, size=(12,)) + d = torch.randint(low=0, high=vocab_size, size=(15,)) + input_ids = [a, b, c, d] + input_ids = torch.nested.as_nested_tensor(input_ids, layout=torch.jagged) + + padded_tensor = torch.randn(4, 10) + non_tensor_dict = {"global_batch_size": "4"} + + data = tu.get_tensordict( + tensor_dict={ + "input_ids": input_ids, + "padded_tensor": padded_tensor, + }, + non_tensor_dict=non_tensor_dict, + ) + + assert data.batch_size == torch.Size([4]) + + # test index select + indices = torch.tensor([1, 3]) + selected_data = tu.index_select_tensor_dict(data, indices) + + assert selected_data.batch_size == torch.Size([2]) + + target_input_ids = torch.nested.as_nested_tensor([input_ids[idx] for idx in indices], layout=torch.jagged) + target_select_data = tu.get_tensordict( + tensor_dict={ + "input_ids": target_input_ids, + "padded_tensor": padded_tensor[indices], + }, + non_tensor_dict=non_tensor_dict, + ) + tu.assert_tensordict_eq(selected_data, target_select_data) + + +def test_index_select_tensor_dict_preserves_3d_nested_tensor_layout_with_equal_seq_len(): + position_ids = tu.nested_tensor_from_tensor_list( + [ + torch.arange(2).expand(4, 2), + torch.arange(5).expand(4, 5), + (torch.arange(5) + 10).expand(4, 5), + torch.arange(3).expand(4, 3), + ] + ) + data = tu.get_tensordict({"position_ids": position_ids}) + + selected = tu.index_select_tensor_dict(data, torch.tensor([1, 2])) + expected = tu.nested_tensor_from_tensor_list([position_ids[1], position_ids[2]], ragged_idx=2) + + assert selected["position_ids"]._ragged_idx == 2 + assert selected["position_ids"].values().shape == torch.Size([4, 10]) + assert torch.equal(selected["position_ids"].values(), expected.values()) + assert torch.equal(selected["position_ids"].offsets(), expected.offsets()) + tu.assert_tensordict_eq(selected, tu.get_tensordict({"position_ids": expected})) + + +def test_tensordict_with_images(): + # each sample contains a sequence with multiple images of different sizes + vocab_size = 128 + a = torch.randint(low=0, high=vocab_size, size=(11,)) + b = torch.randint(low=0, high=vocab_size, size=(13,)) + input_ids = [a, b] + input_ids = torch.nested.as_nested_tensor(input_ids, layout=torch.jagged) + + # must be numpy + # TODO(vermouth1992). We may use nested tensor too. But this requires nested over nested + a_images = [ + torch.randint(low=0, high=255, size=(3, 256, 256), dtype=torch.uint8).numpy(), + torch.randint(low=0, high=255, size=(3, 128, 128), dtype=torch.uint8).numpy(), + ] + b_images = [ + torch.randint(low=0, high=255, size=(3, 256, 256), dtype=torch.uint8).numpy(), + torch.randint(low=0, high=255, size=(3, 128, 128), dtype=torch.uint8).numpy(), + torch.randint(low=0, high=255, size=(3, 64, 64), dtype=torch.uint8).numpy(), + ] + + images = [a_images, b_images] + + data = tu.get_tensordict({"input_ids": input_ids, "images": images}) + + assert np.all(np.equal(data[0]["images"][0], a_images[0])) + assert torch.all(torch.eq(data[0]["input_ids"], a)) + + +def test_tensordict_with_packing(): + vocab_size = 128 + a = torch.randint(low=0, high=vocab_size, size=(11,)) + b = torch.randint(low=0, high=vocab_size, size=(13,)) + input_ids = [a, b] + input_ids = torch.nested.as_nested_tensor(input_ids, layout=torch.jagged) + + data = tu.get_tensordict({"input_ids": input_ids}) + + # test cu_seqlens + cu_seqlens = torch.tensor([0, 11, 24]) + assert torch.all(torch.eq(cu_seqlens, data["input_ids"].offsets())) + + # test index + assert torch.all(torch.eq(data["input_ids"][0], a)) + assert torch.all(torch.eq(data["input_ids"][1], b)) + + assert torch.all(torch.eq(data[0]["input_ids"], a)) + assert torch.all(torch.eq(data[1]["input_ids"], b)) + + data_lst = data.chunk(2) + + assert torch.all(torch.eq(data_lst[0]["input_ids"][0], a)) + assert torch.all(torch.eq(data_lst[1]["input_ids"][0], b)) + + +def test_tensordict_eq(): + obs = torch.tensor([1, 2, 3, 4, 5, 6]) + data_sources = ["abc", "def", "abc", "def", "pol", "klj"] + non_tensor_dict = {"train_sample_kwargs": {"top_p": 1.0}, "val_sample_kwargs": {"top_p": 0.7}} + data = tu.get_tensordict({"obs": obs, "data_sources": data_sources}, non_tensor_dict=non_tensor_dict) + + obs = torch.tensor([1, 2, 3, 4, 5, 6]) + data_sources = ["abc", "def", "abc", "def", "pol", "klj"] + non_tensor_dict = {"train_sample_kwargs": {"top_p": 1.0}, "val_sample_kwargs": {"top_p": 0.7}} + data1 = tu.get_tensordict({"obs": obs, "data_sources": data_sources}, non_tensor_dict=non_tensor_dict) + + tu.assert_tensordict_eq(data, data1) + + data2 = copy.deepcopy(data1) + data2["obs"][0] += 1 + + with pytest.raises(AssertionError): + tu.assert_tensordict_eq(data, data2) + + data2 = copy.deepcopy(data1) + data2["data_sources"][0] = "math" + + with pytest.raises(AssertionError): + tu.assert_tensordict_eq(data, data2) + + data2 = copy.deepcopy(data1) + data2["train_sample_kwargs"]["top_p"] = 0.9 + + with pytest.raises(AssertionError): + tu.assert_tensordict_eq(data, data2) + + tensor_list = [ + torch.tensor([1, 2, 3, 3, 2]), + torch.tensor([4, 5]), + torch.tensor([7, 8, 10, 14]), + torch.tensor([10, 11, 12]), + torch.tensor([13, 14, 15, 18]), + torch.tensor([16, 17]), + ] + obs = torch.nested.as_nested_tensor(tensor_list, layout=torch.jagged) + data_sources = ["abc", "def", "abc", "def", "pol", "klj"] + non_tensor_dict = {"train_sample_kwargs": {"top_p": 1.0}, "val_sample_kwargs": {"top_p": 0.7}} + data3 = tu.get_tensordict({"obs": obs, "data_sources": data_sources}, non_tensor_dict=non_tensor_dict) + + tensor_list[0] = torch.tensor([1, 2, 3, 3, 2]) + obs = torch.nested.as_nested_tensor(tensor_list, layout=torch.jagged) + data4 = tu.get_tensordict({"obs": obs, "data_sources": data_sources}, non_tensor_dict=non_tensor_dict) + tu.assert_tensordict_eq(data3, data4) + + tensor_list[0] = torch.tensor([1, 2, 4]) + obs = torch.nested.as_nested_tensor(tensor_list, layout=torch.jagged) + data5 = tu.get_tensordict({"obs": obs, "data_sources": data_sources}, non_tensor_dict=non_tensor_dict) + with pytest.raises(AssertionError): + tu.assert_tensordict_eq(data3, data5) + + tensor_list[0] = torch.tensor([4, 5]) + tensor_list[1] = torch.tensor([1, 2, 3, 3, 2]) + obs = torch.nested.as_nested_tensor(tensor_list, layout=torch.jagged) + data6 = tu.get_tensordict({"obs": obs, "data_sources": data_sources}, non_tensor_dict=non_tensor_dict) + with pytest.raises(AssertionError): + tu.assert_tensordict_eq(data3, data6) + + +def test_tensor_dict_make_iterator(): + obs = torch.tensor([1, 2, 3, 4, 5, 6]) + input_ids = torch.nested.as_nested_tensor( + [ + torch.tensor([0, 1]), + torch.tensor([2]), + torch.tensor([3, 4]), + torch.tensor([5]), + torch.tensor([6, 7, 8]), + torch.tensor([9]), + ], + layout=torch.jagged, + ) + data_sources = ["abc", "def", "abc", "def", "pol", "klj"] + non_tensor_dict = {"train_sample_kwargs": {"top_p": 1.0}, "val_sample_kwargs": {"top_p": 0.7}} + dataset = tu.get_tensordict( + {"obs": obs, "data_sources": data_sources, "input_ids": input_ids}, non_tensor_dict=non_tensor_dict + ) + + dataloader = tu.make_iterator( + dataset, mini_batch_size=2, epochs=2, seed=0, dataloader_kwargs={"shuffle": False, "drop_last": False} + ) + + expected_tensor_dict = [ + tu.index_select_tensor_dict(dataset, indices=list(range(0, 2))), + tu.index_select_tensor_dict(dataset, indices=list(range(2, 4))), + tu.index_select_tensor_dict(dataset, indices=list(range(4, 6))), + tu.index_select_tensor_dict(dataset, indices=list(range(0, 2))), + tu.index_select_tensor_dict(dataset, indices=list(range(2, 4))), + tu.index_select_tensor_dict(dataset, indices=list(range(4, 6))), + ] + + i = 0 + + for d in dataloader: + tu.assert_tensordict_eq(d, expected_tensor_dict[i]) + i += 1 + + data_iter_1 = tu.make_iterator(dataset, mini_batch_size=3, epochs=1, seed=1, dataloader_kwargs={"shuffle": True}) + data_list_1 = [] + for data in data_iter_1: + data_list_1.append(data) + + data_iter_2 = tu.make_iterator(dataset, mini_batch_size=3, epochs=1, seed=1, dataloader_kwargs={"shuffle": True}) + data_list_2 = [] + for data in data_iter_2: + data_list_2.append(data) + + for data1, data2 in zip(data_list_1, data_list_2, strict=True): + tu.assert_tensordict_eq(data1, data2) + + +def test_reorder(): + obs = torch.tensor([1, 2, 3, 4, 5, 6]) + labels = ["a", "b", "c", "d", "e", "f"] + non_tensor_dict = {"name": "abdce"} + + data = tu.get_tensordict(tensor_dict={"obs": obs, "labels": labels}, non_tensor_dict=non_tensor_dict) + data = data[torch.tensor([3, 4, 2, 0, 1, 5])] + + assert torch.all(torch.eq(data["obs"], torch.tensor([4, 5, 3, 1, 2, 6]))) + assert np.all(data["labels"] == np.array(["d", "e", "c", "a", "b", "f"])) + assert data["name"] == "abdce" + + +def test_chunk_concat(): + obs = torch.tensor([1, 2, 3, 4, 5, 6]) + labels = ["a", "b", "c", "d", "e", "f"] + data = tu.get_tensordict({"obs": obs, "labels": labels}, non_tensor_dict={"name": "abcde"}) + + data_split = data.tensor_split(indices_or_sections=5, dim=0) + + expected_idx_lst = [[0, 1], [2], [3], [4], [5]] + + for d, expected_idx in zip(data_split, expected_idx_lst, strict=False): + tu.assert_tensordict_eq(d, data[expected_idx]) + + data_split = data.chunk(2) + assert len(data_split) == 2 + assert torch.all(torch.eq(data_split[0]["obs"], torch.tensor([1, 2, 3]))) + assert np.all(data_split[0]["labels"] == np.array(["a", "b", "c"])) + assert data_split[0]["name"] == "abcde" + + assert torch.all(torch.eq(data_split[1]["obs"], torch.tensor([4, 5, 6]))) + assert np.all(data_split[1]["labels"] == np.array(["d", "e", "f"])) + assert data_split[1]["name"] == "abcde" + + concat_data = torch.cat(data_split, dim=0) + assert torch.all(torch.eq(concat_data["obs"], data["obs"])) + assert np.all(concat_data["labels"] == data["labels"]) + assert concat_data["name"] == data["name"] + + data1 = tu.get_tensordict(tensor_dict={"obs": obs, "labels": labels}, non_tensor_dict={"name": "abcde"}) + data2 = tu.get_tensordict(tensor_dict={"obs": obs, "labels": labels}, non_tensor_dict={"name": "def"}) + data3 = tu.get_tensordict(tensor_dict={"obs": obs, "labels": labels}, non_tensor_dict={"name": "cfg"}) + + output = torch.cat([data1, data2, data3], dim=0) + + # concat NonTensorData will keep the first one. + assert output["name"] == "abcde" + + +def test_pop(): + obs = torch.randn(3, 10) + act = torch.randn(3, 3) + labels = ["a", ["b"], []] + dataset = tu.get_tensordict({"obs": obs, "act": act, "labels": labels}, non_tensor_dict={"2": 2, "1": 1}) + + dataset1 = copy.deepcopy(dataset) + + # test pop keys + popped_dataset = tu.pop_keys(dataset, keys=["obs", "2"]) + + assert popped_dataset.batch_size[0] == 3 + + assert popped_dataset.keys() == {"obs", "2"} + assert torch.all(torch.eq(popped_dataset["obs"], obs)).item() + assert popped_dataset["2"] == 2 + + assert dataset.keys() == {"act", "1", "labels"} + + # test pop non-exist key + with pytest.raises(KeyError): + tu.pop_keys(dataset, keys=["obs", "2"]) + + # test single pop + # NonTensorData + assert tu.pop(dataset1, key="2") == 2 + # NonTensorStack + assert tu.pop(dataset1, key="labels") == ["a", ["b"], []] + # Tensor + assert torch.all(torch.eq(tu.pop(dataset1, key="obs"), obs)).item() + + +def test_get(): + obs = torch.randn(3, 10) + act = torch.randn(3, 3) + labels = ["a", ["b"], []] + dataset = tu.get_tensordict({"obs": obs, "act": act, "labels": labels}, non_tensor_dict={"2": 2, "1": 1}) + + # test pop keys + popped_dataset = tu.get_keys(dataset, keys=["obs", "2"]) + + assert popped_dataset.batch_size[0] == 3 + + assert torch.all(torch.eq(popped_dataset["obs"], dataset["obs"])).item() + + assert popped_dataset["2"] == dataset["2"] + + # test pop non-exist key + with pytest.raises(KeyError): + tu.get_keys(dataset, keys=["obs", "3"]) + + # test single pop + # NonTensorData + assert tu.get(dataset, key="2") == 2 + # NonTensorStack + assert tu.get(dataset, key="labels") == ["a", ["b"], []] + # Tensor + assert torch.all(torch.eq(tu.get(dataset, key="obs"), obs)).item() + # Non-exist key + assert tu.get(dataset, key="3", default=3) == 3 + + +def test_repeat(): + # Create a DataProto object with some batch and non-tensor data + obs = torch.tensor([[1, 2], [3, 4], [5, 6]]) + labels = ["a", "b", "c"] + data = tu.get_tensordict({"obs": obs, "labels": labels}, non_tensor_dict={"info": "test_info"}) + + # Test interleave=True + repeated_data_interleave = data.repeat_interleave(repeats=2) + expected_obs_interleave = torch.tensor([[1, 2], [1, 2], [3, 4], [3, 4], [5, 6], [5, 6]]) + expected_labels_interleave = ["a", "a", "b", "b", "c", "c"] + + assert torch.all(torch.eq(repeated_data_interleave["obs"], expected_obs_interleave)) + assert repeated_data_interleave["labels"] == expected_labels_interleave + assert repeated_data_interleave["info"] == "test_info" + + # Test interleave=False + repeated_data_no_interleave = data.repeat(2) + expected_obs_no_interleave = torch.tensor([[1, 2], [3, 4], [5, 6], [1, 2], [3, 4], [5, 6]]) + expected_labels_no_interleave = ["a", "b", "c", "a", "b", "c"] + + assert torch.all(torch.eq(repeated_data_no_interleave["obs"], expected_obs_no_interleave)) + assert repeated_data_no_interleave["labels"] == expected_labels_no_interleave + assert repeated_data_no_interleave["info"] == "test_info" + + +def test_dataproto_pad_unpad(): + obs = torch.tensor([[1, 2], [3, 4], [5, 6]]) + labels = ["a", "b", "c"] + data = tu.get_tensordict(tensor_dict={"obs": obs, "labels": labels}, non_tensor_dict={"info": "test_info"}) + + padded_data, pad_size = tu.pad_to_divisor(data, size_divisor=2) + + assert pad_size == 1 + + expected_obs = torch.tensor([[1, 2], [3, 4], [5, 6], [1, 2]]) + expected_labels = ["a", "b", "c", "a"] + + assert torch.all(torch.eq(padded_data["obs"], expected_obs)) + assert padded_data["labels"] == expected_labels + assert padded_data["info"] == "test_info" + + unpadd_data = tu.unpad(padded_data, pad_size=pad_size) + assert torch.all(torch.eq(unpadd_data["obs"], obs)) + assert unpadd_data["labels"] == labels + assert unpadd_data["info"] == "test_info" + + padded_data, pad_size = tu.pad_to_divisor(data, size_divisor=3) + assert pad_size == 0 + + expected_obs = torch.tensor([[1, 2], [3, 4], [5, 6]]) + expected_labels = ["a", "b", "c"] + + assert torch.all(torch.eq(padded_data["obs"], expected_obs)) + assert padded_data["labels"] == expected_labels + assert padded_data["info"] == "test_info" + + unpadd_data = tu.unpad(padded_data, pad_size=pad_size) + assert torch.all(torch.eq(unpadd_data["obs"], obs)) + assert unpadd_data["labels"] == labels + assert unpadd_data["info"] == "test_info" + + padded_data, pad_size = tu.pad_to_divisor(data, size_divisor=7) + assert pad_size == 4 + + expected_obs = torch.tensor([[1, 2], [3, 4], [5, 6], [1, 2], [3, 4], [5, 6], [1, 2]]) + expected_labels = ["a", "b", "c", "a", "b", "c", "a"] + assert torch.all(torch.eq(padded_data["obs"], expected_obs)) + assert padded_data["labels"] == expected_labels + assert padded_data["info"] == "test_info" + + unpadd_data = tu.unpad(padded_data, pad_size=pad_size) + assert torch.all(torch.eq(unpadd_data["obs"], obs)) + assert unpadd_data["labels"] == labels + assert unpadd_data["info"] == "test_info" + + +def test_torch_save_data_proto(): + obs = torch.tensor([[1, 2], [3, 4], [5, 6]]) + labels = ["a", "b", "c"] + data = tu.get_tensordict({"obs": obs, "labels": labels}, non_tensor_dict={"info": "test_info"}) + + filename = "test_data.pt" + torch.save(data, filename) + loaded_data = torch.load(filename, weights_only=False) + + assert torch.all(torch.eq(loaded_data["obs"], data["obs"])) + assert loaded_data["labels"] == data["labels"] + assert loaded_data["info"] == data["info"] + + import os + + os.remove(filename) + + +def test_len(): + obs = torch.tensor([[1, 2], [3, 4], [5, 6]]) + labels = np.array(["a", "b", "c"], dtype=object) + + data = tu.get_tensordict({"obs": obs, "labels": labels.tolist()}, non_tensor_dict={"info": "test_info"}) + assert len(data) == 3 + + data = tu.get_tensordict({"labels": labels.tolist()}, non_tensor_dict={"info": "test_info"}) + assert len(data) == 3 + + data_item = data[0] + assert len(data_item) == 0 + + data = tu.get_tensordict({}, non_tensor_dict={"info": "test_info"}) + assert len(data) == 0 + + +def test_dataproto_index(): + data_len = 100 + idx_num = 10 + + obs = torch.randn(data_len, 10) + labels = [random.choice(["abc", "cde"]) for _ in range(data_len)] + + data = tu.get_tensordict({"obs": obs, "labels": labels}) + + labels_np = np.array(labels) + + idx_np_int = np.random.randint(0, data_len, size=(idx_num,)) + result_np_int = data[idx_np_int] + assert result_np_int.keys() == data.keys() + assert result_np_int["obs"].shape[0] == idx_num + assert len(result_np_int["labels"]) == idx_num + assert np.array_equal(result_np_int["obs"].cpu().numpy(), obs[idx_np_int].numpy()) + assert np.array_equal(result_np_int["labels"], labels_np[idx_np_int]) + + idx_torch_int = torch.randint(0, data_len, size=(idx_num,)) + result_torch_int = data[idx_torch_int] + assert result_torch_int.keys() == data.keys() + assert result_torch_int["obs"].shape[0] == idx_num + assert len(result_torch_int["labels"]) == idx_num + assert np.array_equal(result_torch_int["obs"].cpu().numpy(), obs[idx_torch_int].cpu().numpy()) + assert np.array_equal(result_torch_int["labels"], labels_np[idx_torch_int.cpu().numpy()]) + + idx_list_int = [np.random.randint(0, data_len) for _ in range(idx_num)] + result_list_int = data[idx_list_int] + assert result_list_int.keys() == data.keys() + assert result_list_int["obs"].shape[0] == idx_num + assert len(result_list_int["labels"]) == idx_num + assert np.array_equal(result_list_int["obs"].cpu().numpy(), obs[idx_list_int].cpu().numpy()) + assert np.array_equal(result_list_int["labels"], labels_np[idx_list_int]) + + # idx_np_bool = np.random.randint(0, 2, size=(data_len,), dtype=bool) + # result_np_bool = data[idx_np_bool] + # assert result_np_bool.keys() == data.keys() + # assert result_np_bool["obs"].shape[0] == idx_np_bool.sum() + # assert len(result_np_bool["labels"]) == idx_np_bool.sum() + # assert np.array_equal(result_np_bool["obs"].cpu().numpy(), obs[idx_np_bool].cpu().numpy()) + # assert np.array_equal(result_np_bool["labels"], labels_np[idx_np_bool]) + + idx_torch_bool = torch.randint(0, 2, size=(data_len,), dtype=torch.bool) + result_torch_bool = data[idx_torch_bool] + assert result_torch_bool.keys() == data.keys() + assert result_torch_bool["obs"].shape[0] == idx_torch_bool.sum().item() + assert len(result_torch_bool["labels"]) == idx_torch_bool.sum().item() + assert np.array_equal(result_torch_bool["obs"].cpu().numpy(), obs[idx_torch_bool].cpu().numpy()) + assert np.array_equal(result_torch_bool["labels"], labels_np[idx_torch_bool]) + + # idx_list_bool = [np.random.randint(0, 2, dtype=bool) for _ in range(data_len)] + # result_list_bool = data[idx_list_bool] + # assert result_list_bool.keys() == data.keys() + # assert result_list_bool["obs"].shape[0] == sum(idx_list_bool) + # assert len(result_list_bool["labels"]) == sum(idx_list_bool) + # assert np.array_equal(result_list_bool["obs"].cpu().numpy(), obs[idx_list_bool].cpu().numpy()) + # assert np.array_equal(result_list_bool["labels"], labels_np[idx_list_bool]) + + +def test_select(): + obs = torch.randn(100, 10) + act = torch.randn(100, 3) + dataset = tu.get_tensordict({"obs": obs, "act": act}, non_tensor_dict={"2": 2, "1": 1}) + + subset = dataset.select("obs", "2") + + assert torch.all(torch.eq(subset["obs"], dataset["obs"])) + assert subset["2"] == dataset["2"] + assert "act" not in subset.keys() + assert "1" not in subset.keys() + + +def test_dataproto_no_batch(): + labels = ["a", "b", "c"] + data = tu.get_tensordict(tensor_dict={"labels": labels}, non_tensor_dict={"info": "test_info"}) + selected = data.select("labels") + + assert selected["labels"] == labels + pop_data = tu.pop_keys(data, keys=["labels"]) + assert pop_data["labels"] == labels + assert "labels" not in data + + +def test_sample_level_repeat(): + # Create a DataProto object with some batch and non-tensor data + obs = torch.tensor([[1, 2], [3, 4], [5, 6]]) + labels = ["a", "b", "c"] + + data = tu.get_tensordict({"obs": obs, "labels": labels}, non_tensor_dict={"info": "test_info"}) + + # list + repeated_data_interleave = data.repeat_interleave(repeats=torch.tensor([3, 1, 2])) + expected_obs_interleave = torch.tensor([[1, 2], [1, 2], [1, 2], [3, 4], [5, 6], [5, 6]]) + expected_labels_interleave = ["a", "a", "a", "b", "c", "c"] + + assert torch.all(torch.eq(repeated_data_interleave["obs"], expected_obs_interleave)) + assert repeated_data_interleave["labels"] == expected_labels_interleave + assert repeated_data_interleave["info"] == "test_info" + + # torch.tensor + repeated_data_no_interleave = data.repeat_interleave(repeats=torch.tensor([1, 2, 3])) + expected_obs_no_interleave = torch.tensor([[1, 2], [3, 4], [3, 4], [5, 6], [5, 6], [5, 6]]) + expected_labels_no_interleave = ["a", "b", "b", "c", "c", "c"] + + assert torch.all(torch.eq(repeated_data_no_interleave["obs"], expected_obs_no_interleave)) + assert repeated_data_no_interleave["labels"] == expected_labels_no_interleave + assert repeated_data_no_interleave["info"] == "test_info" + + +def test_dataproto_chunk_after_index(): + data_len = 4 + obs = torch.randn(data_len, 4) + labels = [f"label_{i}" for i in range(data_len)] + + data = tu.get_tensordict(tensor_dict={"obs": obs, "labels": labels}, non_tensor_dict={"name": "abc"}) + # Test with boolean numpy array + bool_mask = torch.tensor([True, False, True, False]) + selected = data[bool_mask] + assert isinstance(selected.batch_size, torch.Size) + assert all(isinstance(d, int) for d in selected.batch_size) # int or List[int] + + # Test with integer numpy array + int_mask = torch.tensor([0, 2]) + selected = data[int_mask] + assert isinstance(selected.batch_size, torch.Size) + assert all(isinstance(d, int) for d in selected.batch_size) + + # Test with boolean list + list_mask = [True, False, True, False] + selected = data[list_mask] + assert isinstance(selected.batch_size, torch.Size) + assert all(isinstance(d, int) for d in selected.batch_size) + + # Test with list + list_mask = [0, 2] + selected = data[list_mask] + assert isinstance(selected.batch_size, torch.Size) + assert all(isinstance(d, int) for d in selected.batch_size) + + # Test with torch tensor (bool) + torch_bool_mask = torch.tensor([True, False, True, False]) + selected = data[torch_bool_mask] + assert isinstance(selected.batch_size, torch.Size) + assert all(isinstance(d, int) for d in selected.batch_size) + + # Test with torch tensor (int) + torch_int_mask = torch.tensor([0, 2]) + selected = data[torch_int_mask] + assert isinstance(selected.batch_size, torch.Size) + assert all(isinstance(d, int) for d in selected.batch_size) + + +def test_concat_nested_tensor(): + # Test 2D nested tensors + vocab_size = 128 + a = torch.randint(low=0, high=vocab_size, size=(11,)) + b = torch.randint(low=0, high=vocab_size, size=(13,)) + c = torch.randint(low=0, high=vocab_size, size=(12,)) + d = torch.randint(low=0, high=vocab_size, size=(15,)) + + nested_a_b = torch.nested.as_nested_tensor([a, b], layout=torch.jagged) + nested_c_d = torch.nested.as_nested_tensor([c, d], layout=torch.jagged) + + output = tu.concat_nested_tensors([nested_a_b, nested_c_d]) + + output_values = output.values() + expected = torch.cat([a, b, c, d], dim=0) + + assert torch.all(torch.eq(output_values, expected)).item() + + # Test 3D nested tensors + a_3d = torch.randint(low=0, high=vocab_size, size=(4, 4)) + b_3d = torch.randint(low=0, high=vocab_size, size=(4, 5)) + c_3d = torch.randint(low=0, high=vocab_size, size=(4, 6)) + d_3d = torch.randint(low=0, high=vocab_size, size=(4, 7)) + + nested_a_b_3d = torch.nested.as_nested_tensor([a_3d, b_3d], layout=torch.jagged) + nested_c_d_3d = torch.nested.as_nested_tensor([c_3d, d_3d], layout=torch.jagged) + + output_3d = tu.concat_nested_tensors([nested_a_b_3d, nested_c_d_3d]) + + assert output_3d.shape[0] == 4 + output_3d_unbind = output_3d.unbind(0) + assert torch.all(torch.eq(output_3d_unbind[0], a_3d)).item() + assert torch.all(torch.eq(output_3d_unbind[1], b_3d)).item() + assert torch.all(torch.eq(output_3d_unbind[2], c_3d)).item() + assert torch.all(torch.eq(output_3d_unbind[3], d_3d)).item() + + # Test 4D nested tensors + a_4d = torch.randint(low=0, high=vocab_size, size=(2, 3, 4)) + b_4d = torch.randint(low=0, high=vocab_size, size=(2, 3, 5)) + c_4d = torch.randint(low=0, high=vocab_size, size=(2, 3, 3)) + d_4d = torch.randint(low=0, high=vocab_size, size=(2, 3, 6)) + + nested_a_b_4d = torch.nested.as_nested_tensor([a_4d, b_4d], layout=torch.jagged) + nested_c_d_4d = torch.nested.as_nested_tensor([c_4d, d_4d], layout=torch.jagged) + + output_4d = tu.concat_nested_tensors([nested_a_b_4d, nested_c_d_4d]) + + assert output_4d.shape[0] == 4 + output_4d_unbind = output_4d.unbind(0) + assert torch.all(torch.eq(output_4d_unbind[0], a_4d)).item() + assert torch.all(torch.eq(output_4d_unbind[1], b_4d)).item() + assert torch.all(torch.eq(output_4d_unbind[2], c_4d)).item() + assert torch.all(torch.eq(output_4d_unbind[3], d_4d)).item() + + +def test_concat_tensordict(): + vocab_size = 128 + a = torch.randint(low=0, high=vocab_size, size=(11,)) + b = torch.randint(low=0, high=vocab_size, size=(13,)) + c = torch.randint(low=0, high=vocab_size, size=(12,)) + d = torch.randint(low=0, high=vocab_size, size=(15,)) + + nested_a_b = torch.nested.as_nested_tensor([a, b], layout=torch.jagged) + nested_c_d = torch.nested.as_nested_tensor([c, d], layout=torch.jagged) + + tensordict1 = tu.get_tensordict( + tensor_dict={"input_ids": nested_a_b, "labels": ["a", "b"]}, non_tensor_dict={"temp": 1.0} + ) + tensordict2 = tu.get_tensordict( + tensor_dict={"input_ids": nested_c_d, "labels": ["c", "d"]}, non_tensor_dict={"temp": 2.0} + ) + + tensordict1_copy = copy.deepcopy(tensordict1) + tensordict2_copy = copy.deepcopy(tensordict2) + + output = tu.concat_tensordict([tensordict1, tensordict2]) + + assert torch.all(torch.eq(output["input_ids"].values(), torch.cat([a, b, c, d]))).item() + assert output["labels"] == ["a", "b", "c", "d"] + assert output["temp"] == 1.0 + + # make sure tensordict1 and tensordict2 is untouched + tu.assert_tensordict_eq(tensordict1, tensordict1_copy) + tu.assert_tensordict_eq(tensordict2, tensordict2_copy) + + # test concat tensordict with only NonTensorStack and NonTensorData + tensordict1 = tu.get_tensordict(tensor_dict={"labels": ["a", "b"]}, non_tensor_dict={"temp": 1.0}) + tensordict2 = tu.get_tensordict(tensor_dict={"labels": ["c", "d"]}, non_tensor_dict={"temp": 2.0}) + + output = tu.concat_tensordict([tensordict1, tensordict2]) + + assert output["labels"] == ["a", "b", "c", "d"] + assert output["temp"] == 1.0 + + assert output.batch_size[0] == 4 + + # test concat tensordict with only NonTensorData + tensordict1 = tu.get_tensordict(tensor_dict={}, non_tensor_dict={"temp": 1.0}) + tensordict2 = tu.get_tensordict(tensor_dict={}, non_tensor_dict={"temp": 2.0}) + + output = tu.concat_tensordict([tensordict1, tensordict2]) + assert len(output.batch_size) == 0 + assert output["temp"] == 1.0 + + +def test_chunk_tensordict(): + # Qwen-VL 3d position_ids + position_ids = torch.nested.as_nested_tensor( + [ + torch.arange(4).expand(4, 4), + torch.arange(5).expand(4, 5), + torch.arange(6).expand(4, 6), + torch.arange(7).expand(4, 7), + ], + layout=torch.jagged, + ) + input_ids = torch.nested.as_nested_tensor( + [torch.arange(4), torch.arange(5), torch.arange(6), torch.arange(7)], layout=torch.jagged + ) + attention_mask = torch.nested.as_nested_tensor( + [ + torch.randint(low=0, high=2, size=[3, 4]), + torch.randint(low=0, high=2, size=[3, 5]), + torch.randint(low=0, high=2, size=[3, 6]), + torch.randint(low=0, high=2, size=[3, 7]), + ], + layout=torch.jagged, + ) + + multi_modal_inputs = torch.stack( + [ + NonTensorData({"pixel_values": torch.randn(3, 224, 224)}), + NonTensorData(None), + NonTensorData({"pixel_values": torch.randn(3, 128, 128)}), + NonTensorData({"pixel_values": torch.randn(3, 128, 128)}), + ] + ) + td = tu.get_tensordict( + { + "input_ids": input_ids, + "position_ids": position_ids, + "attention_mask": attention_mask, + "multi_modal_inputs": multi_modal_inputs, + }, + ) + assert len(td) == 4 + chunks = tu.chunk_tensordict(td, chunks=2) + + for i, chunk in enumerate(chunks): + assert len(chunk) == 2 + for key, val in chunk.items(): + if isinstance(val, torch.Tensor) and val.is_nested: + tensors = td[key].unbind(dim=0) + expected = torch.nested.as_nested_tensor(tensors[i * 2 : (i + 1) * 2], layout=torch.jagged) + assert torch.all(torch.eq(val.values(), expected.values())).item() + else: + expected = td[key][i * 2 : (i + 1) * 2] + for tensor, expect in zip(val, expected, strict=False): + if tensor.data is None: + assert expect is None + else: + assert torch.all(torch.eq(tensor.data["pixel_values"], expect["pixel_values"])).item() + + +def test_chunk_tensordict_preserves_3d_nested_tensor_layout_with_equal_seq_len_per_chunk(): + position_ids = tu.nested_tensor_from_tensor_list( + [ + torch.arange(2).expand(4, 2), + (torch.arange(2) + 10).expand(4, 2), + torch.arange(5).expand(4, 5), + (torch.arange(5) + 20).expand(4, 5), + ] + ) + input_ids = torch.nested.as_nested_tensor( + [torch.arange(2), torch.arange(2) + 10, torch.arange(5), torch.arange(5) + 20], layout=torch.jagged + ) + td = tu.get_tensordict({"input_ids": input_ids, "position_ids": position_ids}) + + chunks = tu.chunk_tensordict(td, chunks=2) + + expected_chunk_0 = tu.nested_tensor_from_tensor_list([position_ids[0], position_ids[1]], ragged_idx=2) + expected_chunk_1 = tu.nested_tensor_from_tensor_list([position_ids[2], position_ids[3]], ragged_idx=2) + + assert chunks[0]["position_ids"]._ragged_idx == 2 + assert chunks[1]["position_ids"]._ragged_idx == 2 + assert torch.equal(chunks[0]["position_ids"].values(), expected_chunk_0.values()) + assert torch.equal(chunks[1]["position_ids"].values(), expected_chunk_1.values()) + assert torch.equal(chunks[0]["position_ids"].offsets(), expected_chunk_0.offsets()) + assert torch.equal(chunks[1]["position_ids"].offsets(), expected_chunk_1.offsets()) + + +def test_chunk_tensordict_preserves_3d_nested_tensor_layout_with_non_last_ragged_idx(): + """Regression test: chunk_tensordict must handle nested tensors where the ragged dimension is not the last one.""" + topk = 64 + elements = [torch.randn(5, topk), torch.randn(8, topk), torch.randn(3, topk), torch.randn(7, topk)] + teacher_logprobs = tu.nested_tensor_from_tensor_list(elements, ragged_idx=1) + + input_ids = torch.nested.as_nested_tensor( + [torch.arange(5), torch.arange(8), torch.arange(3), torch.arange(7)], layout=torch.jagged + ) + td = tu.get_tensordict({"input_ids": input_ids, "teacher_logprobs": teacher_logprobs}) + + chunks = tu.chunk_tensordict(td, chunks=2) + + assert chunks[0]["teacher_logprobs"]._ragged_idx == 1 + assert torch.equal(chunks[0]["teacher_logprobs"].unbind(0)[0], elements[0]) + assert torch.equal(chunks[0]["teacher_logprobs"].unbind(0)[1], elements[1]) + assert chunks[1]["teacher_logprobs"]._ragged_idx == 1 + assert torch.equal(chunks[1]["teacher_logprobs"].unbind(0)[0], elements[2]) + assert torch.equal(chunks[1]["teacher_logprobs"].unbind(0)[1], elements[3]) + + +def test_index_select_tensor_dict_preserves_3d_nested_tensor_layout_with_non_last_ragged_idx(): + """Regression test: index_select_tensor_dict must handle nested tensors where the ragged dim is not the last.""" + topk = 64 + elements = [torch.randn(5, topk), torch.randn(8, topk), torch.randn(3, topk), torch.randn(7, topk)] + teacher_logprobs = tu.nested_tensor_from_tensor_list(elements, ragged_idx=1) + td = tu.get_tensordict({"teacher_logprobs": teacher_logprobs}) + + selected = tu.index_select_tensor_dict(td, torch.tensor([1, 3])) + + assert selected["teacher_logprobs"]._ragged_idx == 1 + assert torch.equal(selected["teacher_logprobs"].unbind(0)[0], elements[1]) + assert torch.equal(selected["teacher_logprobs"].unbind(0)[1], elements[3]) + + +def test_assign_non_tensor_stack_with_nested_lists(): + """Test assign_non_tensor_stack with lists of lists.""" + td = tu.get_tensordict({"obs": torch.randn(3, 4)}, non_tensor_dict={}) + + # Lists of varying lengths (like turn_scores or tool_rewards) + turn_scores = [[], [0.5, 0.8], [0.9]] + tu.assign_non_tensor_stack(td, "turn_scores", turn_scores) + + # Verify data is accessible + assert len(td["turn_scores"]) == 3 + assert list(td["turn_scores"][0]) == [] + assert list(td["turn_scores"][1]) == [0.5, 0.8] + assert list(td["turn_scores"][2]) == [0.9] + + +def test_assign_non_tensor_stack_with_nested_dicts(): + """Test assign_non_tensor_stack with lists of dicts.""" + td = tu.get_tensordict({"obs": torch.randn(3, 4)}, non_tensor_dict={}) + + # Lists of dicts (like reward_extra_info) + reward_extra_info = [{"acc": 1.0, "loss": 0.1}, {"acc": 0.0, "loss": 0.9}, {"acc": 1.0, "loss": 0.05}] + tu.assign_non_tensor_stack(td, "reward_extra_info", reward_extra_info) + + # Verify data is accessible + assert len(td["reward_extra_info"]) == 3 + assert dict(td["reward_extra_info"][0]) == {"acc": 1.0, "loss": 0.1} + assert dict(td["reward_extra_info"][1]) == {"acc": 0.0, "loss": 0.9} + assert dict(td["reward_extra_info"][2]) == {"acc": 1.0, "loss": 0.05} + + +def test_assign_non_tensor_stack_with_complex_nested(): + """Test assign_non_tensor_stack with lists of lists of dicts.""" + td = tu.get_tensordict({"obs": torch.randn(2, 4)}, non_tensor_dict={}) + + # Lists of lists of dicts (like raw_prompt) + raw_prompt = [ + [{"content": "Question 1", "role": "user"}], + [{"content": "Question 2", "role": "user"}, {"content": "Answer 2", "role": "assistant"}], + ] + tu.assign_non_tensor_stack(td, "raw_prompt", raw_prompt) + + # Verify data is accessible + assert len(td["raw_prompt"]) == 2 + assert len(td["raw_prompt"][0]) == 1 + assert dict(td["raw_prompt"][0][0]) == {"content": "Question 1", "role": "user"} + assert len(td["raw_prompt"][1]) == 2 + assert dict(td["raw_prompt"][1][0]) == {"content": "Question 2", "role": "user"} + + +def test_assign_non_tensor_handles_wrappers(): + td = tu.get_tensordict({"obs": torch.randn(3, 4)}, non_tensor_dict={}) + + meta = {"top_p": 0.8} + tu.assign_non_tensor(td, **meta) + assert td["top_p"] == 0.8 + + wrapped = NonTensorData(0.3) + stack = NonTensorStack.from_list([NonTensorData(1.0), NonTensorData(2.0), NonTensorData(3.0)]) + tu.assign_non_tensor(td, wrapped=wrapped, stack=stack) + + assert td["wrapped"] == 0.3 + assert td["stack"] == [1.0, 2.0, 3.0] + + +def test_assign_non_tensor_stack_batch_size_check(): + td = tu.get_tensordict({"obs": torch.randn(3, 4)}, non_tensor_dict={}) + stack = NonTensorStack.from_list([NonTensorData(1.0), NonTensorData(2.0)]) + + with pytest.raises(RuntimeError): + tu.assign_non_tensor(td, stack=stack) + + +def test_assign_non_tensor_with_auto_detection(): + """Test assign_non_tensor automatically detects and handles nested structures.""" + td = tu.get_tensordict({"obs": torch.randn(3, 4)}, non_tensor_dict={}) + + # Mix of simple and nested data + tu.assign_non_tensor( + td, + metadata="experiment_1", # Simple value + turn_scores=[[], [0.5, 0.8], [0.9]], # Nested list + reward_extra_info=[{"acc": 1.0}, {"acc": 0.0}, {"acc": 1.0}], # List of dicts + simple_list=["a", "b", "c"], # Simple list (also uses NonTensorStack for consistency) + ) + + # Verify all data is accessible + assert td["metadata"] == "experiment_1" + assert len(td["turn_scores"]) == 3 + assert list(td["turn_scores"][1]) == [0.5, 0.8] + assert len(td["reward_extra_info"]) == 3 + assert dict(td["reward_extra_info"][0]) == {"acc": 1.0} + assert len(td["simple_list"]) == 3 + assert td["simple_list"][0] == "a" + + +def test_get_tensordict_with_nested_lists(): + """Test get_tensordict automatically handles nested lists.""" + obs = torch.randn(3, 4) + turn_scores = [[], [0.5, 0.8], [0.9]] + + # This should automatically convert turn_scores to NonTensorStack + td = tu.get_tensordict({"obs": obs, "turn_scores": turn_scores}) + + # Verify tensors and nested data are both accessible + assert torch.all(torch.eq(td["obs"], obs)) + assert len(td["turn_scores"]) == 3 + assert list(td["turn_scores"][0]) == [] + assert list(td["turn_scores"][1]) == [0.5, 0.8] + + +def test_get_tensordict_with_nested_dicts(): + """Test get_tensordict automatically handles lists of dicts.""" + obs = torch.randn(3, 4) + reward_extra_info = [{"acc": 1.0}, {"acc": 0.0}, {"acc": 1.0}] + + td = tu.get_tensordict({"obs": obs, "reward_extra_info": reward_extra_info}) + + assert torch.all(torch.eq(td["obs"], obs)) + assert len(td["reward_extra_info"]) == 3 + assert dict(td["reward_extra_info"][0]) == {"acc": 1.0} + + +def test_get_tensordict_with_complex_nested_structures(): + """Test get_tensordict with lists of lists of dicts.""" + obs = torch.randn(2, 4) + raw_prompt = [ + [{"content": "Q1", "role": "user"}], + [{"content": "Q2", "role": "user"}, {"content": "A2", "role": "assistant"}], + ] + + td = tu.get_tensordict({"obs": obs, "raw_prompt": raw_prompt}) + + assert torch.all(torch.eq(td["obs"], obs)) + assert len(td["raw_prompt"]) == 2 + assert dict(td["raw_prompt"][0][0]) == {"content": "Q1", "role": "user"} + + +def test_get_tensordict_agent_loop_scenario(): + """Test the complete agent loop scenario with all nested types. + + This simulates the exact use case from agent loops with: + - turn_scores: lists of lists + - reward_extra_info: lists of dicts + - raw_prompt: lists of lists of dicts + - tool_rewards: lists of lists + """ + prompts = torch.randn(2, 10) + responses = torch.randn(2, 5) + + # Nested structures from agent loop + data_source = ["lighteval/MATH", "lighteval/MATH"] + uid = ["uuid-1", "uuid-2"] + turn_scores = [[], [0.5, 0.8]] # Lists of varying lengths + reward_extra_info = [{"acc": 1.0, "loss": 0.1}, {"acc": 0.0, "loss": 0.9}] + raw_prompt = [ + [{"content": "Compute 4 @ 2", "role": "user"}], + [{"content": "Compute 8 @ 7", "role": "user"}], + ] + tool_rewards = [[0.0], []] # List of lists + + # This should handle all nested structures automatically + td = tu.get_tensordict( + tensor_dict={ + "prompts": prompts, + "responses": responses, + "data_source": data_source, + "uid": uid, + "turn_scores": turn_scores, + "reward_extra_info": reward_extra_info, + "raw_prompt": raw_prompt, + "tool_rewards": tool_rewards, + }, + non_tensor_dict={"global_steps": 42}, + ) + + # Verify all data types are accessible + assert torch.all(torch.eq(td["prompts"], prompts)) + assert torch.all(torch.eq(td["responses"], responses)) + assert td["data_source"] == data_source + assert td["uid"] == uid + + # Verify nested structures + assert len(td["turn_scores"]) == 2 + assert list(td["turn_scores"][0]) == [] + assert list(td["turn_scores"][1]) == [0.5, 0.8] + + assert len(td["reward_extra_info"]) == 2 + assert dict(td["reward_extra_info"][0]) == {"acc": 1.0, "loss": 0.1} + + assert len(td["raw_prompt"]) == 2 + assert dict(td["raw_prompt"][0][0]) == {"content": "Compute 4 @ 2", "role": "user"} + + assert len(td["tool_rewards"]) == 2 + assert list(td["tool_rewards"][0]) == [0.0] + assert list(td["tool_rewards"][1]) == [] + + # Verify metadata + assert td["global_steps"] == 42 + + +def test_contiguous(): + # create a tensordict that contains normal tensor, nested tensor, + # nontensorstack with numpy, nontensorstack with tensor, NonTensorData with numpy and NonTensorData with tensor + + a = torch.randn(3, 4) # contiguous tensor + b = torch.randn(3, 4)[:, :-1] # non contiguous tensor + c = torch.nested.as_nested_tensor([torch.randn(3), torch.randn(4), torch.randn(5)], layout=torch.jagged) + + d = torch.randn(10, 12) + e = torch.randn(11, 12) + f = torch.randn(13, 12) + + data = tu.get_tensordict( + tensor_dict={"a": a, "b": b, "c": c, "nt": [{"pixel": d}, {"pixel": e}, {"pixel": f}]}, + non_tensor_dict={"ntd": a.clone()}, + ) + + with pytest.raises(RuntimeError): + # b is not contiguous + data.consolidate() + + data1 = copy.deepcopy(data) + data_cont = tu.contiguous(data1) + + tu.assert_tensordict_eq(data_cont, data) + + data_cont.consolidate() + + tu.assert_tensordict_eq(data_cont, data) diff --git a/verl/tests/tools/__init__.py b/verl/tests/tools/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d828409b82e82a2f9aae138e1150608cf891a667 --- /dev/null +++ b/verl/tests/tools/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/verl/tests/tools/_stub_search_tools.py b/verl/tests/tools/_stub_search_tools.py new file mode 100644 index 0000000000000000000000000000000000000000..c419a7194c078375e75ea790efacab761e4a0d47 --- /dev/null +++ b/verl/tests/tools/_stub_search_tools.py @@ -0,0 +1,47 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Stub native tools for testing +""" + +from __future__ import annotations + +from typing import Any + +from verl.tools.base_tool import BaseTool +from verl.tools.schemas import OpenAIFunctionToolSchema, ToolResponse + + +class StubSearchTool(BaseTool): + def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema): + super().__init__(config, tool_schema) + self.calls: list[dict[str, Any]] = [] + + async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[ToolResponse, float, dict]: + self.calls.append(parameters) + queries = parameters.get("query_list") or [] + text = "; ".join(f"hits-for:{q}" for q in queries) or "hits-for:" + return ToolResponse(text=text), 0.0, {"num_queries": len(queries)} + + +class StubCrawlTool(BaseTool): + def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema): + super().__init__(config, tool_schema) + self.calls: list[dict[str, Any]] = [] + + async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[ToolResponse, float, dict]: + self.calls.append(parameters) + urls = parameters.get("url_list") or [] + text = "crawled:" + ",".join(urls) + return ToolResponse(text=text), 0.0, {"num_urls": len(urls)} diff --git a/verl/tests/tools/test_function_tool_on_cpu.py b/verl/tests/tools/test_function_tool_on_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..63135634c76c0731baaee128611720b8fc011b0c --- /dev/null +++ b/verl/tests/tools/test_function_tool_on_cpu.py @@ -0,0 +1,644 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the function-based tool API.""" + +from __future__ import annotations + +import asyncio +import textwrap +from pathlib import Path +from typing import Literal + +import pytest + +from verl.tools import function_tool as function_tool_mod +from verl.tools.function_tool import ( + FUNCTION_TOOL_REGISTRY, + FunctionTool, + function_tool, + load_function_tools_from_path, + normalize_function_tool_return, +) +from verl.tools.schemas import OpenAIFunctionToolSchema, ToolResponse + + +@pytest.fixture(autouse=True) +def _clean_registry(): + """Reset both the registry and the per-path cache around every test.""" + FUNCTION_TOOL_REGISTRY.clear() + function_tool_mod._LOADED_FUNCTION_TOOL_PATHS.clear() + yield + FUNCTION_TOOL_REGISTRY.clear() + function_tool_mod._LOADED_FUNCTION_TOOL_PATHS.clear() + + +def _write_tool_file(tmp_path: Path, body: str) -> str: + path = tmp_path / "my_tools.py" + path.write_text(textwrap.dedent(body)) + return str(path) + + +# --------------------------------------------------------------------------- +# @function_tool decorator + schema inference +# --------------------------------------------------------------------------- + + +def test_decorator_registers_with_inferred_schema(): + @function_tool("greet") + def greet(name: str, excited: bool = False) -> str: + """Greet someone. + + Args: + name: Person to greet. + excited: Whether to add an exclamation mark. + """ + return f"hi {name}{'!' if excited else ''}" + + assert "greet" in FUNCTION_TOOL_REGISTRY + tool = FUNCTION_TOOL_REGISTRY["greet"] + fn_schema = tool.tool_schema.model_dump(exclude_unset=True, exclude_none=True)["function"] + + assert fn_schema["name"] == "greet" + assert fn_schema["description"].startswith("Greet someone.") + assert fn_schema["parameters"]["properties"]["name"]["type"] == "string" + assert fn_schema["parameters"]["properties"]["name"]["description"] == "Person to greet." + assert fn_schema["parameters"]["properties"]["excited"]["type"] == "boolean" + assert fn_schema["parameters"]["required"] == ["name"] + + +def test_int_float_union_emits_number_type(): + """Numeric unions must let the LLM produce decimals. + + transformers ``get_json_schema`` returns the JSON Schema-standard + ``{"type": ["integer", "number"]}`` for ``int | float``; verl's loosened + schema accepts list-typed ``type`` fields so the LLM is allowed to emit + decimals rather than being told "must be integer". + """ + + @function_tool("num") + def num(x: int | float) -> str: + """Numeric. + + Args: + x: a numeric value. + """ + return str(x) + + params = FUNCTION_TOOL_REGISTRY["num"].tool_schema.model_dump(exclude_unset=True, exclude_none=True)["function"][ + "parameters" + ] + assert "number" in params["properties"]["x"]["type"] + + +def test_int_literal_emits_int_enum(): + """``Literal[1, 2, 3]`` -> JSON ``enum: [1, 2, 3]``. + + Pins the verl schema loosening that allows non-string ``enum`` values + (otherwise pydantic rejects integer literals). + """ + + @function_tool("pick") + def pick(x: Literal[1, 2, 3]) -> str: + """Pick. + + Args: + x: pick one. + """ + return str(x) + + props = FUNCTION_TOOL_REGISTRY["pick"].tool_schema.model_dump(exclude_unset=True, exclude_none=True)["function"][ + "parameters" + ]["properties"] + assert props["x"]["enum"] == [1, 2, 3] + + +def test_decorator_default_name_uses_function_name(): + @function_tool() + def my_special_tool(x: int) -> int: + """Doc. + + Args: + x: A number. + """ + return x + + assert "my_special_tool" in FUNCTION_TOOL_REGISTRY + + +def test_bare_decorator_without_parentheses(): + """``@function_tool`` (no parens) registers under the function name.""" + + @function_tool + def bare_tool(x: int) -> int: + """Doc. + + Args: + x: A number. + """ + return x + + assert "bare_tool" in FUNCTION_TOOL_REGISTRY + fn_schema = FUNCTION_TOOL_REGISTRY["bare_tool"].tool_schema.model_dump(exclude_unset=True, exclude_none=True)[ + "function" + ] + assert fn_schema["name"] == "bare_tool" + # Schema inference path is the same as the parenthesised form. + assert fn_schema["parameters"]["properties"]["x"]["type"] == "integer" + + +def test_explicit_schema_dict_override_skips_inference(): + custom = { + "type": "function", + "function": { + "name": "x", + "description": "custom desc", + "parameters": {"type": "object", "properties": {}, "required": []}, + }, + } + + @function_tool("x", schema=custom) + def x() -> str: + """This docstring should be ignored. + + Args: + ignored: ignored. + """ + return "" + + schema = FUNCTION_TOOL_REGISTRY["x"].tool_schema.model_dump(exclude_unset=True, exclude_none=True) + assert schema["function"]["description"] == "custom desc" + assert schema["function"]["parameters"]["properties"] == {} + + +def test_explicit_schema_object_override(): + custom = OpenAIFunctionToolSchema.model_validate( + { + "type": "function", + "function": { + "name": "y", + "description": "obj desc", + "parameters": {"type": "object", "properties": {}, "required": []}, + }, + } + ) + + @function_tool("y", schema=custom) + def y() -> str: + return "" + + assert FUNCTION_TOOL_REGISTRY["y"].tool_schema is custom + + +def test_duplicate_name_raises(): + @function_tool("dup") + def fn1(x: str) -> str: + """Doc. + + Args: + x: a string. + """ + return x + + with pytest.raises(ValueError, match="already registered"): + + @function_tool("dup") + def fn2(x: str) -> str: + """Doc. + + Args: + x: a string. + """ + return x + + +def test_async_function_marked_is_async(): + @function_tool("aecho") + async def aecho(text: str) -> str: + """Echo. + + Args: + text: text. + """ + return text + + tool = FUNCTION_TOOL_REGISTRY["aecho"] + assert tool.is_async is True + assert asyncio.run(tool.call({"text": "ok"})) == "ok" + + +def test_sync_function_runs_in_thread(): + @function_tool("secho") + def secho(text: str) -> str: + """Echo. + + Args: + text: text. + """ + return text.upper() + + tool = FUNCTION_TOOL_REGISTRY["secho"] + assert tool.is_async is False + assert asyncio.run(tool.call({"text": "hi"})) == "HI" + + +def test_missing_docstring_raises_at_registration(): + """Schema inference is delegated to ``transformers.get_json_schema``, + which raises ``DocstringParsingException`` when the function has no + docstring at all. We verify the contract surfaces, not the exact + exception type, to avoid coupling tests to transformers internals. + """ + + with pytest.raises(Exception, match=r"no docstring"): + + @function_tool("nodoc") + def nodoc(x: str) -> str: + return x + + +def test_missing_type_hint_raises_at_registration(): + """``transformers.get_json_schema`` raises when a parameter is unannotated.""" + + with pytest.raises(Exception, match=r"missing a type hint"): + + @function_tool("untyped") + def untyped(x) -> str: + """Doc. + + Args: + x: a thing. + """ + return x + + +def test_missing_arg_description_raises_at_registration(): + """``transformers.get_json_schema`` requires every parameter to be + described in the docstring's ``Args:`` block.""" + + with pytest.raises(Exception, match=r"no description for the argument"): + + @function_tool("partial") + def partial(x: int, y: int) -> int: + """Add. + + Args: + x: only x described. + """ + return x + y + + +def test_var_args_raises_at_registration(): + """``*args`` / ``**kwargs`` can't be expressed as fixed JSON properties. + + We catch this before ``get_json_schema`` so the user gets a verl-specific + pointer to the right fix (``param: list[T]``) instead of transformers' + less actionable "missing type hint for args". + """ + + with pytest.raises(ValueError, match=r"variadic parameter"): + + @function_tool("varargs") + def varargs(x: int, *args: int) -> int: + """Sum. + + Args: + x: x. + """ + return x + sum(args) + + with pytest.raises(ValueError, match=r"variadic parameter"): + + @function_tool("varkw") + def varkw(x: int, **kwargs: int) -> int: + """Sum. + + Args: + x: x. + """ + return x + sum(kwargs.values()) + + +# --------------------------------------------------------------------------- +# load_function_tools_from_path +# --------------------------------------------------------------------------- + + +def test_load_basic_returns_registered_tools(tmp_path): + path = _write_tool_file( + tmp_path, + """ + from verl.tools.function_tool import function_tool + + @function_tool("greet") + def greet(name: str) -> str: + '''Greet someone. + + Args: + name: who to greet. + ''' + return f"hi {name}" + """, + ) + + tools = load_function_tools_from_path(path) + assert [t.name for t in tools] == ["greet"] + assert FUNCTION_TOOL_REGISTRY["greet"] is tools[0] + + +def test_load_multiple_tools(tmp_path): + path = _write_tool_file( + tmp_path, + """ + from verl.tools.function_tool import function_tool + + @function_tool("a") + def a(x: str) -> str: + '''A. + + Args: + x: x. + ''' + return x + + @function_tool("b") + def b(x: str) -> str: + '''B. + + Args: + x: x. + ''' + return x + """, + ) + tools = load_function_tools_from_path(path) + assert sorted(t.name for t in tools) == ["a", "b"] + + +def test_missing_path_raises(): + with pytest.raises(FileNotFoundError, match="function_tool_path does not exist"): + load_function_tools_from_path("/nonexistent/path/here.py") + + +def test_no_decorator_logs_warning(tmp_path, caplog): + path = _write_tool_file(tmp_path, "x = 1\n") + with caplog.at_level("WARNING"): + tools = load_function_tools_from_path(path) + assert tools == [] + assert any("no @function_tool decorators found" in rec.getMessage() for rec in caplog.records) + + +def test_load_is_idempotent_across_calls(tmp_path): + """Loading the same path twice in one process must be a no-op. + + Production calls ``load_function_tools_from_path`` exactly once per + worker process (from ``AgentLoopWorker.__init__``), so this is not a + hot-path concern there. The contract still matters for tests, custom + managers, or any code that re-enters the loader: without the + :data:`_LOADED_FUNCTION_TOOL_PATHS` cache, the second call would + re-exec the user file, the ``@function_tool`` decorator would run + again with a *new* function object for the same name, and the + decorator's dup-name guard would raise ``ValueError``. + """ + path = _write_tool_file( + tmp_path, + """ + from verl.tools.function_tool import function_tool + + @function_tool("idem") + def idem(x: str) -> str: + '''Idem. + + Args: + x: x. + ''' + return x + """, + ) + + first = load_function_tools_from_path(path) + second = load_function_tools_from_path(path) + + assert [t.name for t in first] == ["idem"] + assert second[0] is first[0] + assert second[0].fn is first[0].fn + + +def test_load_returns_only_tools_added_by_this_file(tmp_path): + """Pre-registering a tool from elsewhere must not leak into the loader's + return value; the loader attributes only what its file added.""" + + @function_tool("preexisting") + def preexisting(x: str) -> str: + """Pre. + + Args: + x: x. + """ + return x + + path = _write_tool_file( + tmp_path, + """ + from verl.tools.function_tool import function_tool + + @function_tool("only_mine") + def only_mine(x: str) -> str: + '''Mine. + + Args: + x: x. + ''' + return x + """, + ) + + tools = load_function_tools_from_path(path) + assert [t.name for t in tools] == ["only_mine"] + assert "preexisting" in FUNCTION_TOOL_REGISTRY + + +def test_relative_path_resolved_against_cwd(tmp_path, monkeypatch): + path_str = _write_tool_file( + tmp_path, + """ + from verl.tools.function_tool import function_tool + + @function_tool("rel") + def rel(x: str) -> str: + '''Rel. + + Args: + x: x. + ''' + return x + """, + ) + monkeypatch.chdir(tmp_path) + tools = load_function_tools_from_path(Path(path_str).name) + assert [t.name for t in tools] == ["rel"] + + +# --------------------------------------------------------------------------- +# normalize_function_tool_return +# --------------------------------------------------------------------------- + + +def test_normalize_str(): + resp, reward, metrics = normalize_function_tool_return("hello") + assert resp == ToolResponse(text="hello") + assert reward == 0.0 + assert metrics == {} + + +def test_normalize_tool_response_passthrough(): + src = ToolResponse(text="x") + resp, reward, metrics = normalize_function_tool_return(src) + assert resp is src + assert reward == 0.0 + assert metrics == {} + + +def test_normalize_dict_serialized_as_json(): + resp, _, _ = normalize_function_tool_return({"a": 1, "b": "two"}) + assert "a" in resp.text and "two" in resp.text + + +def test_normalize_2_tuple_carries_reward(): + resp, reward, metrics = normalize_function_tool_return(("text", 1.5)) + assert resp.text == "text" + assert reward == 1.5 + assert metrics == {} + + +def test_normalize_3_tuple_carries_metrics(): + resp, reward, metrics = normalize_function_tool_return(("text", 2.0, {"k": "v"})) + assert resp.text == "text" + assert reward == 2.0 + assert metrics == {"k": "v"} + + +def test_normalize_tuple_tolerates_none_reward_and_metrics(): + """Tools may legitimately omit reward/metrics by returning ``None``.""" + resp, reward, metrics = normalize_function_tool_return(("text", None)) + assert resp.text == "text" + assert reward == 0.0 + assert metrics == {} + + resp, reward, metrics = normalize_function_tool_return(("text", None, None)) + assert resp.text == "text" + assert reward == 0.0 + assert metrics == {} + + +def test_normalize_falsy_reward_is_preserved_not_coerced_to_default(): + """Regression: detect ``None`` via ``is None``, not truthiness. + + The earlier ``ret[1] or 0.0`` form swallowed every falsy reward value + including ``False`` and integer ``0``, so a tool that legitimately + reported "no progress this turn" via ``reward=0`` or ``reward=False`` + was indistinguishable from one that returned ``reward=None`` -- and + more importantly, distinct from the ``or``-fallback path semantically. + """ + # int 0 is the canonical "no signal" reward; must round-trip as 0.0, + # and importantly come out of the ``int -> float`` branch, not the + # ``or 0.0`` branch. + _, reward, _ = normalize_function_tool_return(("t", 0)) + assert reward == 0.0 + assert isinstance(reward, float) + + # bool is a subclass of int so ``False`` is technically valid here. + _, reward, _ = normalize_function_tool_return(("t", False)) + assert reward == 0.0 + + +def test_normalize_tuple_of_invalid_length_raises(): + """0-length and >=4-length tuples almost always indicate a bug; we + refuse rather than silently ``str(ret)``-ing the entire tuple, which + would corrupt the ToolResponse shown to the LLM.""" + with pytest.raises(TypeError, match=r"length 1, 2, or 3"): + normalize_function_tool_return(()) + + with pytest.raises(TypeError, match=r"length 1, 2, or 3"): + normalize_function_tool_return(("a", 1, {}, "extra")) + + +def test_normalize_arbitrary_object_falls_back_to_str(): + class Foo: + def __str__(self) -> str: + return "FOO" + + resp, reward, metrics = normalize_function_tool_return(Foo()) + assert resp.text == "FOO" + assert reward == 0.0 + assert metrics == {} + + +# --------------------------------------------------------------------------- +# Rollout config field +# --------------------------------------------------------------------------- + + +def test_rollout_yaml_exposes_function_tool_path(): + """Smoke test the YAML default so ``ToolAgentLoop.__init__`` can read it.""" + from omegaconf import OmegaConf + + cfg = OmegaConf.load("verl/trainer/config/rollout/rollout.yaml") + assert "function_tool_path" in cfg.multi_turn + assert cfg.multi_turn.function_tool_path is None + + +class _HydraProbe: + def __init__(self, tools): + self.tools = tools + + +def test_tool_list_wrap_survives_hydra_instantiate(tmp_path): + """Without ``ToolListWrap``, ``hydra.utils.instantiate`` demotes each + ``FunctionTool`` in a kwarg list to ``DictConfig`` and breaks + ``isinstance(tool, FunctionTool)`` in ``ToolAgentLoop._call_tool``.""" + import hydra + + from verl.experimental.agent_loop.agent_loop import ToolListWrap + + path = _write_tool_file( + tmp_path, + """ + from verl.tools.function_tool import function_tool + + @function_tool + def probe(text: str) -> str: + '''Probe. + + Args: + text: text. + ''' + return text + """, + ) + tools = load_function_tools_from_path(path) + assert all(isinstance(t, FunctionTool) for t in tools) + + # Without the wrap: hydra demotes each FunctionTool to DictConfig. If this + # ever stops being true, ToolListWrap is obsolete and can be deleted. + naked = hydra.utils.instantiate({"_target_": f"{__name__}._HydraProbe"}, tools=tools) + assert not all(isinstance(t, FunctionTool) for t in naked.tools), ( + "hydra.utils.instantiate no longer demotes FunctionTool to DictConfig; ToolListWrap may be obsolete." + ) + + wrapped = hydra.utils.instantiate( + {"_target_": f"{__name__}._HydraProbe"}, + tools=ToolListWrap(tools), + ) + assert isinstance(wrapped.tools, ToolListWrap) + assert all(isinstance(t, FunctionTool) for t in wrapped.tools.tools) + assert callable(wrapped.tools.tools[0].fn) diff --git a/verl/tests/tools/test_mixed_tools_on_cpu.py b/verl/tests/tools/test_mixed_tools_on_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..c0717c7d8806ccc1d233186426a130cfb8eef80a --- /dev/null +++ b/verl/tests/tools/test_mixed_tools_on_cpu.py @@ -0,0 +1,338 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Coexistence test: yaml-defined native tools + ``@function_tool`` tools. + +Pins :func:`verl.tools.tool_registry.load_all_tools`, the loader both +``AgentLoopWorker`` and ``RLHFDataset`` use. The yaml mirrors +``recipe/search_agent/config/all_tool_config.yaml`` but points at CPU-only +``BaseTool`` stubs in ``tests.tools._stub_search_tools``. +""" + +from __future__ import annotations + +import asyncio +import textwrap +from pathlib import Path + +import pytest + +from tests.tools._stub_search_tools import StubCrawlTool, StubSearchTool +from verl.tools import function_tool as function_tool_mod +from verl.tools.base_tool import BaseTool +from verl.tools.function_tool import ( + FUNCTION_TOOL_REGISTRY, + FunctionTool, + normalize_function_tool_return, +) +from verl.tools.tool_registry import load_all_tools + + +@pytest.fixture(autouse=True) +def _clean_registry(): + """Reset the process-global registry + per-path cache around each test.""" + FUNCTION_TOOL_REGISTRY.clear() + function_tool_mod._LOADED_FUNCTION_TOOL_PATHS.clear() + yield + FUNCTION_TOOL_REGISTRY.clear() + function_tool_mod._LOADED_FUNCTION_TOOL_PATHS.clear() + + +_NATIVE_TOOL_YAML = """\ +tools: + - class_name: "tests.tools._stub_search_tools.StubSearchTool" + config: + retrieval_service_url: http://stub/retrieve + topk: 3 + type: native + tool_schema: + type: function + function: + name: search + description: Stub web search. + parameters: + type: object + properties: + query_list: + type: array + description: A list of fully-formed semantic queries. + required: ["query_list"] + + - class_name: "tests.tools._stub_search_tools.StubCrawlTool" + config: + crawl_service_url: http://stub/crawl + type: native + tool_schema: + type: function + function: + name: crawler + description: Stub crawl. + parameters: + type: object + properties: + url_list: + type: array + description: URLs to crawl. + required: ["url_list"] +""" + +_FUNCTION_TOOL_BODY = """\ +from verl.tools.function_tool import function_tool + +@function_tool("get_weather") +def get_weather(city: str) -> dict: + '''Get the current weather for a city. + + Args: + city: the city to look up. + ''' + return {"temperature_c": 17.3, "condition": "drizzle"} + +@function_tool("calculator") +def calculator(expression: str) -> str: + '''Evaluate an arithmetic expression. + + Args: + expression: a python-style arithmetic expression. + ''' + return str(eval(expression, {"__builtins__": {}}, {})) +""" + + +@pytest.fixture +def native_yaml_path(tmp_path: Path) -> str: + p = tmp_path / "all_tool_config.yaml" + p.write_text(_NATIVE_TOOL_YAML) + return str(p) + + +@pytest.fixture +def function_tool_py_path(tmp_path: Path) -> str: + p = tmp_path / "my_function_tools.py" + p.write_text(textwrap.dedent(_FUNCTION_TOOL_BODY)) + return str(p) + + +def _load_as_dict(native_yaml: str | None, function_path: str | None) -> dict[str, BaseTool | FunctionTool]: + """``load_all_tools`` keyed by tool name.""" + tools = load_all_tools(tool_config_path=native_yaml, function_tool_path=function_path) + return {tool.name: tool for tool in tools} + + +def test_no_paths_returns_empty(): + """Both args ``None`` is the "tools disabled" path; must not blow up.""" + assert load_all_tools(tool_config_path=None, function_tool_path=None) == [] + + +def test_native_only_loader(native_yaml_path): + """Sanity: the stub yaml alone yields exactly the two BaseTool instances.""" + tools = _load_as_dict(native_yaml_path, None) + + assert sorted(tools) == ["crawler", "search"] + assert isinstance(tools["search"], StubSearchTool) + assert isinstance(tools["crawler"], StubCrawlTool) + # No function tools should have been registered as a side effect. + assert FUNCTION_TOOL_REGISTRY == {} + + +def test_function_only_loader(function_tool_py_path): + """Sanity: the function-tool file alone yields FunctionTool instances.""" + tools = _load_as_dict(None, function_tool_py_path) + + assert sorted(tools) == ["calculator", "get_weather"] + assert isinstance(tools["get_weather"], FunctionTool) + assert isinstance(tools["calculator"], FunctionTool) + + +def test_native_and_function_tools_coexist(native_yaml_path, function_tool_py_path): + """All four tools land in one mapping with the right concrete types.""" + tools = _load_as_dict(native_yaml_path, function_tool_py_path) + + assert sorted(tools) == ["calculator", "crawler", "get_weather", "search"] + assert isinstance(tools["search"], StubSearchTool) + assert isinstance(tools["crawler"], StubCrawlTool) + assert isinstance(tools["get_weather"], FunctionTool) + assert isinstance(tools["calculator"], FunctionTool) + # ToolAgentLoop._call_tool dispatches via isinstance(tool, FunctionTool), + # so FunctionTool must not subclass BaseTool. + assert not isinstance(tools["get_weather"], BaseTool) + assert not isinstance(tools["calculator"], BaseTool) + + +def test_merged_schemas_are_well_formed(native_yaml_path, function_tool_py_path): + """Each tool exposes a valid OpenAI function schema.""" + tools = _load_as_dict(native_yaml_path, function_tool_py_path) + + schemas = {name: tool.tool_schema.model_dump(exclude_unset=True, exclude_none=True) for name, tool in tools.items()} + + for name, sch in schemas.items(): + assert sch["type"] == "function", name + assert sch["function"]["name"] == name, name + assert sch["function"]["parameters"]["type"] == "object", name + + # Native tool params come from yaml. + assert "query_list" in schemas["search"]["function"]["parameters"]["properties"] + assert set(schemas["crawler"]["function"]["parameters"]["properties"]) == {"url_list"} + + # Function tool params come from signature + docstring inference. + weather_props = schemas["get_weather"]["function"]["parameters"]["properties"] + assert weather_props["city"]["type"] == "string" + assert weather_props["city"]["description"] == "the city to look up." + + +def test_dispatch_branches_match_tool_agent_loop(native_yaml_path, function_tool_py_path): + """Drive both ``ToolAgentLoop._call_tool`` branches end-to-end.""" + tools = _load_as_dict(native_yaml_path, function_tool_py_path) + + async def _drive(): + # function tool branch -- get_weather returns dict, exercising the + # dict -> JSON path of normalize_function_tool_return. + weather_raw = await tools["get_weather"].call({"city": "Tokyo"}) + weather_resp, weather_reward, weather_metrics = normalize_function_tool_return(weather_raw) + assert "17.3" in weather_resp.text and "drizzle" in weather_resp.text + assert weather_reward == 0.0 + assert weather_metrics == {} + + calc_raw = await tools["calculator"].call({"expression": "2 + 3 * 4"}) + calc_resp, _, _ = normalize_function_tool_return(calc_raw) + assert calc_resp.text == "14" + + # --- BaseTool branch (create → execute → release) --- + search_tool: StubSearchTool = tools["search"] + instance_id, _ = await search_tool.create() + s_resp, s_reward, s_metrics = await search_tool.execute(instance_id, {"query_list": ["foo", "bar"]}) + await search_tool.release(instance_id) + assert "hits-for:foo" in s_resp.text and "hits-for:bar" in s_resp.text + assert s_reward == 0.0 + assert s_metrics == {"num_queries": 2} + assert search_tool.calls == [{"query_list": ["foo", "bar"]}] + + crawl_tool: StubCrawlTool = tools["crawler"] + instance_id, _ = await crawl_tool.create() + c_resp, _, c_metrics = await crawl_tool.execute(instance_id, {"url_list": ["http://a", "http://b"]}) + await crawl_tool.release(instance_id) + assert c_resp.text == "crawled:http://a,http://b" + assert c_metrics == {"num_urls": 2} + + asyncio.run(_drive()) + + +def test_loader_is_safe_to_call_repeatedly(native_yaml_path, function_tool_py_path): + """``load_all_tools`` runs at least twice per process (worker + dataset).""" + runs = [_load_as_dict(native_yaml_path, function_tool_py_path) for _ in range(3)] + + for tools in runs: + assert sorted(tools) == ["calculator", "crawler", "get_weather", "search"] + assert isinstance(tools["get_weather"], FunctionTool) + assert isinstance(tools["calculator"], FunctionTool) + assert isinstance(tools["search"], StubSearchTool) + assert isinstance(tools["crawler"], StubCrawlTool) + + # Function tools must be cached (re-execing the file would double-register). + assert runs[0]["get_weather"] is runs[1]["get_weather"] is runs[2]["get_weather"] + assert runs[0]["calculator"] is runs[1]["calculator"] is runs[2]["calculator"] + # Native tools, by contrast, are re-instantiated per call. Pin the asymmetry. + assert runs[0]["search"] is not runs[1]["search"] + + +def test_function_tool_name_collision_with_native_tool_raises(native_yaml_path, tmp_path): + """A function tool sharing a name with a native tool must fail loudly.""" + colliding_path = tmp_path / "colliding.py" + colliding_path.write_text( + textwrap.dedent( + """ + from verl.tools.function_tool import function_tool + + @function_tool("search") + def search(query_list: list) -> str: + '''Stub function tool deliberately reusing the native name. + + Args: + query_list: queries. + ''' + return "from-function-tool" + """ + ) + ) + + with pytest.raises(ValueError, match=r"\['search'\].*collide"): + _load_as_dict(native_yaml_path, str(colliding_path)) + + +def test_function_tool_name_collision_reports_all_offenders(native_yaml_path, tmp_path): + """Multiple collisions are surfaced together, sorted.""" + colliding_path = tmp_path / "many_colliding.py" + colliding_path.write_text( + textwrap.dedent( + """ + from verl.tools.function_tool import function_tool + + @function_tool("crawler") + def crawler(url_list: list) -> str: + '''Collides. + + Args: + url_list: urls. + ''' + return "" + + @function_tool("search") + def search(query_list: list) -> str: + '''Collides. + + Args: + query_list: queries. + ''' + return "" + + @function_tool("ok_unique") + def ok_unique(x: str) -> str: + '''Fine. + + Args: + x: x. + ''' + return x + """ + ) + ) + + with pytest.raises(ValueError, match=r"\['crawler', 'search'\].*collide"): + _load_as_dict(native_yaml_path, str(colliding_path)) + + +def test_dataset_loader_sees_function_tools(native_yaml_path, function_tool_py_path): + """RLHFDataset and AgentLoopWorker must see the same tool schemas, else + prompt-length filtering and rollout disagree.""" + from omegaconf import OmegaConf + + from verl.utils.dataset.rl_dataset import RLHFDataset + + cfg = OmegaConf.create( + { + "tool_config_path": native_yaml_path, + "function_tool_path": function_tool_py_path, + } + ) + # Skip __init__ I/O; just exercise the tool-loading branch. + ds = RLHFDataset.__new__(RLHFDataset) + ds.config = cfg + ds.tool_config_path = cfg.tool_config_path + ds.function_tool_path = cfg.function_tool_path + + tools = load_all_tools( + tool_config_path=ds.tool_config_path, + function_tool_path=ds.function_tool_path, + ) + schema_names = sorted(t.tool_schema.function.name for t in tools) + assert schema_names == ["calculator", "crawler", "get_weather", "search"] diff --git a/verl/tests/trainer/__init__.py b/verl/tests/trainer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6f79d474d156e16ae54bb3d0c8f9ae7d0e16946e --- /dev/null +++ b/verl/tests/trainer/__init__.py @@ -0,0 +1,16 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Tests for the trainer module. +""" diff --git a/verl/tests/trainer/ppo/__init__.py b/verl/tests/trainer/ppo/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..26d7c04fc335c873ef77f8989e82e4239be7dba1 --- /dev/null +++ b/verl/tests/trainer/ppo/__init__.py @@ -0,0 +1,16 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Tests for the PPO trainer module. +""" diff --git a/verl/tests/trainer/ppo/test_core_algos_on_cpu.py b/verl/tests/trainer/ppo/test_core_algos_on_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..36c8bc8d0a8f460929bc3d227a87074633830134 --- /dev/null +++ b/verl/tests/trainer/ppo/test_core_algos_on_cpu.py @@ -0,0 +1,385 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import random +import unittest + +import numpy as np +import pytest +import torch + +import verl.trainer.ppo.core_algos +from verl.trainer.ppo.core_algos import ( + compute_gae_advantage_return, + compute_grpo_outcome_advantage, + compute_grpo_vectorized_outcome_advantage, + compute_rloo_outcome_advantage, + compute_rloo_vectorized_outcome_advantage, + get_adv_estimator_fn, + kl_penalty, + register_adv_est, +) + + +def mock_test_fn(): + pass + + +class TestRegisterAdvEst(unittest.TestCase): + def setUp(self): + """Clear the registry before each test""" + verl.trainer.ppo.core_algos.ADV_ESTIMATOR_REGISTRY.clear() + verl.trainer.ppo.core_algos.ADV_ESTIMATOR_REGISTRY = { + "gae": lambda x: x * 2, + "vtrace": lambda x: x + 1, + } + self.ADV_ESTIMATOR_REGISTRY = verl.trainer.ppo.core_algos.ADV_ESTIMATOR_REGISTRY + + def tearDown(self) -> None: + verl.trainer.ppo.core_algos.ADV_ESTIMATOR_REGISTRY.clear() + return super().tearDown() + + def test_register_new_function(self): + """Test registering a new function with a string name""" + + @register_adv_est("test_estimator") + def test_fn(): + pass + + self.assertIn("test_estimator", self.ADV_ESTIMATOR_REGISTRY) + self.assertEqual(self.ADV_ESTIMATOR_REGISTRY["test_estimator"], test_fn) + + def test_register_with_enum(self): + """Test registering with an enum value (assuming AdvantageEstimator exists)""" + from enum import Enum + + class AdvantageEstimator(Enum): + TEST = "test_enum_estimator" + + @register_adv_est(AdvantageEstimator.TEST) + def test_fn(): + pass + + self.assertIn("test_enum_estimator", self.ADV_ESTIMATOR_REGISTRY) + self.assertEqual(self.ADV_ESTIMATOR_REGISTRY["test_enum_estimator"], test_fn) + + def test_duplicate_registration_same_function(self): + """Test that registering the same function twice doesn't raise an error""" + register_adv_est("duplicate_test")(mock_test_fn) + register_adv_est("duplicate_test")(mock_test_fn) + + self.assertEqual(self.ADV_ESTIMATOR_REGISTRY["duplicate_test"], mock_test_fn) + + def test_duplicate_registration_different_function(self): + """Test that registering different functions with same name raises ValueError""" + + @register_adv_est("conflict_test") + def test_fn1(): + pass + + with self.assertRaises(ValueError): + + @register_adv_est("conflict_test") + def test_fn2(): + pass + + def test_decorator_preserves_function(self): + """Test that the decorator returns the original function""" + + def test_fn(): + return "original" + + decorated = register_adv_est("preserve_test")(test_fn) + self.assertEqual(decorated(), "original") + + def test_multiple_registrations(self): + """Test registering multiple different functions""" + init_adv_count = len(self.ADV_ESTIMATOR_REGISTRY) + + @register_adv_est("estimator1") + def fn1(): + pass + + @register_adv_est("estimator2") + def fn2(): + pass + + self.assertEqual(len(self.ADV_ESTIMATOR_REGISTRY), 2 + init_adv_count) + self.assertEqual(self.ADV_ESTIMATOR_REGISTRY["estimator1"], fn1) + self.assertEqual(self.ADV_ESTIMATOR_REGISTRY["estimator2"], fn2) + + def test_get_adv_estimator_fn_valid_names(self): + """Test that valid names return the correct function from registry.""" + # Test GAE + gae_fn = get_adv_estimator_fn("gae") + assert gae_fn(5) == 10 # 5 * 2 = 10 + + # Test Vtrace + vtrace_fn = get_adv_estimator_fn("vtrace") + assert vtrace_fn(5) == 6 # 5 + 1 = 6 + + def test_get_adv_estimator_fn_invalid_name(self): + """Test that invalid names raise ValueError.""" + with pytest.raises(ValueError) as excinfo: + get_adv_estimator_fn("invalid_name") + assert "Unknown advantage estimator simply: invalid_name" in str(excinfo.value) + + def test_get_adv_estimator_fn_case_sensitive(self): + """Test that name lookup is case-sensitive.""" + with pytest.raises(ValueError): + get_adv_estimator_fn("GAE") # Different case + + +def test_multi_turn_compute_gae_advantage_return(): + """Test multi-turn GAE skip observation tokens.""" + gamma = random.uniform(0.0, 1.0) + lam = random.uniform(0.0, 1.0) + + rewards = torch.tensor([[0.0, 0.0, 0.1, 0.1, 0.1, 0.0, 0.0, 0.1, 1.0, 0.0, 0.0]], dtype=torch.float) + + values1 = torch.tensor( + [ + [ + random.uniform(-100.0, 100.0), + random.random(), + 4.0, + 5.0, + 6.0, + random.uniform(-100.0, 0), + random.random(), + 7.0, + 9.0, + 0.0, + 0.0, + ] + ], + dtype=torch.float, + ) + + values2 = torch.tensor( + [ + [ + random.random(), + random.uniform(-100.0, 100.0), + 4.0, + 5.0, + 6.0, + random.random(), + random.uniform(0.0, 100.0), + 7.0, + 9.0, + 0.0, + 0.0, + ] + ], + dtype=torch.float, + ) + + response_mask = torch.tensor([[0, 0, 1, 1, 1, 0, 0, 1, 1, 0, 0]], dtype=torch.float) + + adv1, ret1 = compute_gae_advantage_return(rewards, values1, response_mask, gamma, lam) + adv2, ret2 = compute_gae_advantage_return(rewards, values2, response_mask, gamma, lam) + + ret1 *= response_mask + ret2 *= response_mask + assert torch.equal(adv1, adv2), f"{adv1=}, {adv2=}" + assert torch.equal(ret1, ret2), f"{ret1=}, {ret2=}" + print(f" [CORRECT] \n\n{adv1=}, \n\n{ret1=}") + + +def _make_group_index(batch_size: int, num_groups: int) -> np.ndarray: + """Create a numpy index array ensuring each group has at least 2 samples.""" + assert num_groups * 2 <= batch_size, "batch_size must allow >=2 samples per group" + counts: list[int] = [2] * num_groups + remaining = batch_size - 2 * num_groups + for _ in range(remaining): + counts[random.randrange(num_groups)] += 1 + index = [] + for gid, c in enumerate(counts): + index.extend([gid] * c) + random.shuffle(index) + return np.asarray(index, dtype=np.int64) + + +def _rand_mask(batch_size: int, seq_len: int) -> torch.Tensor: + mask = torch.randint(0, 2, (batch_size, seq_len), dtype=torch.int64).float() + rows_without_one = (mask.sum(dim=-1) == 0).nonzero(as_tuple=True)[0] + if len(rows_without_one) > 0: + mask[rows_without_one, -1] = 1.0 + return mask + + +@pytest.mark.parametrize( + "batch_size,seq_len,num_groups,seed", + [ + (64, 128, 5, 0), + (128, 256, 8, 1), + (512, 512, 10, 2), + ], +) +def test_rloo_and_vectorized_equivalence(batch_size: int, seq_len: int, num_groups: int, seed: int): + torch.manual_seed(seed) + random.seed(seed) + np.random.seed(seed) + index = _make_group_index(batch_size, num_groups) + response_mask = _rand_mask(batch_size, seq_len) + base_rewards = torch.randn(batch_size, seq_len, dtype=torch.float32) + token_level_rewards = base_rewards * response_mask + adv1, ret1 = compute_rloo_outcome_advantage( + token_level_rewards=token_level_rewards, + response_mask=response_mask, + index=index, + ) + adv2, ret2 = compute_rloo_vectorized_outcome_advantage( + token_level_rewards=token_level_rewards, + response_mask=response_mask, + index=index, + ) + # Print concise diagnostics for visibility during test runs + adv_max_diff = (adv1 - adv2).abs().max().item() + ret_max_diff = (ret1 - ret2).abs().max().item() + total_mask_tokens = int(response_mask.sum().item()) + print( + f"[RLOO] seed={seed} groups={num_groups} shape={adv1.shape} " + f"mask_tokens={total_mask_tokens} adv_max_diff={adv_max_diff:.3e} ret_max_diff={ret_max_diff:.3e}" + ) + assert adv1.shape == adv2.shape == (batch_size, seq_len) + assert ret1.shape == ret2.shape == (batch_size, seq_len) + assert torch.allclose(adv1, adv2, rtol=1e-5, atol=1e-6) + assert torch.allclose(ret1, ret2, rtol=1e-5, atol=1e-6) + + +def test_grpo_vectorized_matches_original_for_low_variance_rewards(): + token_level_rewards = torch.tensor([[1.0], [1.00001], [2.0], [2.00001]], dtype=torch.float32) + response_mask = torch.ones_like(token_level_rewards) + index = np.array(["prompt-a", "prompt-a", "prompt-b", "prompt-b"], dtype=object) + + adv1, ret1 = compute_grpo_outcome_advantage( + token_level_rewards=token_level_rewards, + response_mask=response_mask, + index=index, + ) + adv2, ret2 = compute_grpo_vectorized_outcome_advantage( + token_level_rewards=token_level_rewards, + response_mask=response_mask, + index=index, + ) + + assert torch.allclose(adv1, adv2, rtol=1e-5, atol=1e-6) + assert torch.allclose(ret1, ret2, rtol=1e-5, atol=1e-6) + + +@pytest.mark.parametrize( + "batch_size,seq_len,num_groups,seed", + [ + (64, 128, 5, 0), + (128, 256, 8, 1), + (512, 512, 10, 2), + ], +) +def test_grpo_and_vectorized_equivalence(batch_size: int, seq_len: int, num_groups: int, seed: int): + # Set seeds for reproducibility + torch.manual_seed(seed) + random.seed(seed) + np.random.seed(seed) + + # Generate group indices (numpy array of shape [batch_size]) + index = _make_group_index(batch_size, num_groups) + + # Generate binary response mask (at least one valid token per row) + response_mask = _rand_mask(batch_size, seq_len) + + # Generate token-level rewards and apply mask + base_rewards = torch.randn(batch_size, seq_len, dtype=torch.float32) + token_level_rewards = base_rewards * response_mask + + # Compute GRPO outcome advantage (original implementation) + adv1, ret1 = compute_grpo_outcome_advantage( + token_level_rewards=token_level_rewards, + response_mask=response_mask, + index=index, + ) + + # Compute GRPO outcome advantage (vectorized implementation) + adv2, ret2 = compute_grpo_vectorized_outcome_advantage( + token_level_rewards=token_level_rewards, + response_mask=response_mask, + index=index, + ) + + # Diagnostic info for visibility (same style as RLOO test) + adv_max_diff = (adv1 - adv2).abs().max().item() + ret_max_diff = (ret1 - ret2).abs().max().item() + total_mask_tokens = int(response_mask.sum().item()) + print( + f"[GRPO] seed={seed} groups={num_groups} shape={adv1.shape} " + f"mask_tokens={total_mask_tokens} adv_max_diff={adv_max_diff:.3e} ret_max_diff={ret_max_diff:.3e}" + ) + + # Assert shape and numerical equivalence + assert adv1.shape == adv2.shape == (batch_size, seq_len) + assert ret1.shape == ret2.shape == (batch_size, seq_len) + assert torch.allclose(adv1, adv2, rtol=1e-5, atol=1e-6) + assert torch.allclose(ret1, ret2, rtol=1e-5, atol=1e-6) + + +@pytest.mark.parametrize( + "name,base", + [ + ("k1+", "k1"), + ("kl+", "kl"), + ("abs+", "abs"), + ("k3+", "k3"), + ("low_var_kl+", "low_var_kl"), + ], +) +def test_kl_penalty_straight_through_value_matches_base(name, base): + """The ``+`` suffix is a straight-through trick that swaps in the k2 + gradient while keeping the base estimator's value. Therefore the forward + value of e.g. ``k3+`` must match the value of plain ``k3``. + + Regression test for the bug where ``kl_penalty(..., "k3+")`` raised + ``NotImplementedError`` because the wrapper forwarded the ``+`` suffix to + ``kl_penalty_forward`` without stripping it. + """ + torch.manual_seed(0) + logprob = torch.randn(4, 8, requires_grad=True) + ref_logprob = torch.randn(4, 8) + + plus_value = kl_penalty(logprob, ref_logprob, name) + base_value = kl_penalty(logprob, ref_logprob, base) + assert torch.allclose(plus_value, base_value) + + +def test_kl_penalty_k3_plus_uses_k2_gradient(): + """With ``k3+`` the gradient w.r.t. ``logprob`` should equal the gradient + obtained from the ``k2`` (``0.5 * log_ratio**2``) estimator, since the + straight-through trick routes the backward pass through ``k2``. + """ + torch.manual_seed(0) + logprob = torch.randn(4, 8, requires_grad=True) + ref_logprob = torch.randn(4, 8) + + out_plus = kl_penalty(logprob, ref_logprob, "k3+").sum() + (grad_plus,) = torch.autograd.grad(out_plus, logprob) + + logprob_k2 = logprob.detach().clone().requires_grad_(True) + out_k2 = kl_penalty(logprob_k2, ref_logprob, "k2").sum() + (grad_k2,) = torch.autograd.grad(out_k2, logprob_k2) + + assert torch.allclose(grad_plus, grad_k2) + + +if __name__ == "__main__": + unittest.main() diff --git a/verl/tests/trainer/ppo/test_metric_utils_on_cpu.py b/verl/tests/trainer/ppo/test_metric_utils_on_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..0bdda88df7adf4161e8ed2c80dee61dc7c9f65ac --- /dev/null +++ b/verl/tests/trainer/ppo/test_metric_utils_on_cpu.py @@ -0,0 +1,547 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Tests for the metric utilities in verl.trainer.ppo.metric_utils. +""" + +import unittest +from unittest.mock import MagicMock, patch + +import numpy as np +import torch + +from verl.trainer.ppo.metric_utils import ( + bootstrap_metric, + calc_maj_val, + compute_data_metrics, + compute_throughout_metrics, + compute_timing_metrics, + process_validation_metrics, +) +from verl.utils.metric import ( + reduce_metrics, +) +from verl.utils.metric.utils import ( + AggregationType, + Metric, +) + + +class TestReduceMetrics(unittest.TestCase): + """Tests for the reduce_metrics function.""" + + def test_reduce_metrics_basic(self): + """Test that reduce_metrics correctly computes means.""" + metrics = { + "loss": [1.0, 2.0, 3.0], + "accuracy": [0.0, 0.5, 1.0], + } + result = reduce_metrics(metrics) + + self.assertEqual(result["loss"], 2.0) + self.assertEqual(result["accuracy"], 0.5) + + def test_reduce_metrics_empty(self): + """Test that reduce_metrics handles empty lists.""" + metrics = { + "empty": [], + } + result = reduce_metrics(metrics) + + self.assertTrue(np.isnan(result["empty"])) + + def test_reduce_metrics_single_value(self): + """Test that reduce_metrics works with single values.""" + metrics = { + "single": [5.0], + } + result = reduce_metrics(metrics) + + self.assertEqual(result["single"], 5.0) + + +class TestMetric(unittest.TestCase): + """Tests for the Metric class.""" + + def test_init_with_string_aggregation(self): + """Test Metric initialization with string aggregation type.""" + metric = Metric(aggregation="mean") + self.assertEqual(metric.aggregation, AggregationType.MEAN) + self.assertEqual(metric.values, []) + + def test_init_with_enum_aggregation(self): + """Test Metric initialization with AggregationType enum.""" + metric = Metric(aggregation=AggregationType.SUM) + self.assertEqual(metric.aggregation, AggregationType.SUM) + self.assertEqual(metric.values, []) + + def test_init_with_value(self): + """Test Metric initialization with an initial value.""" + metric = Metric(aggregation="mean", value=5.0) + self.assertEqual(metric.values, [5.0]) + + def test_init_with_invalid_aggregation(self): + """Test Metric initialization with invalid aggregation type.""" + with self.assertRaises(ValueError): + Metric(aggregation="invalid") + + def test_append_float(self): + """Test appending float values.""" + metric = Metric(aggregation="mean") + metric.append(1.0) + metric.append(2.0) + self.assertEqual(metric.values, [1.0, 2.0]) + + def test_append_int(self): + """Test appending int values.""" + metric = Metric(aggregation="mean") + metric.append(1) + metric.append(2) + self.assertEqual(metric.values, [1, 2]) + + def test_append_tensor(self): + """Test appending scalar tensor values.""" + metric = Metric(aggregation="mean") + metric.append(torch.tensor(3.0)) + metric.append(torch.tensor(4.0)) + self.assertEqual(metric.values, [3.0, 4.0]) + + def test_append_non_scalar_tensor_raises(self): + """Test that appending non-scalar tensor raises ValueError.""" + metric = Metric(aggregation="mean") + with self.assertRaises(ValueError): + metric.append(torch.tensor([1.0, 2.0])) + + def test_append_metric(self): + """Test appending another Metric extends values.""" + metric1 = Metric(aggregation="mean", value=1.0) + metric1.append(2.0) + + metric2 = Metric(aggregation="mean", value=3.0) + metric2.append(metric1) + + self.assertEqual(metric2.values, [3.0, 1.0, 2.0]) + + def test_extend_with_list(self): + """Test extending with a list of values.""" + metric = Metric(aggregation="mean") + metric.extend([1.0, 2.0, 3.0]) + self.assertEqual(metric.values, [1.0, 2.0, 3.0]) + + def test_extend_with_metric(self): + """Test extending with another Metric.""" + metric1 = Metric(aggregation="mean") + metric1.extend([1.0, 2.0]) + + metric2 = Metric(aggregation="mean") + metric2.extend([3.0, 4.0]) + metric2.extend(metric1) + + self.assertEqual(metric2.values, [3.0, 4.0, 1.0, 2.0]) + + def test_extend_aggregation_mismatch_raises(self): + """Test that extending with mismatched aggregation raises ValueError.""" + metric1 = Metric(aggregation="mean") + metric2 = Metric(aggregation="sum") + + with self.assertRaises(ValueError): + metric1.extend(metric2) + + def test_aggregate_mean(self): + """Test aggregation with mean.""" + metric = Metric(aggregation="mean") + metric.extend([1.0, 2.0, 3.0, 4.0]) + self.assertEqual(metric.aggregate(), 2.5) + + def test_aggregate_sum(self): + """Test aggregation with sum.""" + metric = Metric(aggregation="sum") + metric.extend([1.0, 2.0, 3.0, 4.0]) + self.assertEqual(metric.aggregate(), 10.0) + + def test_aggregate_min(self): + """Test aggregation with min.""" + metric = Metric(aggregation="min") + metric.extend([3.0, 1.0, 4.0, 2.0]) + self.assertEqual(metric.aggregate(), 1.0) + + def test_aggregate_max(self): + """Test aggregation with max.""" + metric = Metric(aggregation="max") + metric.extend([3.0, 1.0, 4.0, 2.0]) + self.assertEqual(metric.aggregate(), 4.0) + + def test_aggregate_dp_sum_mean(self): + """Test aggregate_dp with SUM and MEAN aggregations.""" + # Test with SUM: mean over DP ranks, then sum + metric1 = Metric(aggregation="sum") + metric1.extend([1.0, 2.0]) + + metric2 = Metric(aggregation="sum") + metric2.extend([3.0, 4.0]) + + result = Metric.aggregate_dp([metric1, metric2]) + + # value_arrays = [[1.0, 2.0], [3.0, 4.0]] + # mean over axis 0 = [2.0, 3.0] + # sum = 5.0 + self.assertEqual(result, 5.0) + + # Test with MEAN: mean over DP ranks, then mean + metric4 = Metric(aggregation="mean") + metric4.extend([1.0, 2.0]) + + metric5 = Metric(aggregation="mean") + metric5.extend([3.0, 4.0]) + + result = Metric.aggregate_dp([metric4, metric5]) + + # value_arrays = [[1.0, 2.0], [3.0, 4.0]] + # mean over axis 0 = [2.0, 3.0] + # mean = 2.5 + self.assertEqual(result, 2.5) + + def test_aggregate_dp_min_max(self): + """Test aggregate_dp with MIN and MAX aggregations.""" + # Test with MAX: flatten, then max + metric1 = Metric(aggregation="max") + metric1.extend([1.0, 2.0]) + + metric2 = Metric(aggregation="max") + metric2.extend([3.0, 4.0]) + + result = Metric.aggregate_dp([metric1, metric2]) + + # value_arrays = [[1.0, 2.0], [3.0, 4.0]] + # flatten = [1.0, 2.0, 3.0, 4.0] + # max = 4.0 + self.assertEqual(result, 4.0) + + # Test with MIN: flatten, then min + metric4 = Metric(aggregation="min") + metric4.extend([1.0, 2.0]) + + metric5 = Metric(aggregation="min") + metric5.extend([3.0, 4.0]) + + result = Metric.aggregate_dp([metric4, metric5]) + + # value_arrays = [[1.0, 2.0], [3.0, 4.0]] + # flatten = [1.0, 2.0, 3.0, 4.0] + # min = 1.0 + self.assertEqual(result, 1.0) + + def test_aggregate_dp_mismatched_lengths(self): + """Test aggregate_dp raises error with mismatched value lengths.""" + metric1 = Metric(aggregation="sum") + metric1.extend([1.0, 2.0]) + + metric2 = Metric(aggregation="sum") + metric2.extend([3.0, 4.0, 5.0]) # Different length + + with self.assertRaises(ValueError): + Metric.aggregate_dp([metric1, metric2]) + + def test_from_dict(self): + """Test from_dict creates Metrics from dictionary.""" + data = {"loss": 1.0, "accuracy": 0.9} + metrics = Metric.from_dict(data, aggregation="mean") + + self.assertIn("loss", metrics) + self.assertIn("accuracy", metrics) + self.assertEqual(metrics["loss"].values, [1.0]) + self.assertEqual(metrics["accuracy"].values, [0.9]) + self.assertEqual(metrics["loss"].aggregation, AggregationType.MEAN) + + def test_init_list(self): + """Test init_list creates new empty Metric with same aggregation.""" + metric = Metric(aggregation="max") + metric.extend([1.0, 2.0]) + + new_metric = metric.init_list() + + self.assertEqual(new_metric.aggregation, AggregationType.MAX) + self.assertEqual(new_metric.values, []) + + def test_reduce_metrics_with_metric(self): + """Test reduce_metrics correctly handles Metric objects.""" + metric = Metric(aggregation="mean") + metric.extend([1.0, 2.0, 3.0]) + + metrics = { + "custom_metric": metric, + "list_metric": [4.0, 5.0, 6.0], + } + result = reduce_metrics(metrics) + + self.assertEqual(result["custom_metric"], 2.0) + self.assertEqual(result["list_metric"], 5.0) + + +class TestComputeDataMetrics(unittest.TestCase): + """Tests for the compute_data_metrics function.""" + + def setUp(self): + """Set up common test data.""" + # Create a mock DataProto object + self.batch = MagicMock() + self.batch.batch = { + "token_level_scores": torch.tensor([[1.0, 2.0], [3.0, 4.0]]), + "token_level_rewards": torch.tensor([[0.5, 1.0], [1.5, 2.0]]), + "advantages": torch.tensor([[0.1, 0.2], [0.3, 0.4]]), + "returns": torch.tensor([[1.1, 1.2], [1.3, 1.4]]), + "prompts": torch.zeros((2, 2)), # 2 samples, 2 tokens for each prompt + "responses": torch.zeros((2, 2)), # 2 samples, 2 tokens for each response + "attention_mask": torch.tensor( + [ + [1, 1, 1, 1], # 2 prompt tokens, 2 response tokens + [1, 1, 1, 1], + ] + ), + "response_mask": torch.tensor( + [ + [1, 1], # 2 response tokens + [1, 1], + ] + ), + "values": torch.tensor([[0.9, 1.0], [1.1, 1.2]]), + } + + def test_compute_data_metrics_with_critic(self): + """Test compute_data_metrics with critic enabled.""" + metrics = compute_data_metrics(self.batch, use_critic=True) + + # Check that all expected metrics are present + self.assertIn("critic/score/mean", metrics) + self.assertIn("critic/rewards/mean", metrics) + self.assertIn("critic/advantages/mean", metrics) + self.assertIn("critic/returns/mean", metrics) + self.assertIn("critic/values/mean", metrics) + self.assertIn("critic/vf_explained_var", metrics) + self.assertIn("response_length/mean", metrics) + self.assertIn("prompt_length/mean", metrics) + + # Check some specific values + self.assertAlmostEqual(metrics["critic/score/mean"], 5.0) # Sum of token_level_scores + self.assertAlmostEqual(metrics["critic/rewards/mean"], 2.5) # Sum of token_level_rewards + + def test_compute_data_metrics_without_critic(self): + """Test compute_data_metrics with critic disabled.""" + metrics = compute_data_metrics(self.batch, use_critic=False) + + # Check that critic-specific metrics are not present + self.assertNotIn("critic/values/mean", metrics) + self.assertNotIn("critic/vf_explained_var", metrics) + + # Check that other metrics are still present + self.assertIn("critic/score/mean", metrics) + self.assertIn("critic/rewards/mean", metrics) + self.assertIn("response_length/mean", metrics) + + +class TestComputeTimingMetrics(unittest.TestCase): + """Tests for the compute_timing_metrics function.""" + + def setUp(self): + """Set up common test data.""" + # Create a mock DataProto object + self.batch = MagicMock() + self.batch.batch = { + "responses": torch.zeros((2, 3)), # 2 samples, 3 response tokens each + "attention_mask": torch.tensor( + [ + [1, 1, 1, 1, 1, 1], # 3 prompt tokens, 3 response tokens + [1, 1, 1, 1, 1, 1], + ] + ), + } + + # Mock the _compute_response_info function to return known values + self.response_info = { + "prompt_length": torch.tensor([3.0, 3.0]), + "response_length": torch.tensor([3.0, 3.0]), + "response_mask": torch.ones((2, 3)), + } + + @patch("verl.trainer.ppo.metric_utils._compute_response_info") + def test_compute_timing_metrics(self, mock_compute_response_info): + """Test compute_timing_metrics with various timing data.""" + mock_compute_response_info.return_value = self.response_info + + timing_raw = { + "gen": 0.5, # 500ms + "ref": 0.3, # 300ms + "values": 0.2, # 200ms + } + + metrics = compute_timing_metrics(self.batch, timing_raw) + + # Check raw timing metrics + self.assertEqual(metrics["timing_s/gen"], 0.5) + self.assertEqual(metrics["timing_s/ref"], 0.3) + self.assertEqual(metrics["timing_s/values"], 0.2) + + # Check per-token timing metrics + # gen uses only response tokens (6 tokens) + self.assertAlmostEqual(metrics["timing_per_token_ms/gen"], 0.5 * 1000 / 6, places=5) + + # ref and values use all tokens (12 tokens) + self.assertAlmostEqual(metrics["timing_per_token_ms/ref"], 0.3 * 1000 / 12, places=5) + self.assertAlmostEqual(metrics["timing_per_token_ms/values"], 0.2 * 1000 / 12, places=5) + + +class TestComputeThroughputMetrics(unittest.TestCase): + """Tests for the compute_throughout_metrics function.""" + + def setUp(self): + """Set up common test data.""" + # Create a mock DataProto object + self.batch = MagicMock() + self.batch.meta_info = { + "global_token_num": [100, 200, 300], # 600 tokens total + } + + def test_compute_throughout_metrics(self): + """Test compute_throughout_metrics with various timing data.""" + timing_raw = { + "step": 2.0, # 2 seconds per step + } + + # Test with 1 GPU + metrics = compute_throughout_metrics(self.batch, timing_raw, n_gpus=1) + + self.assertEqual(metrics["perf/total_num_tokens"], 600) + self.assertEqual(metrics["perf/time_per_step"], 2.0) + self.assertEqual(metrics["perf/throughput"], 600 / 2.0) # 300 tokens/sec + + # Test with 2 GPUs + metrics = compute_throughout_metrics(self.batch, timing_raw, n_gpus=2) + + self.assertEqual(metrics["perf/total_num_tokens"], 600) + self.assertEqual(metrics["perf/time_per_step"], 2.0) + self.assertEqual(metrics["perf/throughput"], 600 / (2.0 * 2)) # 150 tokens/sec/GPU + + +class TestBootstrapMetric(unittest.TestCase): + """Tests for the bootstrap_metric function.""" + + def test_bootstrap_metric_basic(self): + """Test bootstrap_metric with simple data and functions.""" + data = [1, 2, 3, 4, 5] + reduce_fns = [np.mean, np.max] + + # Use a fixed seed for reproducibility + result = bootstrap_metric(data, subset_size=3, reduce_fns=reduce_fns, n_bootstrap=100, seed=42) + + # Check that we get two results (one for each reduce_fn) + self.assertEqual(len(result), 2) + + # Each result should be a tuple of (mean, std) + mean_result, max_result = result + self.assertEqual(len(mean_result), 2) + self.assertEqual(len(max_result), 2) + + # The mean of means should be close to the true mean (3.0) + self.assertAlmostEqual(mean_result[0], 3.0, delta=0.3) + + # The mean of maxes should be close to the expected value for samples of size 3 + # For samples of size 3 from [1,2,3,4,5], the expected max is around 4.0-4.5 + self.assertGreater(max_result[0], 3.5) + self.assertLess(max_result[0], 5.0) + + def test_bootstrap_metric_empty(self): + """Test bootstrap_metric with empty data.""" + with self.assertRaises(ValueError): + bootstrap_metric([], subset_size=1, reduce_fns=[np.mean]) + + +class TestCalcMajVal(unittest.TestCase): + """Tests for the calc_maj_val function.""" + + def test_calc_maj_val_basic(self): + """Test calc_maj_val with simple data.""" + data = [ + {"pred": "A", "val": 0.9}, + {"pred": "B", "val": 0.8}, + {"pred": "A", "val": 0.7}, + ] + + result = calc_maj_val(data, vote_key="pred", val_key="val") + + # "A" is the majority vote, so we should get the first "val" for "A" + self.assertEqual(result, 0.9) + + def test_calc_maj_val_tie(self): + """Test calc_maj_val with tied votes.""" + data = [ + {"pred": "A", "val": 0.9}, + {"pred": "B", "val": 0.8}, + {"pred": "B", "val": 0.7}, + {"pred": "A", "val": 0.6}, + ] + + # In case of a tie, the first key in sorted order wins + # This depends on Python's dict implementation, but for this test + # we just verify that one of the valid values is returned + result = calc_maj_val(data, vote_key="pred", val_key="val") + + self.assertTrue(result in [0.9, 0.8]) + + +class TestProcessValidationMetrics(unittest.TestCase): + """Tests for the process_validation_metrics function.""" + + def test_process_validation_metrics_basic(self): + """Test process_validation_metrics with simple data.""" + data_sources = ["source1", "source1", "source2"] + sample_inputs = ["prompt1", "prompt1", "prompt2"] + infos_dict = { + "score": [0.8, 0.9, 0.7], + } + + result = process_validation_metrics(data_sources, sample_inputs, infos_dict, seed=42) + + # Check the structure of the result + self.assertIn("source1", result) + self.assertIn("source2", result) + + # Check that source1 has metrics for score + self.assertIn("score", result["source1"]) + + # Check that mean@2 is present for source1/score + self.assertIn("mean@2", result["source1"]["score"]) + + # Check the value of mean@2 for source1/score + self.assertAlmostEqual(result["source1"]["score"]["mean@2"], 0.85) + + def test_process_validation_metrics_with_pred(self): + """Test process_validation_metrics with prediction data.""" + data_sources = ["source1", "source1", "source1"] + sample_inputs = ["prompt1", "prompt1", "prompt1"] + infos_dict = { + "score": [0.8, 0.9, 0.7], + "pred": ["A", "B", "A"], + } + + result = process_validation_metrics(data_sources, sample_inputs, infos_dict, seed=42) + + # Check that majority voting metrics are present + self.assertIn("maj@2/mean", result["source1"]["score"]) + + # For bootstrap with n=2, the majority vote could be either A or B + # depending on the random sampling, so we don't check the exact value + + +if __name__ == "__main__": + unittest.main() diff --git a/verl/tests/trainer/ppo/test_rollout_corr.py b/verl/tests/trainer/ppo/test_rollout_corr.py new file mode 100644 index 0000000000000000000000000000000000000000..a081d8b73d3d98e431df3086682502cc9b046327 --- /dev/null +++ b/verl/tests/trainer/ppo/test_rollout_corr.py @@ -0,0 +1,433 @@ +#!/usr/bin/env python3 +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Quick Sanity Test for Rollout Correction + +This is a standalone test script that can be run without pytest to quickly verify +the rollout correction implementation is working correctly. For comprehensive integration +tests, see: tests/trainer/ppo/test_rollout_corr_integration.py + +Usage: + python test_rollout_corr.py + +This tests: +- Basic rollout correction functionality (IS weights + rejection sampling) +- Metrics completeness (IS metrics + rejection metrics + off-policy metrics) +- Edge cases +""" + +import pytest +import torch + +from verl.trainer.ppo.rollout_corr_helper import ( + SUPPORTED_ROLLOUT_RS_OPTIONS, + compute_offpolicy_metrics, + compute_rollout_correction_and_rejection_mask, +) + + +def test_basic_rollout_correction(): + """Test basic rollout correction functionality.""" + print("Testing basic rollout correction functionality...") + + # Create test data + batch_size, seq_length = 4, 10 + device = "cuda" if torch.cuda.is_available() else "cpu" + + # Create slightly different log probs (simulating BF16 vs FP32 mismatch) + old_log_prob = torch.randn(batch_size, seq_length, device=device) + rollout_log_prob = old_log_prob + torch.randn(batch_size, seq_length, device=device) * 0.1 + eos_mask = torch.ones(batch_size, seq_length, device=device) + + # Test token-level truncate mode + print("\n1. Testing token-level truncate mode...") + weights_proto, modified_response_mask, metrics = compute_rollout_correction_and_rejection_mask( + old_log_prob=old_log_prob, + rollout_log_prob=rollout_log_prob, + response_mask=eos_mask, + rollout_is="token", # Compute IS weights at token level + rollout_is_threshold=2.0, + rollout_rs=None, # No rejection sampling (truncate mode) + ) + + weights = weights_proto.batch["rollout_is_weights"] + print(f" Weights shape: {weights.shape}") + print(f" Mean weight: {metrics['rollout_corr/rollout_is_mean']:.4f}") + print(f" Max weight: {metrics['rollout_corr/rollout_is_max']:.4f}") + print(f" Min weight: {metrics['rollout_corr/rollout_is_min']:.4f}") + assert weights.shape == old_log_prob.shape + assert weights.max() <= 2.0, "Weights should be capped at threshold" + print(" ✓ Token-level truncate mode passed") + + # Test sequence-level mode + print("\n2. Testing sequence-level mode...") + weights_seq_proto, _, metrics_seq = compute_rollout_correction_and_rejection_mask( + old_log_prob=old_log_prob, + rollout_log_prob=rollout_log_prob, + response_mask=eos_mask, + rollout_is="sequence", # Compute IS weights at sequence level + rollout_is_threshold=5.0, + rollout_rs=None, # No rejection sampling (truncate mode) + ) + + weights_seq = weights_seq_proto.batch["rollout_is_weights"] + print(f" Mean weight: {metrics_seq['rollout_corr/rollout_is_mean']:.4f}") + print(f" Effective sample size: {metrics_seq['rollout_corr/rollout_is_eff_sample_size']:.4f}") + # Check that all tokens in a sequence have the same weight + for i in range(batch_size): + seq_weights = weights_seq[i, eos_mask[i].bool()] + assert torch.allclose(seq_weights, seq_weights[0]), "All tokens in sequence should have same weight" + print(" ✓ Sequence-level mode passed") + + # Test K1 sequence mean rejection sampling (mask mode) + print("\n3. Testing K1 (sequence mean) rejection sampling...") + weights_geo_proto, modified_mask_geo, metrics_geo = compute_rollout_correction_and_rejection_mask( + old_log_prob=old_log_prob, + rollout_log_prob=rollout_log_prob, + response_mask=eos_mask, + rollout_is=None, # No IS weights (pure mask mode) + rollout_rs="seq_mean_k1", # Rejection sampling with sequence-mean log ratio bounds + rollout_rs_threshold="0.5_1.5", + ) + + print(f" Masked fraction: {metrics_geo['rollout_corr/rollout_rs_masked_fraction']:.4f}") + print(" ✓ K1 sequence mean rejection sampling passed") + + # Test disabled IS (rollout_is=None, rollout_rs=None) + print("\n4. Testing disabled IS...") + weights_disabled, modified_response_mask_disabled, metrics_disabled = compute_rollout_correction_and_rejection_mask( + old_log_prob=old_log_prob, + rollout_log_prob=rollout_log_prob, + response_mask=eos_mask, + rollout_is=None, + rollout_rs=None, + ) + + assert weights_disabled is None, "Should return None when IS is disabled" + assert torch.equal(modified_response_mask_disabled, eos_mask), "Should return original mask unchanged" + # Note: off-policy metrics are still computed even when IS/RS are disabled + assert "rollout_corr/kl" in metrics_disabled, "Should still compute off-policy metrics" + print(" ✓ Disabled IS passed") + + print("\n✓ All tests passed!") + + +@pytest.mark.parametrize( + ("option", "threshold"), + [ + ("token_k1", "0.5_1.5"), + ("token_k2", 2.0), + ("token_k3", 2.0), + ("seq_sum_k1", "0.6_1.4"), + ("seq_sum_k2", 2.5), + ("seq_sum_k3", 2.5), + ("seq_mean_k1", "0.5_1.5"), + ("seq_mean_k2", 2.0), + ("seq_mean_k3", 2.0), + ("seq_max_k2", 2.0), + ("seq_max_k3", 2.0), + ], +) +def test_each_supported_rollout_rs_option(option: str, threshold): + """Ensure every supported RS option produces metrics without error.""" + assert option in SUPPORTED_ROLLOUT_RS_OPTIONS + + batch_size, seq_length = 3, 7 + device = "cuda" if torch.cuda.is_available() else "cpu" + + old_log_prob = torch.randn(batch_size, seq_length, device=device) + rollout_log_prob = old_log_prob + torch.randn(batch_size, seq_length, device=device) * 0.15 + response_mask = torch.ones(batch_size, seq_length, device=device) + + _, modified_mask, metrics = compute_rollout_correction_and_rejection_mask( + old_log_prob=old_log_prob, + rollout_log_prob=rollout_log_prob, + response_mask=response_mask, + rollout_is=None, + rollout_rs=option, + rollout_rs_threshold=threshold, + ) + + expected_key = f"rollout_corr/rollout_rs_{option}_mean" + assert expected_key in metrics, f"Missing metric for {option}" + assert modified_mask.shape == response_mask.shape + + +def test_rollout_rs_multiple_options(): + """Verify multiple RS options with mixed threshold formats.""" + batch_size, seq_length = 2, 6 + device = "cuda" if torch.cuda.is_available() else "cpu" + + old_log_prob = torch.randn(batch_size, seq_length, device=device) + rollout_log_prob = old_log_prob + torch.randn(batch_size, seq_length, device=device) * 0.2 + response_mask = torch.ones(batch_size, seq_length, device=device) + + rollout_rs = "token_k1,seq_max_k3" + rollout_rs_threshold = "0.4_1.8,3.0" + + _, _, metrics = compute_rollout_correction_and_rejection_mask( + old_log_prob=old_log_prob, + rollout_log_prob=rollout_log_prob, + response_mask=response_mask, + rollout_is=None, + rollout_rs=rollout_rs, + rollout_rs_threshold=rollout_rs_threshold, + ) + + for option in rollout_rs.split(","): + key = f"rollout_corr/rollout_rs_{option}_mean" + assert key in metrics, f"Metrics missing for chained option {option}" + + +def test_metrics_completeness(): + """Test that all expected metrics are returned.""" + print("\nTesting metrics completeness...") + + batch_size, seq_length = 3, 8 + device = "cuda" if torch.cuda.is_available() else "cpu" + + old_log_prob = torch.randn(batch_size, seq_length, device=device) + rollout_log_prob = old_log_prob + torch.randn(batch_size, seq_length, device=device) * 0.2 + eos_mask = torch.ones(batch_size, seq_length, device=device) + + _, _, metrics = compute_rollout_correction_and_rejection_mask( + old_log_prob=old_log_prob, + rollout_log_prob=rollout_log_prob, + response_mask=eos_mask, + rollout_is="token", + rollout_is_threshold=2.5, + rollout_rs=None, + ) + + # Expected IS metrics + expected_is_metrics = [ + "rollout_corr/rollout_is_mean", + "rollout_corr/rollout_is_max", + "rollout_corr/rollout_is_min", + "rollout_corr/rollout_is_std", + "rollout_corr/rollout_is_eff_sample_size", + "rollout_corr/rollout_is_ratio_fraction_high", + "rollout_corr/rollout_is_ratio_fraction_low", + ] + + # Expected off-policy diagnostic metrics (also included now) + expected_offpolicy_metrics = [ + "rollout_corr/training_ppl", + "rollout_corr/training_log_ppl", + "rollout_corr/kl", + "rollout_corr/k3_kl", + "rollout_corr/rollout_ppl", + "rollout_corr/rollout_log_ppl", + "rollout_corr/log_ppl_diff", + "rollout_corr/log_ppl_abs_diff", + "rollout_corr/log_ppl_diff_max", + "rollout_corr/log_ppl_diff_min", + "rollout_corr/ppl_ratio", + "rollout_corr/chi2_token", + "rollout_corr/chi2_seq", + ] + + expected_metrics = expected_is_metrics + expected_offpolicy_metrics + + missing_metrics = [m for m in expected_metrics if m not in metrics] + if missing_metrics: + print(f" ✗ Missing metrics: {missing_metrics}") + return False + + print(f" ✓ All {len(expected_metrics)} expected metrics present") + print(f" Total metrics returned: {len(metrics)}") + return True + + +def test_offpolicy_metrics(): + """Test off-policy metrics computation.""" + print("\nTesting off-policy metrics computation...") + + batch_size, seq_length = 4, 12 + device = "cuda" if torch.cuda.is_available() else "cpu" + + # Create test data with some mismatch + old_log_prob = torch.randn(batch_size, seq_length, device=device) - 2.0 # training policy + rollout_log_prob = torch.randn(batch_size, seq_length, device=device) - 1.5 # rollout policy (more confident) + response_mask = torch.ones(batch_size, seq_length, device=device) + + # Test with rollout log probs + metrics = compute_offpolicy_metrics( + old_log_prob=old_log_prob, + rollout_log_prob=rollout_log_prob, + response_mask=response_mask, + ) + + expected_metrics = [ + "training_ppl", + "training_log_ppl", + "kl", + "k3_kl", + "rollout_ppl", + "rollout_log_ppl", + "log_ppl_diff", + "log_ppl_abs_diff", + "log_ppl_diff_max", + "log_ppl_diff_min", + "ppl_ratio", + "chi2_token", + "chi2_seq", + ] + + for metric in expected_metrics: + assert metric in metrics, f"Missing metric: {metric}" + + print(f" Training PPL: {metrics['training_ppl']:.4f}") + print(f" Rollout PPL: {metrics['rollout_ppl']:.4f}") + print(f" KL divergence: {metrics['kl']:.6f}") + print(f" K3 KL: {metrics['k3_kl']:.6f}") + print(f" PPL ratio: {metrics['ppl_ratio']:.4f}") + print(f" ✓ All {len(expected_metrics)} off-policy metrics present") + + # Test without rollout log probs + metrics_no_rollout = compute_offpolicy_metrics( + old_log_prob=old_log_prob, + rollout_log_prob=None, + response_mask=response_mask, + ) + + assert "training_ppl" in metrics_no_rollout + assert "rollout_ppl" not in metrics_no_rollout + print(" ✓ Off-policy metrics work without rollout log probs") + + +def test_mask_mode(): + """Test mask mode applies rejection via response_mask, keeps true IS weights.""" + print("\nTesting mask mode behavior...") + + batch_size = 2 + seq_length = 5 + device = "cuda" if torch.cuda.is_available() else "cpu" + + # Sequence 0: ratio ≈ 0.37 (below 0.5, should be rejected) + # Sequence 1: ratio ≈ 1.65 (in [0.5, 2.0], should be accepted) + old_log_prob = torch.tensor([[-2.0] * seq_length, [-2.0] * seq_length], device=device) + rollout_log_prob = torch.tensor( + [ + [-1.0] * seq_length, # exp(-2.0 - (-1.0)) = exp(-1.0) ≈ 0.37 + [-2.5] * seq_length, # exp(-2.0 - (-2.5)) = exp(0.5) ≈ 1.65 + ], + device=device, + ) + response_mask = torch.ones(batch_size, seq_length, device=device) + + weights_proto, modified_response_mask, metrics = compute_rollout_correction_and_rejection_mask( + old_log_prob=old_log_prob, + rollout_log_prob=rollout_log_prob, + response_mask=response_mask, + rollout_is="token", # Compute IS weights + rollout_is_threshold=2.0, + rollout_rs="token_k1", # Also apply rejection sampling (mask mode) + rollout_rs_threshold="0.5_2.0", + ) + + weights = weights_proto.batch["rollout_is_weights"] + + # KEY FIX: Weights should be safety-bounded ratios (NOT zeroed) + assert torch.all(weights[0, :] > 0), "Weights should remain as safety-bounded ratios (not zeroed)" + assert torch.allclose(weights[0, 0], torch.tensor(0.368, device=device), atol=0.01), ( + "First seq ratio should be ≈0.37" + ) + assert torch.allclose(weights[1, 0], torch.tensor(1.649, device=device), atol=0.01), ( + "Second seq ratio should be ≈1.65" + ) + + # Rejection should be applied via response_mask + assert torch.all(modified_response_mask[0, :] == 0), "First sequence should be rejected via mask" + assert torch.all(modified_response_mask[1, :] == 1), "Second sequence should be accepted" + + # Verify rejection sampling metrics exist + assert "rollout_corr/rollout_rs_masked_fraction" in metrics, "Should have rollout_rs_masked_fraction metric" + assert abs(metrics["rollout_corr/rollout_rs_masked_fraction"] - 0.5) < 0.01, "Should reject 50% of tokens" + + print(f" First seq IS weight: {weights[0, 0]:.4f} (expected ≈0.37)") + print(f" Second seq IS weight: {weights[1, 0]:.4f} (expected ≈1.65)") + print(f" First seq mask: {modified_response_mask[0, 0]:.0f} (expected 0 - rejected)") + print(f" Second seq mask: {modified_response_mask[1, 0]:.0f} (expected 1 - accepted)") + print(f" Masked fraction: {metrics['rollout_corr/rollout_rs_masked_fraction']:.2f}") + print(" ✓ Mask mode correctly separates IS weights from rejection") + + +def test_exact_icepop_zeroes_weights_without_changing_mask(): + """IcePop should zero OOB IS weights while preserving response_mask.""" + device = "cuda" if torch.cuda.is_available() else "cpu" + + raw_is_weights = torch.tensor([[0.4, 0.8, 6.0]], device=device) + old_log_prob = torch.zeros_like(raw_is_weights) + rollout_log_prob = -torch.log(raw_is_weights) + response_mask = torch.ones_like(raw_is_weights) + + weights_proto, modified_response_mask, metrics = compute_rollout_correction_and_rejection_mask( + old_log_prob=old_log_prob, + rollout_log_prob=rollout_log_prob, + response_mask=response_mask, + rollout_is="token", + rollout_is_threshold="0.5_5.0", + rollout_rs=None, + ) + + weights = weights_proto.batch["rollout_is_weights"] + expected_weights = torch.tensor([[0.0, 0.8, 0.0]], device=device) + + torch.testing.assert_close(weights, expected_weights, atol=1e-6, rtol=1e-6) + assert torch.equal(modified_response_mask, response_mask) + assert metrics["rollout_corr/rollout_is_oob_ratio"] == pytest.approx(2.0 / 3.0, abs=1e-6) + assert metrics["rollout_corr/rollout_is_std"] == pytest.approx(0.3771236, abs=1e-6) + assert metrics["rollout_corr/rollout_is_eff_sample_size"] == pytest.approx(1.0 / 3.0, abs=1e-6) + + +def test_bool_rollout_is_threshold_is_rejected(): + """Boolean thresholds should not be silently accepted via bool <: int.""" + device = "cuda" if torch.cuda.is_available() else "cpu" + + old_log_prob = torch.zeros(1, 2, device=device) + rollout_log_prob = torch.zeros(1, 2, device=device) + response_mask = torch.ones(1, 2, device=device) + + with pytest.raises(TypeError, match="not a boolean"): + compute_rollout_correction_and_rejection_mask( + old_log_prob=old_log_prob, + rollout_log_prob=rollout_log_prob, + response_mask=response_mask, + rollout_is="token", + rollout_is_threshold=True, + rollout_rs=None, + ) + + +if __name__ == "__main__": + print("=" * 60) + print("Rollout Correction Test Suite") + print("=" * 60) + + try: + test_basic_rollout_correction() + test_metrics_completeness() + test_offpolicy_metrics() + test_mask_mode() + print("\n" + "=" * 60) + print("ALL TESTS PASSED ✓") + print("=" * 60) + except Exception as e: + print(f"\n✗ Test failed with error: {e}") + import traceback + + traceback.print_exc() + exit(1) diff --git a/verl/tests/trainer/ppo/test_rollout_corr_integration.py b/verl/tests/trainer/ppo/test_rollout_corr_integration.py new file mode 100644 index 0000000000000000000000000000000000000000..dfdc022b221835ce41addda4cfaeef5b07499d69 --- /dev/null +++ b/verl/tests/trainer/ppo/test_rollout_corr_integration.py @@ -0,0 +1,305 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Integration tests for Rollout Correction.""" + +import pytest +import torch + +from verl.trainer.config.algorithm import RolloutCorrectionConfig +from verl.trainer.ppo.core_algos import compute_policy_loss_vanilla +from verl.trainer.ppo.rollout_corr_helper import ( + compute_offpolicy_metrics, + compute_rollout_correction_and_rejection_mask, +) +from verl.workers.config.actor import ActorConfig + + +class TestRolloutISIntegration: + """Integration tests for Rollout Correction with PPO.""" + + @pytest.fixture + def sample_data(self): + """Create sample training data.""" + batch_size, seq_length = 4, 16 + device = "cuda" if torch.cuda.is_available() else "cpu" + + return { + "old_log_prob": torch.randn(batch_size, seq_length, device=device), + "log_prob": torch.randn(batch_size, seq_length, device=device), + "rollout_log_prob": torch.randn(batch_size, seq_length, device=device), + "advantages": torch.randn(batch_size, seq_length, device=device), + "response_mask": torch.ones(batch_size, seq_length, device=device), + } + + @pytest.fixture + def config_with_rollout_is(self): + """Create config for policy loss computation. + + Note: rollout_is config has been moved to algorithm config. + This config only needs fields used by policy loss (clip_ratio, etc). + """ + config = ActorConfig( + strategy="fsdp", + rollout_n=1, + ppo_micro_batch_size=2, + clip_ratio=0.2, + ) + return config + + def test_policy_loss_with_rollout_is(self, sample_data, config_with_rollout_is): + """Test that policy loss computation works with rollout correction weights. + + Note: In production, IS weights are computed centrally in the trainer + (before advantage computation) and passed to policy loss. + This test simulates that workflow. + """ + # First compute IS weights (as trainer would do centrally) + rollout_is_weights_proto, _, _ = compute_rollout_correction_and_rejection_mask( + old_log_prob=sample_data["old_log_prob"], + rollout_log_prob=sample_data["rollout_log_prob"], + response_mask=sample_data["response_mask"], + rollout_is="token", + rollout_is_threshold=2.0, + rollout_rs=None, + ) + + rollout_is_weights = rollout_is_weights_proto.batch["rollout_is_weights"] + + # Policy loss function receives pre-computed IS weights + pg_loss, _ = compute_policy_loss_vanilla( + old_log_prob=sample_data["old_log_prob"], + log_prob=sample_data["log_prob"], + advantages=sample_data["advantages"], + response_mask=sample_data["response_mask"], + loss_agg_mode="token-mean", + config=config_with_rollout_is, + rollout_is_weights=rollout_is_weights, + ) + + # Check loss is valid + assert isinstance(pg_loss, torch.Tensor) + assert pg_loss.ndim == 0 # Scalar + assert not torch.isnan(pg_loss) + assert not torch.isinf(pg_loss) + + def test_rollout_is_weights_computation(self, sample_data): + """Test rollout correction weights and metrics computation.""" + weights_proto, _, metrics = compute_rollout_correction_and_rejection_mask( + old_log_prob=sample_data["old_log_prob"], + rollout_log_prob=sample_data["rollout_log_prob"], + response_mask=sample_data["response_mask"], + rollout_is="token", + rollout_is_threshold=2.0, + rollout_rs=None, + ) + + # Check weights + from verl.protocol import DataProto + + assert isinstance(weights_proto, DataProto) + weights = weights_proto.batch["rollout_is_weights"] + assert isinstance(weights, torch.Tensor) + assert weights.shape == sample_data["old_log_prob"].shape + + # Check metrics are returned + assert isinstance(metrics, dict) + assert len(metrics) > 0 + assert "rollout_corr/rollout_is_mean" in metrics + + def test_all_aggregation_levels(self, sample_data): + """Test all aggregation levels (token, sequence for IS; K1 for RS).""" + # Test IS weight levels + is_levels = ["token", "sequence"] + for level in is_levels: + _, _, metrics = compute_rollout_correction_and_rejection_mask( + old_log_prob=sample_data["old_log_prob"], + rollout_log_prob=sample_data["rollout_log_prob"], + response_mask=sample_data["response_mask"], + rollout_is=level, + rollout_is_threshold=2.0, + rollout_rs=None, + ) + assert "rollout_corr/rollout_is_mean" in metrics + + # Test rejection sampling with K1 sequence mean level + _, _, metrics_geo = compute_rollout_correction_and_rejection_mask( + old_log_prob=sample_data["old_log_prob"], + rollout_log_prob=sample_data["rollout_log_prob"], + response_mask=sample_data["response_mask"], + rollout_is=None, + rollout_rs="seq_mean_k1", + rollout_rs_threshold="0.999_1.001", + ) + assert "rollout_corr/rollout_rs_seq_mean_k1_mean" in metrics_geo + + def test_both_bounding_modes(self, sample_data): + """Test both truncate and mask modes.""" + # Test truncate mode (IS weights only) + _, _, metrics_truncate = compute_rollout_correction_and_rejection_mask( + old_log_prob=sample_data["old_log_prob"], + rollout_log_prob=sample_data["rollout_log_prob"], + response_mask=sample_data["response_mask"], + rollout_is="token", + rollout_is_threshold=2.0, + rollout_rs=None, + ) + assert "rollout_corr/rollout_is_mean" in metrics_truncate + + # Test mask mode (rejection sampling) + _, _, metrics_mask = compute_rollout_correction_and_rejection_mask( + old_log_prob=sample_data["old_log_prob"], + rollout_log_prob=sample_data["rollout_log_prob"], + response_mask=sample_data["response_mask"], + rollout_is="token", # Can also compute IS weights in mask mode + rollout_is_threshold=2.0, + rollout_rs="token_k1", # Enable rejection sampling + rollout_rs_threshold=1.3, # Float upper bound (lower inferred automatically) + ) + assert "rollout_corr/rollout_is_mean" in metrics_mask + assert "rollout_corr/rollout_rs_token_k1_mean" in metrics_mask + + def test_offpolicy_metrics(self, sample_data): + """Test off-policy diagnostic metrics computation.""" + metrics = compute_offpolicy_metrics( + old_log_prob=sample_data["old_log_prob"], + rollout_log_prob=sample_data["rollout_log_prob"], + response_mask=sample_data["response_mask"], + ) + + # Check key metrics are present + assert "training_ppl" in metrics + assert "rollout_ppl" in metrics + assert "kl" in metrics + assert isinstance(metrics["kl"], float) + + def test_metrics_only_mode(self, sample_data, config_with_rollout_is): + """Test metrics-only mode: compute IS weights/metrics but don't apply to loss. + + This tests the use case where rollout_is_threshold is set (enables computation) + but rollout_is=False (disables weight application to policy loss). + """ + # Compute IS weights (as trainer would do) + rollout_is_weights_proto, _, is_metrics = compute_rollout_correction_and_rejection_mask( + old_log_prob=sample_data["old_log_prob"], + rollout_log_prob=sample_data["rollout_log_prob"], + response_mask=sample_data["response_mask"], + rollout_is="token", + rollout_is_threshold=2.0, + rollout_rs=None, + ) + + # Metrics should be computed + assert len(is_metrics) > 0 + assert "rollout_corr/rollout_is_mean" in is_metrics + + # In metrics-only mode, we compute loss WITHOUT applying weights + # (simulating rollout_is=False) + pg_loss_no_weights, _ = compute_policy_loss_vanilla( + old_log_prob=sample_data["old_log_prob"], + log_prob=sample_data["log_prob"], + advantages=sample_data["advantages"], + response_mask=sample_data["response_mask"], + loss_agg_mode="token-mean", + config=config_with_rollout_is, + rollout_is_weights=None, # Don't apply weights + ) + + # Compare to loss WITH weights (rollout_is=True) + rollout_is_weights = rollout_is_weights_proto.batch["rollout_is_weights"] + pg_loss_with_weights, _ = compute_policy_loss_vanilla( + old_log_prob=sample_data["old_log_prob"], + log_prob=sample_data["log_prob"], + advantages=sample_data["advantages"], + response_mask=sample_data["response_mask"], + loss_agg_mode="token-mean", + config=config_with_rollout_is, + rollout_is_weights=rollout_is_weights, + ) + + # Losses should be different (weights have an effect) + assert not torch.allclose(pg_loss_no_weights, pg_loss_with_weights) + + def test_exact_icepop_matches_filtered_weighted_ppo_loss(self, config_with_rollout_is): + """IcePop should match the local RL zero-weight semantics.""" + device = "cuda" if torch.cuda.is_available() else "cpu" + + old_log_prob = torch.tensor([[-1.0, -1.0, -1.0]], device=device) + log_prob = old_log_prob.clone() + rollout_log_prob = torch.tensor([[-0.5, -3.5, -0.8]], device=device) + advantages = torch.tensor([[1.0, -1.0, 2.0]], device=device) + response_mask = torch.ones_like(old_log_prob) + + rollout_is_weights_proto, modified_response_mask, metrics = compute_rollout_correction_and_rejection_mask( + old_log_prob=old_log_prob, + rollout_log_prob=rollout_log_prob, + response_mask=response_mask, + rollout_is="token", + rollout_is_threshold="0.5_5.0", + rollout_rs=None, + ) + + rollout_is_weights = rollout_is_weights_proto.batch["rollout_is_weights"] + expected_weights = torch.tensor([[0.60653067, 0.0, 0.81873075]], device=device) + expected_loss = torch.mean(expected_weights * (-advantages)) + + pg_loss, _ = compute_policy_loss_vanilla( + old_log_prob=old_log_prob, + log_prob=log_prob, + advantages=advantages, + response_mask=response_mask, + loss_agg_mode="token-mean", + config=config_with_rollout_is, + rollout_is_weights=rollout_is_weights, + ) + + assert torch.equal(modified_response_mask, response_mask) + assert metrics["rollout_corr/rollout_is_oob_ratio"] == pytest.approx(1.0 / 3.0, abs=1e-6) + torch.testing.assert_close(rollout_is_weights, expected_weights, atol=1e-6, rtol=1e-6) + torch.testing.assert_close(pg_loss, expected_loss, atol=1e-6, rtol=1e-6) + + +class TestRolloutCorrectionConfigNormalization: + """Unit tests for RolloutCorrectionConfig canonicalization logic.""" + + def test_alias_normalization_and_threshold_parsing(self): + config = RolloutCorrectionConfig( + rollout_is="token", + rollout_is_threshold=2.5, + rollout_rs="seq_mean_k1,seq_max_k3", + rollout_rs_threshold="0.8_1.2,3.0", + ) + + assert config.rollout_is == "token" + assert config.rollout_is_threshold == pytest.approx(2.5) + assert config.rollout_rs == "seq_mean_k1,seq_max_k3" + assert config.rollout_rs_threshold == "0.8_1.2,3.0" + + def test_missing_threshold_raises(self): + config = RolloutCorrectionConfig(rollout_rs="token_k1") + assert config.rollout_rs == "token_k1" + assert config.rollout_rs_threshold is None + + def test_float_threshold_conversion_in_factory(self): + config = RolloutCorrectionConfig.decoupled_geo_rs_seq_tis(rs_threshold=1.001) + assert config.rollout_rs == "seq_mean_k1" + assert config.rollout_rs_threshold == 1.001 + + def test_icepop_factory(self): + config = RolloutCorrectionConfig.decoupled_token_icepop() + assert config.rollout_is == "token" + assert config.rollout_is_threshold == "0.5_5.0" + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/verl/tests/trainer/test_multi_trajectories_advantage.py b/verl/tests/trainer/test_multi_trajectories_advantage.py new file mode 100644 index 0000000000000000000000000000000000000000..93ca773801e3d8032cb898dc2dbee970aeee940e --- /dev/null +++ b/verl/tests/trainer/test_multi_trajectories_advantage.py @@ -0,0 +1,86 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import numpy as np +import pytest +import torch + +from verl.protocol import DataProto +from verl.trainer.main_ppo_sync import compute_advantage, compute_advantage_for_multi_trajectories +from verl.trainer.ppo.core_algos import AdvantageEstimator + + +@pytest.fixture +def batch_data() -> DataProto: + tensors = { + "token_level_rewards": torch.tensor( + [ + [100, 0, 0, 0], + [1, 2, 3, 4], + [200, 200, 0, 0], + [150, 0, 0, 150], + [4, 5, 6, 7], + [8, 9, 10, 11], + ], + dtype=torch.float32, + ), + "response_mask": torch.tensor( + [ + [1, 0, 0, 0], + [1, 1, 1, 1], + [1, 0, 1, 0], + [0, 1, 0, 1], + [1, 1, 0, 1], + [0, 0, 0, 0], + ], + dtype=torch.long, + ), + } + non_tensors = { + "uid": np.array(["prompt_a"] * 6, dtype=object), + } + return DataProto.from_dict(tensors=tensors, non_tensors=non_tensors) + + +def test_compute_advantage_for_single_trajectory(batch_data: DataProto): + result = compute_advantage_for_multi_trajectories( + data=batch_data, + batch_keys=[f"prompt_a_{i}_0" for i in range(len(batch_data))], + adv_estimator=AdvantageEstimator.GRPO, + ) + expected = compute_advantage( + batch_data, + adv_estimator=AdvantageEstimator.GRPO, + ) + assert torch.equal(result.batch["advantages"], expected.batch["advantages"]) + assert torch.equal(result.batch["returns"], expected.batch["returns"]) + + +def test_compute_advantage_for_multi_trajectories(batch_data: DataProto): + result = compute_advantage_for_multi_trajectories( + data=batch_data, + batch_keys=["prompt_a_0_0", "prompt_a_0_1", "prompt_a_2_0", "prompt_a_2_1", "prompt_a_3_0", "prompt_a_4_0"], + adv_estimator=AdvantageEstimator.GRPO, + ) + expected = compute_advantage( + batch_data.select_idxs([1, 3, 4, 5]), + adv_estimator=AdvantageEstimator.GRPO, + ) + gather_row_indices = [0, 0, 1, 1, 2, 3] + gather_col_indices = [0, 0, 1, 1, 0, 0] + adv_expected = ( + expected.batch["advantages"][gather_row_indices, gather_col_indices].unsqueeze(-1) + * result.batch["response_mask"] + ) + assert torch.equal(result.batch["advantages"], adv_expected) + assert torch.equal(result.batch["returns"], adv_expected) diff --git a/verl/tests/utils/_test_module.py b/verl/tests/utils/_test_module.py new file mode 100644 index 0000000000000000000000000000000000000000..5e10e65cff07378514d72920e63a77da28509537 --- /dev/null +++ b/verl/tests/utils/_test_module.py @@ -0,0 +1,31 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +# Test module for import_utils.load_extern_object testing +class TestClass: + """A test class to be imported by load_extern_object""" + + def __init__(self, value=None): + self.value = value or "default" + + def get_value(self): + return self.value + + +TEST_CONSTANT = "test_constant_value" + + +def test_function(): + return "test_function_result" diff --git a/verl/tests/utils/ckpt/test_checkpoint_cleanup_on_cpu.py b/verl/tests/utils/ckpt/test_checkpoint_cleanup_on_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..166208a4fc3493342818d6b530686b7e84015816 --- /dev/null +++ b/verl/tests/utils/ckpt/test_checkpoint_cleanup_on_cpu.py @@ -0,0 +1,139 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import shutil +import tempfile + +import pytest + + +class TestCheckpointCleanupLogic: + """Tests for checkpoint cleanup methods in BaseCheckpointManager.""" + + @pytest.fixture(autouse=True) + def setup(self): + """Set up test fixtures.""" + self.test_dir = tempfile.mkdtemp() + yield + shutil.rmtree(self.test_dir, ignore_errors=True) + + @pytest.fixture + def manager(self, monkeypatch): + """Create a minimal BaseCheckpointManager for testing.""" + import torch.distributed + + monkeypatch.setattr(torch.distributed, "get_rank", lambda: 0) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 1) + + from verl.utils.checkpoint.checkpoint_manager import BaseCheckpointManager + + class MockModel: + pass + + class MockOptimizer: + pass + + return BaseCheckpointManager( + model=MockModel(), + optimizer=MockOptimizer(), + lr_scheduler=None, + processing_class=None, + checkpoint_config=None, + ) + + def _create_checkpoint_dir(self, step: int) -> str: + """Create a mock checkpoint directory.""" + path = os.path.join(self.test_dir, f"global_step_{step}") + os.makedirs(path, exist_ok=True) + with open(os.path.join(path, "checkpoint.txt"), "w") as f: + f.write(f"step={step}") + return path + + def test_max_ckpt_1_preserves_existing_before_save(self, manager): + """ + Regression test: max_ckpt_to_keep=1 must NOT delete existing checkpoint before save. + """ + ckpt_100 = self._create_checkpoint_dir(100) + manager.previous_saved_paths = [ckpt_100] + + manager.ensure_checkpoint_capacity(max_ckpt_to_keep=1) + + assert os.path.exists(ckpt_100), "Bug: checkpoint deleted before save!" + assert manager.previous_saved_paths == [ckpt_100] + + def test_max_ckpt_1_deletes_old_after_save(self, manager): + """After save succeeds, old checkpoint should be deleted.""" + ckpt_100 = self._create_checkpoint_dir(100) + manager.previous_saved_paths = [ckpt_100] + + ckpt_200 = self._create_checkpoint_dir(200) + manager.register_checkpoint(ckpt_200, max_ckpt_to_keep=1) + + assert not os.path.exists(ckpt_100) + assert os.path.exists(ckpt_200) + assert manager.previous_saved_paths == [ckpt_200] + + def test_max_ckpt_2_keeps_one_before_save(self, manager): + """With max_ckpt_to_keep=2, pre-save cleanup keeps 1 checkpoint.""" + ckpt_100 = self._create_checkpoint_dir(100) + ckpt_200 = self._create_checkpoint_dir(200) + manager.previous_saved_paths = [ckpt_100, ckpt_200] + + manager.ensure_checkpoint_capacity(max_ckpt_to_keep=2) + + assert not os.path.exists(ckpt_100) + assert os.path.exists(ckpt_200) + assert len(manager.previous_saved_paths) == 1 + + def test_max_ckpt_0_keeps_all(self, manager): + """max_ckpt_to_keep=0 means unlimited - no deletions.""" + ckpt_100 = self._create_checkpoint_dir(100) + ckpt_200 = self._create_checkpoint_dir(200) + manager.previous_saved_paths = [ckpt_100, ckpt_200] + + manager.ensure_checkpoint_capacity(max_ckpt_to_keep=0) + ckpt_300 = self._create_checkpoint_dir(300) + manager.register_checkpoint(ckpt_300, max_ckpt_to_keep=0) + + assert os.path.exists(ckpt_100) + assert os.path.exists(ckpt_200) + assert os.path.exists(ckpt_300) + assert len(manager.previous_saved_paths) == 3 + + def test_full_save_cycle_max_ckpt_1(self, manager): + """Simulate multiple save cycles with max_ckpt_to_keep=1.""" + # First save + manager.ensure_checkpoint_capacity(1) + ckpt_100 = self._create_checkpoint_dir(100) + manager.register_checkpoint(ckpt_100, 1) + assert manager.previous_saved_paths == [ckpt_100] + + # Second save - existing checkpoint must survive pre-save + manager.ensure_checkpoint_capacity(1) + assert os.path.exists(ckpt_100), "Bug: checkpoint deleted before save!" + + ckpt_200 = self._create_checkpoint_dir(200) + manager.register_checkpoint(ckpt_200, 1) + assert not os.path.exists(ckpt_100) + assert manager.previous_saved_paths == [ckpt_200] + + # Third save + manager.ensure_checkpoint_capacity(1) + assert os.path.exists(ckpt_200), "Bug: checkpoint deleted before save!" + + ckpt_300 = self._create_checkpoint_dir(300) + manager.register_checkpoint(ckpt_300, 1) + assert not os.path.exists(ckpt_200) + assert manager.previous_saved_paths == [ckpt_300] diff --git a/verl/tests/utils/ckpt/test_esi_save_ckpt_on_cpu.py b/verl/tests/utils/ckpt/test_esi_save_ckpt_on_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..203494bd90bd9676fd615f5db5576e94c0219ee9 --- /dev/null +++ b/verl/tests/utils/ckpt/test_esi_save_ckpt_on_cpu.py @@ -0,0 +1,70 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os +import time +from datetime import datetime, timedelta +from unittest import TestCase + +from verl.utils.checkpoint.checkpoint_manager import should_save_ckpt_esi + + +class TestShouldSaveCkptEsi(TestCase): + def test_no_expiration_timestamp(self): + """Test case when no expiration timestamp is set""" + os.environ.pop("MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP", None) + os.environ.pop("SAGEMAKER_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP", None) + self.assertFalse(should_save_ckpt_esi(100)) + + def test_mlp_expiration_valid(self): + """Test valid MLP expiration timestamp requiring save""" + current_time = time.time() + os.environ["MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP"] = str(current_time + 90) + self.assertTrue(should_save_ckpt_esi(30)) # max_steps_duration=30 seconds + + def test_mlp_expiration_passed(self): + """Test expired MLP timestamp""" + current_time = time.time() + os.environ["MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP"] = str(current_time - 10) + self.assertFalse(should_save_ckpt_esi(30)) + + def test_mlp_invalid_timestamp(self): + """Test invalid MLP timestamp format""" + os.environ["MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP"] = "invalid" + self.assertFalse(should_save_ckpt_esi(30)) + + def test_mlp_expiration_not_reached(self): + """Test MLP expiration timestamp with insufficient remaining time""" + current_time = time.time() + os.environ["MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP"] = str(current_time + 200) + self.assertFalse(should_save_ckpt_esi(30)) # max_steps_duration=30 + + def test_aws_expiration_not_reached(self): + """Test AWS expiration timestamp with sufficient remaining time""" + now = datetime.now() + expiration = now + timedelta(minutes=100) # Exceeds 90-minute threshold + os.environ["SAGEMAKER_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP"] = str(int(expiration.timestamp())) + self.assertFalse(should_save_ckpt_esi(30 * 60)) + + def test_redundant_time(self): + """Test redundant_time parameter effect""" + current_time = time.time() + # Total required: 60+30+30=120 seconds + os.environ["MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP"] = str(current_time + 120) + self.assertTrue(should_save_ckpt_esi(30, redundant_time=30)) + + def test_zero_max_steps_duration(self): + """Test zero max_steps_duration""" + current_time = time.time() + os.environ["MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP"] = str(current_time + 60) + self.assertFalse(should_save_ckpt_esi(0)) diff --git a/verl/tests/utils/dataset/test_multiturn_sft_dataset_on_cpu.py b/verl/tests/utils/dataset/test_multiturn_sft_dataset_on_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..6430fd8a083589a015bdd9578204c4a9baaa70d5 --- /dev/null +++ b/verl/tests/utils/dataset/test_multiturn_sft_dataset_on_cpu.py @@ -0,0 +1,535 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Test the MultiTurnSFTDataset implementation +""" + +import os +from io import BytesIO +from pathlib import Path + +import pandas as pd +import pytest +import torch +from PIL import Image +from tensordict import TensorDict +from torch.utils.data import DistributedSampler +from torchdata.stateful_dataloader import StatefulDataLoader +from transformers.utils import get_json_schema + +from verl.utils import hf_processor, hf_tokenizer +from verl.utils.dataset.dataset_utils import DatasetPadMode, SFTTensorCollator +from verl.utils.dataset.multiturn_sft_dataset import MultiTurnSFTDataset +from verl.utils.model import extract_multi_modal_inputs + +custom_model_prefix = Path("~/models").expanduser().resolve() + + +@pytest.mark.parametrize( + "model_path, ignore_input_ids_mismatch", + [ + (f"{custom_model_prefix}/Qwen/Qwen2.5-0.5B", False), + (f"{custom_model_prefix}/Qwen/Qwen3-0.6B", True), + (f"{custom_model_prefix}/Qwen/Qwen3.5-0.8B", False), + ], +) +def test_multiturn_sft_dataset(model_path: str, ignore_input_ids_mismatch: bool): + print(f"Starting test... model_path={model_path}, ignore_input_ids_mismatch={ignore_input_ids_mismatch}") + # Create a temporary parquet file with test data + test_data = { + "messages": [ + [ + {"role": "user", "content": "What is 2+2?"}, + {"role": "assistant", "content": "2+2 equals 4."}, + {"role": "tool", "content": "And what is 4+4?"}, + {"role": "assistant", "content": "4+4 equals 8."}, + ], + [ + # {"role": "system", "content": "You are a powerful assistant."}, + {"role": "user", "content": "Tell me a joke."}, + {"role": "assistant", "content": "Why did the chicken cross the road?"}, + {"role": "tool", "content": "Why?"}, + {"role": "assistant", "content": "To get to the other side!"}, + ], + ] + } + + # Create test directory if it doesn't exist + os.makedirs("test_data", exist_ok=True) + test_file = "test_data/test.parquet" + + # Save test data to parquet + df = pd.DataFrame(test_data) + df.to_parquet(test_file) + + # Initialize tokenizer and dataset + tokenizer = hf_tokenizer(model_path) + # processor = hf_processor(model_path) + processor = None + config = { + "max_length": 512, + "truncation": "error", + "multiturn": {"messages_key": "messages"}, + "ignore_input_ids_mismatch": ignore_input_ids_mismatch, + } + dataset = MultiTurnSFTDataset(parquet_files=test_file, tokenizer=tokenizer, processor=processor, config=config) + + # Test 1: Dataset Length + assert len(dataset) == 2, f"Expected dataset length 2, got {len(dataset)}" + + # Get items for testing + item0 = dataset[0] # Math conversation + item1 = dataset[1] # Joke conversation + + # Test 2: Required Keys and Types + required_keys = ["input_ids", "attention_mask", "position_ids", "loss_mask"] + for key in required_keys: + assert key in item0, f"Missing key {key} in dataset item" + assert isinstance(item0[key], torch.Tensor), f"Expected torch.Tensor for {key}" + assert item0[key].dtype == torch.long, f"Expected torch.long for {key}, got {item0[key].dtype}" + + # Test 3: Shape Consistency + assert item0["loss_mask"].shape == item0["input_ids"].shape, "Loss mask shape doesn't match input_ids shape" + assert item0["attention_mask"].shape == item0["input_ids"].shape, ( + "Attention mask shape doesn't match input_ids shape" + ) + assert item0["position_ids"].shape == item0["input_ids"].shape, "Position IDs shape doesn't match input_ids shape" + + # Test 4: Loss Mask Pattern - Math Conversation + loss_mask0 = item0["loss_mask"] + input_ids0 = item0["input_ids"] + + # Find assistant response positions + assistant_positions0 = torch.where(loss_mask0 == 1)[0] + assert len(assistant_positions0) > 0, "No assistant positions found in loss mask" + + # Decode and verify assistant responses + assistant_text0 = tokenizer.decode(input_ids0[loss_mask0 == 1]) + print(f"Math conversation assistant text: {assistant_text0}") + assert "2+2 equals 4" in assistant_text0, "First assistant response not found" + assert "4+4 equals 8" in assistant_text0, "Second assistant response not found" + + # Test 5: Loss Mask Pattern - Joke Conversation + loss_mask1 = item1["loss_mask"] + input_ids1 = item1["input_ids"] + + # Find assistant response positions + assistant_positions1 = torch.where(loss_mask1 == 1)[0] + assert len(assistant_positions1) > 0, "No assistant positions found in loss mask" + + # Decode and verify assistant responses + assistant_text1 = tokenizer.decode(input_ids1[loss_mask1 == 1]) + print(f"Joke conversation assistant text: {assistant_text1}") + assert "chicken cross the road" in assistant_text1, "First assistant response not found" + assert "other side" in assistant_text1, "Second assistant response not found" + + # Test 6: Attention Mask Pattern + attention_mask0 = item0["attention_mask"] + sequence_length = torch.sum(attention_mask0) + assert sequence_length > 0, "No tokens marked as attended in attention mask" + assert torch.all(attention_mask0[:sequence_length] == 1), "Incorrect attention mask pattern" + if sequence_length < len(attention_mask0): + assert torch.all(attention_mask0[sequence_length:] == 0), "Padding not properly masked" + + # Test 7: Position IDs Pattern + position_ids0 = item0["position_ids"] + assert torch.equal(position_ids0[:sequence_length], torch.arange(sequence_length)), ( + "Position IDs not sequential for non-padded tokens" + ) + if sequence_length < len(position_ids0): + assert torch.all(position_ids0[sequence_length:] == 0), "Padding position IDs not zero" + + # Test 8: Verify loss mask for assistant responses + # Get the full conversation text + full_text = tokenizer.decode(input_ids0) + print(f"\nFull conversation text:\n{full_text}") + + # Get the assistant responses + assistant_text = tokenizer.decode(input_ids0[loss_mask0 == 1]) + print(f"\nAssistant responses (from loss mask):\n{assistant_text}") + + # Verify that loss mask is set for all assistant responses + for msg in test_data["messages"][0]: # First conversation + if msg["role"] == "assistant": + # The content should appear in the masked text + assert msg["content"] in assistant_text, f"Assistant message '{msg['content']}' not found in masked text" + + # The content should NOT appear in the non-masked text + non_assistant_text = tokenizer.decode(input_ids0[loss_mask0 == 0]) + assert msg["content"] not in non_assistant_text, ( + f"Assistant message '{msg['content']}' found in non-assistant text" + ) + + # Test 9: Verify non-assistant parts have loss_mask=0 + # Get non-assistant text + non_assistant_text = tokenizer.decode(input_ids0[loss_mask0 == 0]) + print(f"\nNon-assistant text (from loss mask):\n{non_assistant_text}") + + # Verify that system and user messages are in the non-assistant text + for msg in test_data["messages"][0]: # First conversation + if msg["role"] in ["system", "user"]: + assert msg["content"] in non_assistant_text, ( + f"{msg['role'].title()} message '{msg['content']}' not found in non-assistant text" + ) + + # And verify they're NOT in the assistant text + assert msg["content"] not in assistant_text, ( + f"{msg['role'].title()} message '{msg['content']}' found in assistant text" + ) + + # Test 10: Verify padding behavior + padding_config = { + "max_length": 1024, + "truncation": "error", + "multiturn": {"messages_key": "messages"}, + "ignore_input_ids_mismatch": ignore_input_ids_mismatch, + } + small_dataset = MultiTurnSFTDataset( + parquet_files=test_file, tokenizer=tokenizer, processor=processor, config=padding_config + ) + padded_item = small_dataset[0] + + # Get actual sequence length (before padding) + actual_length = torch.sum(padded_item["attention_mask"]) + + # Verify padding tokens + assert torch.all(padded_item["input_ids"][actual_length:] == tokenizer.pad_token_id), ( + "Padding tokens not set correctly" + ) + assert torch.all(padded_item["attention_mask"][actual_length:] == 0), "Attention mask not set correctly for padding" + assert torch.all(padded_item["loss_mask"][actual_length:] == 0), "Loss mask not set correctly for padding" + + # test no-padding + config = { + "max_length": 512, + "truncation": "error", + "multiturn": {"messages_key": "messages"}, + "pad_mode": "no_padding", + "ignore_input_ids_mismatch": ignore_input_ids_mismatch, + } + dataset = MultiTurnSFTDataset(parquet_files=test_file, tokenizer=tokenizer, processor=processor, config=config) + + item0 = dataset[0] + + # Verify that the output contains expected keys for no-padding mode + required_keys = ["input_ids", "position_ids", "loss_mask"] + for key in required_keys: + assert key in item0, f"Missing key {key} in no-padding mode dataset item" + assert isinstance(item0[key], torch.Tensor), f"Expected torch.Tensor for {key} in no-padding mode" + + # make sure assistant_text matches with expected + assistant_text = tokenizer.decode(item0["input_ids"][item0["loss_mask"] == 1]) + assert assistant_text == "2+2 equals 4.<|im_end|>\n4+4 equals 8.<|im_end|>\n" + + print("All tests passed!") + print("Starting test...") + + +@pytest.mark.parametrize( + "model_path, apply_chat_template_kwargs", + [ + (f"{custom_model_prefix}/openai/gpt-oss-20b", {"model_identity": "You are a helpful assistant."}), + ], +) +def test_multiturn_sft_dataset_with_chat_template_kwargs(model_path: str, apply_chat_template_kwargs: dict): + """Test that custom apply_chat_template_kwargs are forwarded to system prompt + measurement so the loss mask is not shifted when kwargs change tokenization. + + Some chat templates embed configurable fields (e.g. model_identity) in the + system prompt. If these kwargs are not forwarded to system prompt length + measurement, the per-turn strip length is wrong, causing role markers to be + removed and the loss mask to shift. + """ + test_data = { + "messages": [ + [ + {"role": "user", "content": "What is 2+2?"}, + {"role": "assistant", "content": "2+2 equals 4."}, + ], + [ + {"role": "user", "content": "Tell me a joke."}, + {"role": "assistant", "content": "Why did the chicken cross the road?"}, + ], + ] + } + + os.makedirs("test_data", exist_ok=True) + test_file = "test_data/test_kwargs.parquet" + df = pd.DataFrame(test_data) + df.to_parquet(test_file) + + tokenizer = hf_tokenizer(model_path) + + config = { + "max_length": 1024, + "truncation": "error", + "pad_mode": "no_padding", + "apply_chat_template_kwargs": apply_chat_template_kwargs, + } + dataset = MultiTurnSFTDataset(parquet_files=test_file, tokenizer=tokenizer, processor=None, config=config) + + for idx in range(len(dataset)): + item = dataset[idx] + input_ids = item["input_ids"] + loss_mask = item["loss_mask"] + + assistant_text = tokenizer.decode(input_ids[loss_mask == 1]) + non_assistant_text = tokenizer.decode(input_ids[loss_mask == 0]) + + for msg in test_data["messages"][idx]: + if msg["role"] == "assistant": + assert msg["content"] in assistant_text, ( + f"Assistant message '{msg['content']}' not found in masked text. " + f"This may indicate system prompt length mismatch when using " + f"custom chat_template_kwargs." + ) + assert msg["content"] not in non_assistant_text, ( + f"Assistant message '{msg['content']}' found in non-assistant text" + ) + elif msg["role"] in ["system", "user"]: + assert msg["content"] in non_assistant_text, ( + f"{msg['role'].title()} message '{msg['content']}' not found in non-assistant text" + ) + assert msg["content"] not in assistant_text, ( + f"{msg['role'].title()} message '{msg['content']}' found in assistant text" + ) + + print("All chat_template_kwargs tests passed!") + + +def generate_image(description: str, size: str = "256x256"): + """Generate a simple image based on description. + + Args: + description: The description of the image to generate. + size: The size of the image. Defaults to "256x256". (choices: ["256x256", "512x512"]) + + Returns: + A generated image + """ + ... + + +@pytest.fixture +def vlm_data_file(): + test_data = [ + # sample 0: single turn with image input + { + "messages": [ + { + "role": "user", + "content": "Describe this image.", + }, + { + "role": "assistant", + "content": "The image is a red square.", + }, + ], + "images": [Image.new("RGB", (300, 300), color="red")], + "tools": [], + }, + # sample 1: single turn with multiple images input + { + "messages": [ + { + "role": "user", + "content": "Compare these images.", + }, + { + "role": "assistant", + "content": "The first image is a red square and the second image is a green square.", + }, + ], + "images": [Image.new("RGB", (100, 100), color="red"), Image.new("RGB", (100, 300), color="green")], + "tools": [], + }, + # sample 2: multi turn with image input and tool generated image + { + "messages": [ + { + "role": "user", + "content": "Describe this image.", + }, + { + "role": "assistant", + "content": "Let's generate a zoom-in image.", + "tool_calls": [ + { + "function": {"arguments": {"bbox_2d": "[0, 1, 2, 4]"}, "name": "image_zoom_in_tool"}, + "type": "function", + } + ], + }, + { + "role": "tool", + "content": "Generated image.", + }, + {"role": "assistant", "content": "The zoom-in image is a red square."}, + ], + "images": [Image.new("RGB", (300, 500), color="red"), Image.new("RGB", (100, 100), color="red")], + "tools": [get_json_schema(generate_image)], + }, + # sample 3: single turn without image input + { + "messages": [ + {"role": "user", "content": "How is the weather today?"}, + {"role": "assistant", "content": "The weather is sunny."}, + ], + "images": [], + "tools": [], + }, + ] + + # Create test directory if it doesn't exist + os.makedirs("test_data", exist_ok=True) + test_file = "test_data/test_vlm.parquet" + + # Save test data to parquet + df = pd.DataFrame(test_data) + + def serialize_image(img): + if isinstance(img, Image.Image): + img_byte_arr = BytesIO() + img.save(img_byte_arr, format="PNG") + return {"bytes": img_byte_arr.getvalue()} + return img + + df["images"] = df["images"].apply(lambda x: [serialize_image(img) for img in x]) + + df.to_parquet(test_file) + return test_file + + +@pytest.mark.parametrize( + "model_path", + [ + f"{custom_model_prefix}/Qwen/Qwen3-VL-2B-Instruct", + f"{custom_model_prefix}/Qwen/Qwen3.5-0.8B", + ], +) +def test_multiturn_sft_vlm_dataset_on_cpu(model_path, vlm_data_file): + df = pd.read_parquet(vlm_data_file) + tokenizer = hf_tokenizer(model_path) + processor = hf_processor(model_path) + config = {"max_length": 1024, "pad_mode": "no_padding", "truncation": "error", "messages_key": "messages"} + dataset = MultiTurnSFTDataset(parquet_files=vlm_data_file, tokenizer=tokenizer, processor=processor, config=config) + assert dataset.pad_mode == DatasetPadMode.NO_PADDING + + for i in range(len(dataset)): + item = dataset[i] + input_ids = item["input_ids"] + loss_mask = item["loss_mask"] + position_ids = item["position_ids"] + pixel_values = item.get("multi_modal_inputs", {}).get("pixel_values") + image_grid_thw = item.get("multi_modal_inputs", {}).get("image_grid_thw") + + assert input_ids.shape == loss_mask.shape, "Shapes of input_ids and loss_mask must be equal" + assert position_ids.dim() == 2, "position_ids must be 2-dimensional" + assert position_ids.shape[0] == 4, f"position_ids[0] should be 4: {position_ids[0]}" + assert position_ids.shape[1] == input_ids.shape[0] + + # 1. verify input_ids without assistant text + text = tokenizer.decode(input_ids[loss_mask == 0], skip_special_tokens=True) + print(f"Text without assistant: {repr(text)}") + for message in df["messages"][i]: + if message["role"] != "assistant": + content = message["content"].replace("", "") + assert content in text, f"user/tool text should be in the input_ids: {text}" + + # 2. verify input_ids with assistant text + text = tokenizer.decode(input_ids[loss_mask == 1], skip_special_tokens=True) + print(f"Text with assistant: {repr(text)}") + for message in df["messages"][i]: + if message["role"] == "assistant": + assert message["content"] in text, f"Assistant text should be in the input_ids: {text}" + assert "assistant" not in text, f"Assistant token should not be in the input_ids: {text}" + + # 3. verify image token match with image_grid_thw + if len(df["images"][i]) > 0: + patch_size = processor.image_processor.patch_size + temporal_patch_size = processor.image_processor.temporal_patch_size + merge_size = processor.image_processor.merge_size + num_patches = image_grid_thw.prod(dim=1).sum() + assert image_grid_thw.shape == (len(df["images"][i]), 3), ( + f"image_grid_thw: {image_grid_thw.shape} should have shape ({len(df['images'][i])}, 3)" + ) + assert pixel_values.shape == (num_patches, 3 * temporal_patch_size * patch_size * patch_size), ( + f"pixel_values: {pixel_values.shape} should have shape ({num_patches}, {3 * patch_size * patch_size})" + ) + assert (input_ids == processor.image_token_id).sum() == num_patches // (merge_size**2) + else: + assert pixel_values is None, "pixel_values should be None when no image is provided" + assert image_grid_thw is None, "image_grid_thw should be None when no image is provided" + + +@pytest.mark.parametrize( + "model_path", + [ + f"{custom_model_prefix}/Qwen/Qwen3-VL-2B-Instruct", + ], +) +def test_multiturn_sft_vlm_dataloader_on_cpu(model_path, vlm_data_file): + df = pd.read_parquet(vlm_data_file) + tokenizer = hf_tokenizer(model_path) + processor = hf_processor(model_path) + config = {"max_length": 1024, "pad_mode": "no_padding", "truncation": "error", "messages_key": "messages"} + dataset = MultiTurnSFTDataset(parquet_files=vlm_data_file, tokenizer=tokenizer, processor=processor, config=config) + assert dataset.pad_mode == DatasetPadMode.NO_PADDING + + collate_fn = SFTTensorCollator(DatasetPadMode.NO_PADDING) + sampler = DistributedSampler(dataset, shuffle=False, num_replicas=1, rank=0, drop_last=True) + batch_size = 2 + dataloader = StatefulDataLoader( + dataset=dataset, + batch_size=batch_size, + sampler=sampler, + collate_fn=collate_fn, + num_workers=0, + pin_memory=False, + drop_last=True, + ) + + for i, batch in enumerate(dataloader): + # 1. verify input_ids, loss_mask + input_ids = batch["input_ids"] + loss_mask = batch["loss_mask"] + assert input_ids.is_nested, "input_ids should be a nested tensor" + assert loss_mask.is_nested, "loss_mask should be a nested tensor" + assert input_ids.shape[0] == loss_mask.shape[0] == batch_size, "Shapes of input_ids, loss_mask must be equal" + + # 2. verify position_ids: (bs, 4, seq_len) + position_ids = batch["position_ids"] + assert position_ids.is_nested, "position_ids should be a nested tensor" + assert position_ids.dim() == 3, "position_ids must be 3-dimensional" + assert position_ids.shape[0] == batch_size + values = position_ids.values() + assert values.shape == (4, len(input_ids.values())) + + # 3. verify multi-modal data + td = TensorDict(**batch, batch_size=batch_size) + multi_modal_inputs = extract_multi_modal_inputs(td["multi_modal_inputs"]) + pixel_values = multi_modal_inputs["pixel_values"] + image_grid_thw = multi_modal_inputs["image_grid_thw"] + + num_images = sum([len(images) for images in df["images"][i * batch_size : (i + 1) * batch_size]]) + assert image_grid_thw.shape == (num_images, 3), ( + f"image_grid_thw: {image_grid_thw.shape} should have shape ({num_images}, 3)" + ) + patch_size = processor.image_processor.patch_size + temporal_patch_size = processor.image_processor.temporal_patch_size + num_patches = image_grid_thw.prod(dim=1).sum() + assert pixel_values.shape[0] == num_patches, ( + f"pixel_values: {pixel_values.shape} should have shape " + f"({num_patches}, 3 * {temporal_patch_size} * {patch_size} * {patch_size})" + ) diff --git a/verl/tests/utils/dataset/test_rl_collate_fn_on_cpu.py b/verl/tests/utils/dataset/test_rl_collate_fn_on_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..415595295e7fde5d4de648284091bc87c53b4a10 --- /dev/null +++ b/verl/tests/utils/dataset/test_rl_collate_fn_on_cpu.py @@ -0,0 +1,72 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import torch + + +def test_rl_collate_fn(): + from verl.utils.dataset.rl_dataset import collate_fn + + max_prompt_length = 5 + + test_data = [ + { + # test tensor + "input_ids": torch.randint(0, 10, (max_prompt_length,)), + # test fixed length (1) list within a batch + "messages": [{"role": "user", "content": "Hi."}], + # test variable length list within a batch + "raw_prompt_ids": [1, 2, 3, 4], + # test string + "ability": "math", + # test dict + "reward_model": {"ground_truth": 5, "style": "rule"}, + # test empty dict + "tools_kwargs": {}, + }, + { + "input_ids": torch.randint(0, 10, (max_prompt_length,)), + "messages": [{"role": "user", "content": "Hello."}], + "raw_prompt_ids": [1, 2, 3], + "ability": "toolcall", + "reward_model": { + "ground_truth": '[{"name": "rgb_to_cmyk", "arguments": {"r": 0, "g": 0, "b": 255}}]', + "style": "rule", + }, + "tools_kwargs": {}, + }, + ] + + batch_size = len(test_data) + batch = collate_fn(test_data) + + # Tensor part + assert batch["input_ids"].shape == (batch_size, max_prompt_length) + assert isinstance(batch["input_ids"], torch.Tensor) + + # Non-tensor parts + expected_types = { + "messages": list, + "raw_prompt_ids": list, + "ability": str, + "reward_model": dict, + "tools_kwargs": dict, + } + + for key, dtype in expected_types.items(): + assert batch[key].shape == (batch_size,), ( + f"Expected shape {(batch_size,)} for '{key}', but got {batch[key].shape}" + ) + assert isinstance(batch[key][0], dtype), ( + f"'{key}' should contain elements of type {dtype}, but got {type(batch[key][0])}" + ) diff --git a/verl/tests/utils/dataset/test_rl_dataset_on_cpu.py b/verl/tests/utils/dataset/test_rl_dataset_on_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..05ebdbab98ce217da3c84aaf450c5dd2fe1b5abf --- /dev/null +++ b/verl/tests/utils/dataset/test_rl_dataset_on_cpu.py @@ -0,0 +1,197 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import json +import os + +import pytest +import torch +from omegaconf import OmegaConf +from PIL import Image +from torch.utils.data import DataLoader + +from verl import DataProto +from verl.utils import hf_processor, hf_tokenizer +from verl.utils.dataset.rl_dataset import RLHFDataset, collate_fn + + +def get_gsm8k_data(): + # prepare test dataset + local_folder = os.path.expanduser("~/data/gsm8k/") + local_path = os.path.join(local_folder, "train.parquet") + os.makedirs(local_folder, exist_ok=True) + return local_path + + +def test_rl_dataset(): + tokenizer = hf_tokenizer(os.path.expanduser("~/models/deepseek-ai/deepseek-coder-1.3b-instruct")) + local_path = get_gsm8k_data() + config = OmegaConf.create( + { + "prompt_key": "prompt", + "max_prompt_length": 256, + "filter_overlong_prompts": True, + "filter_overlong_prompts_workers": 2, + } + ) + dataset = RLHFDataset(data_files=local_path, tokenizer=tokenizer, config=config) + + dataloader = DataLoader(dataset=dataset, batch_size=16, shuffle=True, drop_last=True, collate_fn=collate_fn) + + a = next(iter(dataloader)) + + tensors = {} + non_tensors = {} + + for key, val in a.items(): + if isinstance(val, torch.Tensor): + tensors[key] = val + else: + non_tensors[key] = val + + data_proto = DataProto.from_dict(tensors=tensors, non_tensors=non_tensors) + assert len(data_proto) == 16 + assert "raw_prompt" in data_proto.non_tensor_batch + + +def test_rl_dataset_with_max_samples(): + tokenizer = hf_tokenizer(os.path.expanduser("~/models/deepseek-ai/deepseek-coder-1.3b-instruct")) + local_path = get_gsm8k_data() + config = OmegaConf.create( + { + "prompt_key": "prompt", + "max_prompt_length": 256, + "filter_overlong_prompts": True, + "filter_overlong_prompts_workers": 2, + "max_samples": 5, + } + ) + dataset = RLHFDataset(data_files=local_path, tokenizer=tokenizer, config=config, max_samples=5) + assert len(dataset) == 5 + + +def test_image_rl_data(): + tokenizer = hf_tokenizer(os.path.expanduser("~/models/Qwen/Qwen2-VL-2B-Instruct")) + processor = hf_processor(os.path.expanduser("~/models/Qwen/Qwen2-VL-2B-Instruct")) + config = OmegaConf.create( + { + "prompt_key": "prompt", + "max_prompt_length": 1024, + "filter_overlong_prompts": True, + "filter_overlong_prompts_workers": None, # num_workers=1 hang in ci + } + ) + dataset = RLHFDataset( + data_files=os.path.expanduser("~/data/geo3k/train.parquet"), + tokenizer=tokenizer, + config=config, + processor=processor, + ) + + dataloader = DataLoader(dataset=dataset, batch_size=16, shuffle=True, drop_last=True, collate_fn=collate_fn) + + a = next(iter(dataloader)) + + tensors = {} + non_tensors = {} + + for key, val in a.items(): + if isinstance(val, torch.Tensor): + tensors[key] = val + else: + non_tensors[key] = val + + data_proto = DataProto.from_dict(tensors=tensors, non_tensors=non_tensors) + assert len(data_proto) == 16 + assert "images" not in data_proto.non_tensor_batch + + for prompt in data_proto.non_tensor_batch["raw_prompt"]: + assert len(prompt) == 1 + prompt = prompt[0] + role, content = prompt["role"], prompt["content"] + assert role == "user" + assert len(content) == 2 + assert content[0]["type"] == "image" and isinstance(content[0]["image"], Image.Image) + assert content[1]["type"] == "text" and isinstance(content[1]["text"], str) + + print("raw_prompt", data_proto.non_tensor_batch["raw_prompt"][0]) + + +@pytest.fixture +def video_data_file(): + data = [ + { + "problem_id": 17, + "problem": "How does the crowd's excitement change as the match progresses?", + "data_type": "video", + "prompt": [ + { + "role": "user", + "content": [ + {"type": "video", "video": "LLaVA-Video-178K/academic_source/activitynet/v_2g9GrshWQrU.mp4"}, + { + "type": "text", + "text": "How does the crowd's excitement change as the match progresses? " + "A. It fluctuates; B. It decreases; C. It builds up; D. It remains the same. " + "Put your answer in ", + }, + ], + } + ], + "problem_type": "multiple choice", + "solution": "C", + "data_source": "LLaVA-Video-178K/2_3_m_academic_v0_1", + } + ] * 30 + + # Create test directory if it doesn't exist + os.makedirs("test_data", exist_ok=True) + test_file = "test_data/test_video.json" + with open(test_file, "w") as f: + json.dump(data, f, indent=2) + + return test_file + + +def test_video_rl_data(video_data_file): + tokenizer = hf_tokenizer(os.path.expanduser("~/models/Qwen/Qwen2-VL-2B-Instruct")) + processor = hf_processor(os.path.expanduser("~/models/Qwen/Qwen2-VL-2B-Instruct")) + config = OmegaConf.create( + { + "prompt_key": "prompt", + "max_prompt_length": 1024, + "filter_overlong_prompts": False, + } + ) + dataset = RLHFDataset( + data_files=video_data_file, + tokenizer=tokenizer, + config=config, + processor=processor, + ) + + dataloader = DataLoader(dataset=dataset, batch_size=16, shuffle=True, drop_last=True, collate_fn=collate_fn) + batch = next(iter(dataloader)) + tensors = {} + non_tensors = {} + for key, val in batch.items(): + if isinstance(val, torch.Tensor): + tensors[key] = val + else: + non_tensors[key] = val + + data_proto = DataProto.from_dict(tensors=tensors, non_tensors=non_tensors) + assert len(data_proto) == 16 + assert "images" not in data_proto.non_tensor_batch + + print("raw_prompt", data_proto.non_tensor_batch["raw_prompt"][0]) diff --git a/verl/tests/utils/debug/test_metrics.py b/verl/tests/utils/debug/test_metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..1b2f7f8faa17dd024df4b92c3a3b1b81d48923e0 --- /dev/null +++ b/verl/tests/utils/debug/test_metrics.py @@ -0,0 +1,48 @@ +# Copyright 2025 Individual Contributor: TomQunChaoA +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import torch + +from verl.protocol import DataProto +from verl.utils.debug.metrics import calculate_debug_metrics + + +class TestMetrics(unittest.TestCase): + def test_calculate_debug_metrics(self): + data = DataProto.from_dict( + { + "rollout_log_probs": torch.tensor( + [ + [-1.5085, -0.1200, -0.6650, -0.4823, -0.1426, -1.5557, -2.8532, -0.3919, -0.4294, -0.4700], + [-0.0585, -0.0573, -0.4681, -0.5187, -0.7451, -1.2737, -0.0682, -0.4284, -0.5754, -0.0611], + ] + ), + "old_log_probs": torch.tensor( + [ + [-1.8636, -0.7863, -0.2136, -0.4376, -2.0257, -0.2579, -1.1547, -0.5203, -0.3802, -0.9872], + [-0.3507, -0.5426, -0.2725, -0.4637, -0.3577, -0.3733, -1.7560, -1.9542, -0.4229, -1.3098], + ] + ), + "loss_mask": torch.tensor([[1, 0, 0, 0, 1, 1, 0, 1, 1, 0], [1, 0, 1, 0, 1, 1, 1, 0, 1, 1]]), + "responses": torch.zeros((2, 10)), + } + ) + metrics = calculate_debug_metrics(data) + print(metrics) + assert metrics["training/rollout_probs_diff_valid"] == 1 + + +if __name__ == "__main__": + unittest.main() diff --git a/verl/tests/utils/megatron/test_pipeline_parallel.py b/verl/tests/utils/megatron/test_pipeline_parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..24a416987dae68089a3d26d18f34d5defbd14245 --- /dev/null +++ b/verl/tests/utils/megatron/test_pipeline_parallel.py @@ -0,0 +1,70 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from verl.model_merger.megatron_model_merger import get_dynamic_pipeline_shards +from verl.utils.megatron.pipeline_parallel import make_batch_generator + + +def test_make_batch_generator_no_vpp(): + batches = [1, 2, 3] + vpp_size = 1 + generator = make_batch_generator(batches, vpp_size) + assert list(generator) == batches + + +def test_make_batch_generator_with_vpp(): + batches = [{"data": 1}, {"data": 2}] + vpp_size = 2 + generators = make_batch_generator(batches, vpp_size) + assert isinstance(generators, list) + assert len(generators) == vpp_size + + # Check each generator yields the original batches + for gen in generators: + assert list(gen) == batches + + +def test_make_batch_generator_empty(): + batches = [] + vpp_size = 1 + generator = make_batch_generator(batches, vpp_size) + assert list(generator) == [] + + vpp_size = 3 + generators = make_batch_generator(batches, vpp_size) + assert len(generators) == vpp_size + for gen in generators: + assert list(gen) == [] + + +@pytest.mark.parametrize( + "layer_num,pp_size,gt", + [ + (61, 8, [6, 8, 8, 8, 8, 8, 8, 7]), + (61, 7, [8, 9, 9, 9, 9, 9, 8]), + (61, 1, [61]), + (61, 0, ValueError), + (10, 16, ValueError), + ], +) +def test_get_dynamic_pipeline_shards(layer_num, pp_size, gt): + if isinstance(gt, list): + shards = get_dynamic_pipeline_shards(layer_num, pp_size) + assert len(shards) == len(gt) == pp_size, f"Expected {pp_size} shards, got {len(shards)}" + assert all([shard == gt[i] for i, shard in enumerate(shards)]), f"Expected shards {gt}, got {shards}" + elif issubclass(gt, Exception): + with pytest.raises(gt): + shards = get_dynamic_pipeline_shards(layer_num, pp_size) diff --git a/verl/tests/utils/reward_score/reward_score/test_sandbox_fusion_on_cpu.py b/verl/tests/utils/reward_score/reward_score/test_sandbox_fusion_on_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..83aed24d054ddce33bc8fd311de2705fcca24776 --- /dev/null +++ b/verl/tests/utils/reward_score/reward_score/test_sandbox_fusion_on_cpu.py @@ -0,0 +1,747 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import multiprocessing +import os +import time +from concurrent.futures import ProcessPoolExecutor +from unittest.mock import patch + +import pytest + +# Import the function to be tested +from verl.utils.reward_score.sandbox_fusion.utils import check_correctness + +# Get SANDBOX_URL from environment variable +SANDBOX_URL = os.environ.get("SANDBOX_FUSION_URL") +# Define skip condition and reason +skip_reason = "SANDBOX_FUSION_URL environment variable not set" +skip_condition = not SANDBOX_URL + +# --- Test code (for real API calls) --- +CODE_SUCCESS = """ +import sys +data = sys.stdin.read() +if data == 'input1': + print('output1\\n', end='') +elif data == 'input2': + print('output2\\n', end='') +else: + print('unexpected input', end='') +""" + +CODE_WRONG_OUTPUT = """ +print('wrong_output\\n', end='') +""" + +CODE_COMPILE_ERROR = """ +a=b +""" + +CODE_RUNTIME_ERROR = """ +import sys +print("About to raise error", file=sys.stderr) +raise ValueError("This is a runtime error") +""" + +CODE_TIMEOUT = """ +import time +import sys +print("Sleeping...", file=sys.stderr) +time.sleep(10) # Sleep time should be longer than the timeout set in the test +print("Finished sleeping", file=sys.stderr) +""" + +# --- Test input/output data --- +INPUT_OUTPUT_VALID = {"inputs": ["input1", "input2"], "outputs": ["output1\n", "output2\n"]} + +INPUT_OUTPUT_SINGLE = {"inputs": ["input1"], "outputs": ["output1\n"]} + +INPUT_OUTPUT_MISMATCH = {"inputs": ["input1"], "outputs": ["output1\n", "output2\n"]} + +INPUT_OUTPUT_INVALID_MISSING_KEY = {"inputs": ["input1"]} + +# --- Integration test cases (calling real API) --- + + +@pytest.mark.skipif(skip_condition, reason=skip_reason) +def test_integration_success_correct(): + """Integration test: Code is correct, output is correct""" + results, metadata_list = check_correctness(SANDBOX_URL, INPUT_OUTPUT_VALID, CODE_SUCCESS) + assert results == [True, True] + assert metadata_list[0]["status"] == "success" + assert metadata_list[0]["stdout"] == "output1\n" + assert metadata_list[1]["status"] == "success" + assert metadata_list[1]["stdout"] == "output2\n" + + +@pytest.mark.skipif(skip_condition, reason=skip_reason) +def test_integration_success_wrong_output(): + """Integration test: Code runs successfully, but output is wrong""" + results, metadata_list = check_correctness(SANDBOX_URL, INPUT_OUTPUT_VALID, CODE_WRONG_OUTPUT) + assert results == [False, False] + assert metadata_list[0]["status"] == "wrong_answer" + assert metadata_list[0]["stdout"] == "wrong_output\n" + assert metadata_list[1]["status"] == "wrong_answer" + + +@pytest.mark.skipif(skip_condition, reason=skip_reason) +def test_integration_compile_error(): + """Integration test: Code causes compile error""" + results, metadata_list = check_correctness(SANDBOX_URL, INPUT_OUTPUT_VALID, CODE_COMPILE_ERROR, language="cpp") + assert results == [-4, -4] + assert metadata_list[0]["status"] == "compile_error" + assert metadata_list[1]["status"] == "compile_error" + + +@pytest.mark.skipif(skip_condition, reason=skip_reason) +def test_integration_runtime_error(): + """Integration test: Code causes runtime error""" + results, metadata_list = check_correctness(SANDBOX_URL, INPUT_OUTPUT_SINGLE, CODE_RUNTIME_ERROR) + assert results == [-2] + assert metadata_list[0]["status"] == "runtime_error" + # More assertions can be added based on the actual API response, e.g., exit_code, stderr + + +@pytest.mark.skipif(skip_condition, reason=skip_reason) +def test_integration_runtime_timeout(): + """Integration test: Code causes runtime timeout""" + test_timeout = 5 # Set a timeout shorter than the sleep time in CODE_TIMEOUT + results, metadata_list = check_correctness(SANDBOX_URL, INPUT_OUTPUT_SINGLE, CODE_TIMEOUT, timeout=test_timeout) + assert results == [-3] + assert metadata_list[0]["status"] == "timeout" + # More assertions can be added based on the actual API response, e.g., run_status + + +@pytest.mark.skipif(skip_condition, reason=skip_reason) +def test_integration_concurrency_high_load(): + """Integration test: High concurrency (100 cases) against real API with mixed results (success, wrong + answer, timeout)""" + concurrency_level = 100 + # Indices for different expected outcomes + wrong_answer_indices = {10, 25, 50} + timeout_indices = {5, 30, 60, 90} # Indices where we expect a timeout + + # Generate 100 input/output pairs and code + high_load_inputs = [] + high_load_outputs = [] + expected_results_map = {} # Store expected result for each index + + for i in range(concurrency_level): + if i in timeout_indices: + # Use a special input to trigger timeout in the code + high_load_inputs.append(f"input_timeout_{i}") + # Output doesn't matter for timeout, but keep it consistent + high_load_outputs.append(f"output_{i}\n") + expected_results_map[i] = -3 # Expect timeout + elif i in wrong_answer_indices: + high_load_inputs.append(f"input_{i}") + # Intentionally set wrong expected output + high_load_outputs.append(f"wrong_output_{i}\n") + expected_results_map[i] = False # Expect wrong answer + else: + high_load_inputs.append(f"input_{i}") + # Correct expected output + high_load_outputs.append(f"output_{i}\n") + expected_results_map[i] = True # Expect success + + high_load_in_outs = {"inputs": high_load_inputs, "outputs": high_load_outputs} + + # Code that handles normal inputs, and sleeps on specific "timeout" inputs + code_mixed_concurrent = """ +import sys +import time +data = sys.stdin.read() +if data.startswith('input_timeout_'): + time.sleep(20) # Sleep longer than the test timeout + print(f"output_{data.split('_')[-1]}\\n", end='') # Still print something in case it finishes early +elif data.startswith('input_'): + print(f"output_{data.split('_')[-1]}\\n", end='') +else: + print("unknown_input\\n", end='') +""" + # Set a reasonable timeout per case (must be less than the sleep time in the code) + test_timeout = 15 # Allow slightly more time due to potential API load, but less than 20s sleep + + start_time = time.time() + results, metadata_list = check_correctness( + SANDBOX_URL, + high_load_in_outs, + code_mixed_concurrent, # Use the new code + timeout=test_timeout, + ) + end_time = time.time() + duration = end_time - start_time + print( + f"\nHigh concurrency test ({concurrency_level} cases with {len(wrong_answer_indices)} wrong answers, " + f"{len(timeout_indices)} timeouts) duration: {duration:.2f} seconds" + ) + + # Verify results against the expected map + assert len(results) == concurrency_level, f"Expected {concurrency_level} results, got {len(results)}" + + correct_count = 0 + wrong_count = 0 + timeout_count = 0 + unexpected_results = [] + for i, r in enumerate(results): + expected = expected_results_map[i] + if r == expected: + if expected is True: + correct_count += 1 + elif expected is False: + wrong_count += 1 + elif expected == -3: + timeout_count += 1 + else: + unexpected_results.append((i, r, f"Expected {expected}")) + + print( + f"Correct results (True): {correct_count}/" + f"{concurrency_level - len(wrong_answer_indices) - len(timeout_indices)}" + ) + print(f"Expected wrong answers (False, correctly identified): {wrong_count}/{len(wrong_answer_indices)}") + print(f"Expected timeouts (-3, correctly identified): {timeout_count}/{len(timeout_indices)}") + + if unexpected_results: + print("Unexpected results found:") + for idx, res, expected_str in unexpected_results[:10]: # Print first 10 unexpected + print(f" Index {idx}: Got {res}, {expected_str}. Metadata: {metadata_list[idx]}") + raise AssertionError(f"Found {len(unexpected_results)} unexpected results.") + + assert correct_count == concurrency_level - len(wrong_answer_indices) - len(timeout_indices), ( + "Incorrect number of successful results" + ) + assert wrong_count == len(wrong_answer_indices), "Incorrect number of identified wrong answers" + assert timeout_count == len(timeout_indices), "Incorrect number of identified timeouts" + + # Verify metadata count and basic status of one of each type + assert len(metadata_list) == concurrency_level + # Find the first correct index + first_correct_index = next( + i for i in range(concurrency_level) if i not in wrong_answer_indices and i not in timeout_indices + ) + assert metadata_list[first_correct_index]["status"] == "success" + assert metadata_list[first_correct_index]["stdout"] == f"output_{first_correct_index}\n" + + # Check the status of the first intentionally wrong case + first_wrong_index = min(wrong_answer_indices) + assert metadata_list[first_wrong_index]["status"] == "wrong_answer" + assert metadata_list[first_wrong_index]["stdout"] == f"output_{first_wrong_index}\n" + assert metadata_list[first_wrong_index]["expected_output"] == f"wrong_output_{first_wrong_index}\n" + + # Check the status of the first intentionally timeout case + first_timeout_index = min(timeout_indices) + assert metadata_list[first_timeout_index]["status"] == "timeout" + # For timeout, stdout might be None or empty depending on when the timeout occurred + # assert metadata_list[first_timeout_index]["stdout"] is None or metadata_list[first_timeout_index]["stdout"] == "" + + +# --- Unit test cases (using mock) --- + + +@patch("verl.utils.reward_score.sandbox_fusion.utils.call_sandbox_api") +def test_unit_concurrency_order(mock_call_sandbox_api): + sandbox_url = "mock_url" + generation = "print(input())" + language = "python" + timeout = 5 + in_outs = {"inputs": ["input1", "input2", "input3"], "outputs": ["output1", "output2", "output3"]} + + def side_effect(*args, **kwargs): + stdin = kwargs.get("stdin") + if stdin == "input1": + return ( + {"status": "Success", "run_result": {"status": "Finished", "stdout": "output1", "return_code": 0}}, + None, + ) + elif stdin == "input2": + time.sleep(0.1) + return ( + {"status": "Success", "run_result": {"status": "Finished", "stdout": "output2", "return_code": 0}}, + None, + ) + elif stdin == "input3": + return ( + {"status": "Success", "run_result": {"status": "Finished", "stdout": "output3", "return_code": 0}}, + None, + ) + else: + return (None, "Unknown input in mock") + + mock_call_sandbox_api.side_effect = side_effect + + results, metadata_list = check_correctness(sandbox_url, in_outs, generation, timeout, language) + + assert results == [True, True, True] + assert len(metadata_list) == 3 + assert metadata_list[0]["case_index"] == 0 + assert metadata_list[0]["status"] == "success" + assert metadata_list[1]["case_index"] == 1 + assert metadata_list[1]["status"] == "success" + assert metadata_list[2]["case_index"] == 2 + assert metadata_list[2]["status"] == "success" + assert mock_call_sandbox_api.call_count == 3 + + +@patch("verl.utils.reward_score.sandbox_fusion.utils.call_sandbox_api") +def test_unit_api_timeout_error_concurrent(mock_call_sandbox_api): + sandbox_url = "mock_url" + generation = "print(input())" + language = "python" + timeout = 5 + in_outs = {"inputs": ["input1", "input2_timeout", "input3"], "outputs": ["output1", "output2", "output3"]} + + api_error_message = "API Call Failed: Gateway Timeout (504) on attempt 3/3" + + def side_effect(*args, **kwargs): + stdin = kwargs.get("stdin") + if stdin == "input1": + return ( + {"status": "Success", "run_result": {"status": "Finished", "stdout": "output1", "return_code": 0}}, + None, + ) + elif stdin == "input2_timeout": + return (None, api_error_message) + elif stdin == "input3": + return ( + {"status": "Success", "run_result": {"status": "Finished", "stdout": "output3", "return_code": 0}}, + None, + ) + else: + return (None, "Unknown input in mock") + + mock_call_sandbox_api.side_effect = side_effect + + results, metadata_list = check_correctness(sandbox_url, in_outs, generation, timeout, language) + + assert results == [True, -1, True] + assert len(metadata_list) == 3 + assert metadata_list[0]["status"] == "success" + assert metadata_list[1]["status"] == "api_error" + assert metadata_list[1]["api_request_error"] == api_error_message + assert metadata_list[2]["status"] == "success" + assert mock_call_sandbox_api.call_count == 3 + + +# --- Constants for the new concurrency test --- +# Define a low global concurrency limit to test the semaphore's effect +MAX_GLOBAL_CONCURRENCY_LIMIT_TEST = 5 +# Define the number of processes used in the test +NUM_PROCESSES_TEST = 4 +# Define the number of tasks processed by check_correctness in each process (i.e., internal +# ThreadPoolExecutor's concurrency potential) +NUM_TASKS_PER_PROCESS_TEST = 3 +# Simulate API call duration to ensure calls can overlap +SIMULATED_API_CALL_DURATION_TEST = 0.2 # seconds + + +# --- Mock API call function for concurrency tracking --- +# This function will replace the real call_sandbox_api and use shared variables to track concurrency +def _mock_api_call_for_concurrency_tracking( + active_calls_counter, # multiprocessing.Value + max_calls_tracker, # multiprocessing.Value + call_lock, # multiprocessing.Lock + # Standard call_sandbox_api parameters + sandbox_fusion_url, + code, + stdin, + compile_timeout, + run_timeout, + memory_limit_mb, + language, +): + # entry_time = time.time() # For detailed logging + with call_lock: + active_calls_counter.value += 1 + if active_calls_counter.value > max_calls_tracker.value: + max_calls_tracker.value = active_calls_counter.value + # Optional debug log: + # print(f"[PID:{os.getpid()}-TID:{threading.get_ident()}] API Call Start. Active: " + # f"{active_calls_counter.value}, Max Observed: {max_calls_tracker.value}, Input: {stdin}") + + time.sleep(SIMULATED_API_CALL_DURATION_TEST) # Simulate actual work duration + + # exit_time = time.time() # For detailed logging + with call_lock: + active_calls_counter.value -= 1 + # Optional debug log: + # print(f"[PID:{os.getpid()}-TID:{threading.get_ident()}] API Call End. Active: " + # f"{active_calls_counter.value}, Input: {stdin}, Duration: {exit_time - entry_time:.2f}s") + + # Return a simulated successful API response + return { + "status": "Success", + "run_result": {"status": "Finished", "stdout": f"mock_output_for_{stdin}", "return_code": 0}, + }, None + + +# --- Worker function for ProcessPoolExecutor --- +# This function runs in each child process of ProcessPoolExecutor +def _process_pool_worker_for_concurrency_test( + sandbox_url, + in_outs, + generation, + memory_limit_mb, + language, + timeout, + mp_semaphore_for_check_correctness, + active_calls_counter, + max_calls_tracker, + call_lock, +): + # Corrected lambda to accept keyword arguments matching call_sandbox_api's usage + curried_mock_api_call = ( + lambda sandbox_fusion_url, code, stdin, compile_timeout, run_timeout, memory_limit_mb, language: ( + _mock_api_call_for_concurrency_tracking( + active_calls_counter, + max_calls_tracker, + call_lock, + sandbox_fusion_url, + code, + stdin, + compile_timeout, + run_timeout, + memory_limit_mb, + language, + ) + ) + ) + + # ---- START DEBUG PRINTS ---- + import os + + import verl.utils.reward_score.sandbox_fusion.utils + + print( + f"[Worker PID:{os.getpid()}] Original call_sandbox_api: " + f"{verl.utils.reward_score.sandbox_fusion.utils.call_sandbox_api}", + flush=True, + ) + # ---- END DEBUG PRINTS ---- + + with patch( + "verl.utils.reward_score.sandbox_fusion.utils.call_sandbox_api", side_effect=curried_mock_api_call + ) as mock_obj: + # ---- START DEBUG PRINTS ---- + print( + f"[Worker PID:{os.getpid()}] Patched call_sandbox_api: " + f"{verl.utils.reward_score.sandbox_fusion.utils.call_sandbox_api}", + flush=True, + ) + print(f"[Worker PID:{os.getpid()}] Mock object: {mock_obj}", flush=True) + # ---- END DEBUG PRINTS ---- + results, metadata_list = check_correctness( + sandbox_fusion_url=sandbox_url, + in_outs=in_outs, + generation=generation, + timeout=timeout, + memory_limit_mb=memory_limit_mb, + language=language, + concurrent_semaphore=mp_semaphore_for_check_correctness, # Pass multiprocessing.Semaphore + ) + # print(f"Process {os.getpid()} finished check_correctness. Processed {len(results)} tasks.") + return len(results) # Return the number of processed tasks for basic validation + + +# --- The actual test case for multiprocess concurrency control --- +def test_multiprocess_global_concurrency_limit_with_semaphore(): + """ + Tests that the global concurrent_semaphore (multiprocessing.Semaphore) + correctly limits the number of concurrent calls to call_sandbox_api + across multiple processes, each potentially running multiple threads + via check_correctness's internal ThreadPoolExecutor. + """ + manager = multiprocessing.Manager() + active_calls_counter = manager.Value("i", 0) # Current active mock API calls + max_calls_tracker = manager.Value("i", 0) # Observed maximum concurrent mock API calls + call_lock = manager.Lock() # Lock to protect counters + + # Create a multiprocessing.Semaphore instance, this is the global semaphore we are testing. + # It will be passed to check_correctness and used by _process_single_case to limit calls to call_sandbox_api. + global_mp_semaphore = manager.Semaphore(MAX_GLOBAL_CONCURRENCY_LIMIT_TEST) + + mock_sandbox_url = "mock_url_for_concurrency_test" + mock_generation = "pass" # Specific code content is not important as API call is mocked + mock_memory_limit_mb = 1024 + mock_language = "python" + mock_timeout = 5 # Timeout setting, not critical for mock calls + + # Input/output data for each process + # NUM_TASKS_PER_PROCESS_TEST tasks will be handled by check_correctness's internal ThreadPoolExecutor + process_in_outs = { + "inputs": [f"task_input_{i}" for i in range(NUM_TASKS_PER_PROCESS_TEST)], + "outputs": [f"task_output_{i}" for i in range(NUM_TASKS_PER_PROCESS_TEST)], + } + + futures = [] + total_tasks_expected_to_run = NUM_PROCESSES_TEST * NUM_TASKS_PER_PROCESS_TEST + + test_start_time = time.time() + + with ProcessPoolExecutor(max_workers=NUM_PROCESSES_TEST) as executor: + for i in range(NUM_PROCESSES_TEST): + future = executor.submit( + _process_pool_worker_for_concurrency_test, # Worker function + mock_sandbox_url, + process_in_outs, + mock_generation, + mock_memory_limit_mb, + mock_language, + mock_timeout, + global_mp_semaphore, # Global semaphore to test + active_calls_counter, # Shared variables for tracking + max_calls_tracker, + call_lock, + ) + futures.append(future) + + # Wait for all processes to complete and collect results + num_tasks_processed_per_worker = [f.result() for f in futures] + test_end_time = time.time() + total_execution_time = test_end_time - test_start_time + + # Print some test statistics for debugging and validation + print("\n--- Global Concurrency Test Stats ---") + print(f"Semaphore Limit (MAX_GLOBAL_CONCURRENCY_LIMIT_TEST): {MAX_GLOBAL_CONCURRENCY_LIMIT_TEST}") + print(f"Number of Processes (NUM_PROCESSES_TEST): {NUM_PROCESSES_TEST}") + print(f"Tasks per Process (NUM_TASKS_PER_PROCESS_TEST): {NUM_TASKS_PER_PROCESS_TEST}") + print(f"Total Tasks Submitted: {total_tasks_expected_to_run}") + print(f"Simulated API Call Duration: {SIMULATED_API_CALL_DURATION_TEST}s") + print(f"Total Test Execution Time: {total_execution_time:.2f}s") + print(f"Max Concurrent Mock API Calls Observed: {max_calls_tracker.value}") + # print(f"Tasks processed per worker: {num_tasks_processed_per_worker}") + + # Verify that all submitted tasks have been processed + assert sum(num_tasks_processed_per_worker) == total_tasks_expected_to_run, ( + "Mismatch in the number of tasks processed." + ) + + # Verify that the mock API was called at least once + assert max_calls_tracker.value > 0, "The mocked API call_sandbox_api was not called." + + # Core assertion: Observed maximum concurrent calls should not exceed the semaphore's limit + assert max_calls_tracker.value <= MAX_GLOBAL_CONCURRENCY_LIMIT_TEST, ( + f"Observed concurrency ({max_calls_tracker.value}) exceeded semaphore limit " + f"({MAX_GLOBAL_CONCURRENCY_LIMIT_TEST})." + ) + + # Optional: Rough check on execution time to verify semaphore is working to limit concurrency + # Theoretical minimum execution time = (Total tasks / Concurrency limit) * Single task duration + # Actual time will be longer due to various overheads + min_expected_duration = ( + total_tasks_expected_to_run * SIMULATED_API_CALL_DURATION_TEST + ) / MAX_GLOBAL_CONCURRENCY_LIMIT_TEST + # print(f"Minimum Expected Execution Time (approx): {min_expected_duration:.2f}s") + # Allow some margin, e.g., 80% of theoretical minimum time + assert total_execution_time >= min_expected_duration * 0.8, ( + f"Total execution time ({total_execution_time:.2f}s) was unexpectedly short, suggesting the " + f"semaphore might not be effectively limiting concurrency as expected " + f"(min expected: {min_expected_duration * 0.8:.2f}s)." + ) + + +# Ensure there is no more code after this point if these were the last functions. +# If there was other code, it would follow here. +def test_unit_invalid_input_format(): + """Unit test: Invalid in_outs format passed""" + results, metadata_list = check_correctness(SANDBOX_URL, None, CODE_SUCCESS) + assert results == [-1] + assert metadata_list[0]["error"] == "Invalid input/output data" + + results, metadata_list = check_correctness(SANDBOX_URL, {}, CODE_SUCCESS) + assert results == [-1] + assert metadata_list[0]["error"] == "Invalid input/output data" + + results, metadata_list = check_correctness(SANDBOX_URL, INPUT_OUTPUT_INVALID_MISSING_KEY, CODE_SUCCESS) + assert results == [-1] + assert metadata_list[0]["error"] == "Invalid input/output data" + + +@pytest.mark.skipif(skip_condition, reason=skip_reason) +def test_unit_input_output_mismatch(): + """Unit test: Mismatch between the number of inputs and outputs""" + results, metadata_list = check_correctness(SANDBOX_URL, INPUT_OUTPUT_MISMATCH, CODE_SUCCESS) + assert results == [-1] + assert len(metadata_list) == 1 + assert metadata_list[0]["error"] == "Input/output count mismatch" + + +@pytest.mark.skipif(skip_condition, reason=skip_reason) +def test_integration_concurrency_all_timeout(): + """Integration test: High concurrency (100 cases) against real API, all causing timeout""" + concurrency_level = 100 + code_infinite_loop = """ +def knight_moves(X, Y): + MOD = 10**9 + 7 + dp = [[0] * (Y + 1) for _ in range(X + 1)] + dp[0][0] = 1 + for i in range(1, X + 1): + for j in range(1, Y + 1): + dp[i][j] = (dp[i - 1][j] + dp[i][j - 1]) % MOD + return dp[X][Y] + +def solve(): + X, Y = map(int, input().split()) + print(knight_moves(X, Y)) + +if __name__ == "__main__": + solve() + """ + + # Generate 100 simple input/output pairs (content doesn't matter) + timeout_inputs = ["324 384429" for i in range(concurrency_level)] + timeout_outputs = [f"output_{i}\n" for i in range(concurrency_level)] + timeout_in_outs = {"inputs": timeout_inputs, "outputs": timeout_outputs} + + # Set a timeout for the test cases + test_timeout = 10 # Set a timeout value + + start_time = time.time() + results, metadata_list = check_correctness(SANDBOX_URL, timeout_in_outs, code_infinite_loop, timeout=test_timeout) + end_time = time.time() + duration = end_time - start_time + print(f"\nHigh concurrency all timeout test ({concurrency_level} cases) duration: {duration:.2f} seconds") + + # Verify all results are -3 (timeout) + assert len(results) == concurrency_level, f"Expected {concurrency_level} results, got {len(results)}" + all_timed_out = all(r == -3 for r in results) + if not all_timed_out: + non_timeout_indices = [i for i, r in enumerate(results) if r != -3] + print(f"Indices that did not time out: {non_timeout_indices}") + # Print metadata for the first few non-timeout cases for debugging + for i in non_timeout_indices[:5]: + print(f"Metadata for non-timeout case {i}: {metadata_list[i]}") + assert all_timed_out, f"Not all {concurrency_level} concurrent tests resulted in timeout (-3). Results: {results}" + + # Verify metadata count and status of the first case + assert len(metadata_list) == concurrency_level + assert metadata_list[0]["status"] == "timeout" + + +@pytest.mark.skipif(skip_condition, reason=skip_reason) +def test_fn_name_success_single_case(): + """Tests successful execution for a single test case with fn_name. + from livecodebench/code_generation_lite test 510 + """ + generation_code = """ +class Solution: + def occurrencesOfElement(self, nums: List[int], queries: List[int], x: int) -> List[int]: + positions = defaultdict(list) + for idx, num in enumerate(nums): + positions[num].append(idx) + + x_positions = positions[x] + answer = [] + for k in queries: + if k > len(x_positions): + answer.append(-1) + else: + answer.append(x_positions[k-1]) + return answer +""" + in_outs = { + "fn_name": "occurrencesOfElement", + "inputs": ["[1, 3, 1, 7]\n[1, 3, 2, 4]\n1", "[1, 2, 3]\n[10]\n5"], + "outputs": ["[0, -1, 2, -1]", "[-1]"], + } + + # Use a short timeout for fast tests + results, metadata_list = check_correctness(SANDBOX_URL, in_outs, generation_code, timeout=5) + # from verl.utils.reward_score.prime_code import apps_check_correctness + # results, metadata_list = apps_check_correctness(in_outs=in_outs, generation=generation_code, + # timeout=50000, debug=True) + + assert results == [True, True] + assert "error" not in metadata_list[0] + assert metadata_list[0].get("status") != "compile_error" + assert metadata_list[0].get("status") != "runtime_error" + + +@pytest.mark.skipif(skip_condition, reason=skip_reason) +def test_none_and_empty_stdin_passed_correctly(): + """ + Tests that when stdin data is set to an empty string or None, it is still + is passed correctly to Sandbox Fusion as an empty string. + """ + echo_code = """ +import sys +print(f"You said '{sys.stdin.readline().strip()}'") +""" + in_outs = { + "inputs": [None, "", "hello"], + "outputs": ["You said ''", "You said ''", "You said 'hello'"], + } + + # Use a short timeout for fast tests + results, metadata_list = check_correctness(SANDBOX_URL, in_outs, echo_code, timeout=5) + + assert results == [True, True, True] + assert "error" not in metadata_list[0] + assert metadata_list[0].get("status") != "compile_error" + assert metadata_list[0].get("status") != "runtime_error" + + +@pytest.mark.skipif(skip_condition, reason=skip_reason) +def test_assert_case_success(): + """Tests successful execution for assert case. + from KodCode + """ + generation_code = """ +from typing import List, Tuple + +def merge_intervals(intervals: List[Tuple[int, int]]) -> List[Tuple[int, int]]: + if not intervals: + return [] + + # Sort intervals by the start time + intervals.sort(key=lambda x: x[0]) + + merged = [intervals[0]] + + for current in intervals[1:]: + last = merged[-1] + # If intervals overlap, merge them + if current[0] <= last[1]: + merged[-1] = (last[0], max(last[1], current[1])) + else: + merged.append(current) + + return merged +""" + test_cases = { + "fn_name": "merge_intervals", + "assert_case": [ + "assert merge_intervals([(0, 1), (3, 5), (4, 7), (6, 8), (10, 12)," + " (12, 14)]) == [(0, 1), (3, 8), (10, 14)]", + "assert merge_intervals([(1, 2), (2, 3), (3, 4)]) == [(1, 4)]", + "assert merge_intervals([(1, 2), (3, 4), (5, 6)]) == [(1, 2), (3, 4), (5, 5)]", + ], + } + + assert_cases = test_cases.get("assert_case") + test_cases.setdefault("inputs", ["" for _ in assert_cases]) + test_cases.setdefault("outputs", [None for _ in assert_cases]) + + # Use a short timeout for fast tests + results, metadata_list = check_correctness(SANDBOX_URL, test_cases, generation_code, timeout=5) + assert results == [True, True, -2] + for i in range(2): + assert "error" not in metadata_list[i] + assert metadata_list[i].get("status") == "success" + assert metadata_list[i].get("expected_output") is None + assert metadata_list[i].get("status") != "runtime_error" + assert "error" not in metadata_list[2] + assert metadata_list[2].get("status") != "success" + assert metadata_list[2].get("expected_output") is None + assert metadata_list[2].get("status") == "runtime_error" diff --git a/verl/tests/utils/reward_score/test_sandbox_on_cpu.py b/verl/tests/utils/reward_score/test_sandbox_on_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..ff8508de255bc24d9c9be6e13ede0eecb81e1459 --- /dev/null +++ b/verl/tests/utils/reward_score/test_sandbox_on_cpu.py @@ -0,0 +1,190 @@ +# Copyright 2024 PRIME team and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +import json +import os + +import pytest + +from verl.utils.reward_score import default_compute_score, sandbox_fusion +from verl.workers.reward_manager.prime import parallel_compute_score_async + +prime_math_answers = [ + """\\begin{bmatrix}\n -7 & 6 & -8 \\\\\n 11 & -9 & 12 \\\\\n 15 & -16 & 19 \n \\end{bmatrix}""", + """\\frac{\\sqrt{505}}{7}""", + """x^2 + y^2 + 4x - 6y + 13""", +] +prime_math_gts = [ + """\\begin{pmatrix}\n -7 & 6 & -8 \\\\\n 11 & -9 & 12 \\\\\n 15 & -16 & 19\n \\end{pmatrix}""", # mat test + """\\frac{\\sqrt{505}}{7}""", # frac test + """(x + 2)^2 + (y - 3)^2 """, # symbolic test +] + +prime_code_answers = [ + """import sys +from collections import deque + +def main(): + data = sys.stdin.read().split() + it = iter(data) + + # Read start and target positions + x0, y0, x1, y1 = int(next(it)), int(next(it)), int(next(it)), int(next(it)) + + n = int(next(it)) + allowed = set() + # The total number of allowed cells is at most 10^5. + for _ in range(n): + r = int(next(it)) + a = int(next(it)) + b = int(next(it)) + for c in range(a, b + 1): + allowed.add((r, c)) + + # Directions for the king (8 neighboring cells) + directions = [(-1, -1), (-1, 0), (-1, 1), + (0, -1), (0, 1), + (1, -1), (1, 0), (1, 1)] + + start = (x0, y0) + target = (x1, y1) + + # BFS initialization + queue = deque() + queue.append((x0, y0, 0)) + # Mark the starting cell as visited by removing it from allowed set. + allowed.discard(start) + + while queue: + x, y, moves = queue.popleft() + if (x, y) == target: + print(moves) + return + for dx, dy in directions: + nx, ny = x + dx, y + dy + if (nx, ny) in allowed: + allowed.remove((nx, ny)) + queue.append((nx, ny, moves + 1)) + + print(-1) + +if __name__ == '__main__': + main() +""" +] * 2 +prime_code_gts = [ + """{\n \"inputs\": [\n \"5 7 6 11\\n3\\n5 3 8\\n6 7 11\\n5 2 5\\n\",\n \"3 4 3 10\\n3\\n3 1 4\\n4 5 9\\n3 10 10\\n\",\n \"1 1 2 10\\n2\\n1 1 3\\n2 6 10\\n\",\n \"9 8 7 8\\n9\\n10 6 6\\n10 6 6\\n7 7 8\\n9 5 6\\n8 9 9\\n9 5 5\\n9 8 8\\n8 5 6\\n9 10 10\\n\",\n \"6 15 7 15\\n9\\n6 15 15\\n7 14 14\\n6 15 15\\n9 14 14\\n7 14 16\\n6 15 15\\n6 15 15\\n7 14 14\\n8 15 15\\n\",\n \"13 16 20 10\\n18\\n13 16 16\\n20 10 10\\n19 10 10\\n12 15 15\\n20 10 10\\n18 11 11\\n19 10 10\\n19 10 10\\n20 10 10\\n19 10 10\\n20 10 10\\n20 10 10\\n19 10 10\\n18 11 11\\n13 16 16\\n12 15 15\\n19 10 10\\n19 10 10\\n\",\n \"89 29 88 30\\n16\\n87 31 31\\n14 95 95\\n98 88 89\\n96 88 88\\n14 97 97\\n13 97 98\\n100 88 88\\n88 32 32\\n99 88 89\\n90 29 29\\n87 31 31\\n15 94 96\\n89 29 29\\n88 32 32\\n97 89 89\\n88 29 30\\n\",\n \"30 14 39 19\\n31\\n35 7 11\\n37 11 12\\n32 13 13\\n37 5 6\\n46 13 13\\n37 14 14\\n31 13 13\\n43 13 19\\n45 15 19\\n46 13 13\\n32 17 17\\n41 14 19\\n30 14 14\\n43 13 17\\n34 16 18\\n44 11 19\\n38 13 13\\n40 12 20\\n37 16 18\\n46 16 18\\n34 10 14\\n36 9 10\\n36 15 19\\n38 15 19\\n42 13 19\\n33 14 15\\n35 15 19\\n33 17 18\\n39 12 20\\n36 5 7\\n45 12 12\\n\",\n \"2 1 1 1\\n2\\n1 1 2\\n2 1 2\\n\",\n \"1 1 1 2\\n5\\n1000000000 1 10000\\n19920401 1188 5566\\n1000000000 1 10000\\n1 1 10000\\n5 100 200\\n\",\n \"1 1 1000000000 2\\n5\\n1000000000 1 10000\\n19920401 1188 5566\\n1000000000 1 10000\\n1 1 10000\\n5 100 200\\n\"\n ],\n \"outputs\": [\n \"4\\n\",\n \"6\\n\",\n \"-1\\n\",\n \"2\\n\",\n \"1\\n\",\n \"-1\\n\",\n \"1\\n\",\n \"9\\n\",\n \"1\\n\",\n \"1\\n\",\n \"-1\\n\"\n ]\n}""", # A correct sample # noqa: E501 + """{\n \"inputs\": [\n \"5 7 6 11\\n3\\n5 3 8\\n6 7 11\\n5 2 5\\n\",\n \"3 4 3 10\\n3\\n3 1 4\\n4 5 9\\n3 10 10\\n\",\n \"1 1 2 10\\n2\\n1 1 3\\n2 6 10\\n\",\n \"9 8 7 8\\n9\\n10 6 6\\n10 6 6\\n7 7 8\\n9 5 6\\n8 9 9\\n9 5 5\\n9 8 8\\n8 5 6\\n9 10 10\\n\",\n \"6 15 7 15\\n9\\n6 15 15\\n7 14 14\\n6 15 15\\n9 14 14\\n7 14 16\\n6 15 15\\n6 15 15\\n7 14 14\\n8 15 15\\n\",\n \"13 16 20 10\\n18\\n13 16 16\\n20 10 10\\n19 10 10\\n12 15 15\\n20 10 10\\n18 11 11\\n19 10 10\\n19 10 10\\n20 10 10\\n19 10 10\\n20 10 10\\n20 10 10\\n19 10 10\\n18 11 11\\n13 16 16\\n12 15 15\\n19 10 10\\n19 10 10\\n\",\n \"89 29 88 30\\n16\\n87 31 31\\n14 95 95\\n98 88 89\\n96 88 88\\n14 97 97\\n13 97 98\\n100 88 88\\n88 32 32\\n99 88 89\\n90 29 29\\n87 31 31\\n15 94 96\\n89 29 29\\n88 32 32\\n97 89 89\\n88 29 30\\n\",\n \"30 14 39 19\\n31\\n35 7 11\\n37 11 12\\n32 13 13\\n37 5 6\\n46 13 13\\n37 14 14\\n31 13 13\\n43 13 19\\n45 15 19\\n46 13 13\\n32 17 17\\n41 14 19\\n30 14 14\\n43 13 17\\n34 16 18\\n44 11 19\\n38 13 13\\n40 12 20\\n37 16 18\\n46 16 18\\n34 10 14\\n36 9 10\\n36 15 19\\n38 15 19\\n42 13 19\\n33 14 15\\n35 15 19\\n33 17 18\\n39 12 20\\n36 5 7\\n45 12 12\\n\",\n \"2 1 1 1\\n2\\n1 1 2\\n2 1 2\\n\",\n \"1 1 1 2\\n5\\n1000000000 1 10000\\n19920401 1188 5566\\n1000000000 1 10000\\n1 1 10000\\n5 100 200\\n\",\n \"1 1 1000000000 2\\n5\\n1000000000 1 10000\\n19920401 1188 5566\\n1000000000 1 10000\\n1 1 10000\\n5 100 200\\n\"\n ],\n \"outputs\": [\n \"4\\n\",\n \"6\\n\",\n \"-1\\n\",\n \"-1\\n\",\n \"1\\n\",\n \"-1\\n\",\n \"1\\n\",\n \"9\\n\",\n \"1\\n\",\n \"1\\n\",\n \"-1\\n\"\n ]\n}""", # noqa: E501 +] # A failed sample with first several in-out passed + +prime_code_scores = [1.0, 0.9] + + +def test_parallelism(): + """ + Test if process pool works properly + """ + sequences_str = [] + ground_truth = [] + data_sources = [] + while len(sequences_str) < 32: + sequences_str.extend(prime_code_answers) + ground_truth.extend(prime_code_gts) + data_sources.extend(["codecontests"] * len(prime_code_answers)) + + sequences_str.extend(prime_math_answers) + ground_truth.extend(prime_math_gts) + data_sources.extend(["numina_aops_forum"] * len(prime_math_answers)) + + scores = asyncio.run( + parallel_compute_score_async(default_compute_score, sequences_str, ground_truth, data_sources, num_processes=16) + ) + print(scores) + + +@pytest.mark.skip("pyext not compatible with python 3.12") +def test_prime_code(): + """ + Test PRIME code sandbox. + """ + data_source = "codecontests" + for completion, ground_truth, score_ in zip(prime_code_answers, prime_code_gts, prime_code_scores, strict=True): + score = default_compute_score(data_source, completion, ground_truth) + assert float(score) == score_ + + +# Use the pytest.mark.skipif decorator to skip the test +@pytest.mark.skipif(not os.environ.get("SANDBOX_FUSION_URL"), reason="SANDBOX_FUSION_URL environment variable not set") +def test_prime_code_sandbox_fusion(): + """ + Test PRIME code on sandbox fusion. Skips if SANDBOX_FUSION_URL is not set. + """ + data_source = "codecontests" + # Get the URL from the environment variable, as skipif ensures it is set at this point + sandbox_fusion_url = os.environ.get("SANDBOX_FUSION_URL") + # Removed the previous 'if not sandbox_url' check block + + for completion, ground_truth, score_ in zip(prime_code_answers, prime_code_gts, prime_code_scores, strict=True): + score = default_compute_score( + data_source, completion, ground_truth, extra_info={"sandbox_fusion_url": sandbox_fusion_url} + ) # <-- Use the URL obtained from the environment variable + assert float(score) == score_ + + +@pytest.mark.skipif(not os.environ.get("SANDBOX_FUSION_URL"), reason="SANDBOX_FUSION_URL environment variable not set") +def test_continuous_score_consistency(): + """ + Verify that continuous score calculation is consistent between prime_code and sandbox_fusion. + Uses a test case where the first 9 out of 11 sub-cases pass (expected score 0.9). + """ + from verl.utils.reward_score import prime_code + + completion = prime_code_answers[1] # Use the second sample + ground_truth = prime_code_gts[1] # Use the second sample (9/11 pass, first 9 pass) + expected_continuous_score = 0.9 + + # 1. Calculate score using prime_code (default) with continuous=True + prime_score, _ = sandbox_fusion.compute_score( + os.environ.get("SANDBOX_FUSION_URL"), None, completion, ground_truth, continuous=True + ) + + # 2. Calculate score using sandbox_fusion with continuous=True + # Ensure the extra_info key triggers the sandbox_fusion path in default_compute_score + fusion_score, _ = prime_code.compute_score(completion, ground_truth, continuous=True) + + # 3. Assert scores are equal (using pytest.approx for float comparison) + assert float(prime_score) == pytest.approx(expected_continuous_score) + assert float(fusion_score) == pytest.approx(expected_continuous_score) + assert float(prime_score) == pytest.approx(float(fusion_score)) + print(f"Continuous Score (Prime Code): {prime_score}") + print(f"Continuous Score (Sandbox Fusion): {fusion_score}") + + +@pytest.mark.skip("pyext not compatible with python 3.12") +def test_check_correctness(): + from verl.utils.reward_score.prime_code import apps_check_correctness + + completion = prime_code_answers[0] + ground_truth = json.loads(prime_code_gts[0]) + ground_truth_single = {"inputs": ground_truth["inputs"][:1], "outputs": ground_truth["outputs"][:1]} + res, meta = apps_check_correctness(in_outs=ground_truth_single, generation=completion, timeout=5, debug=False) + print(res, meta) + + +def test_prime_math(): + data_source = "numina_aops_forum" + for completion, ground_truth in zip(prime_math_answers, prime_math_gts, strict=True): + score = default_compute_score(data_source, completion, ground_truth) + assert float(score) == 1.0 diff --git a/verl/tests/utils/test_activation_offload.py b/verl/tests/utils/test_activation_offload.py new file mode 100644 index 0000000000000000000000000000000000000000..c8a827cfab4eacfd89b5cf8f3b473c3a288673a3 --- /dev/null +++ b/verl/tests/utils/test_activation_offload.py @@ -0,0 +1,175 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os +import shutil +import tempfile + +import pytest +import torch +import torch.distributed +import torch.multiprocessing as mp +from torch.distributed import init_device_mesh +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +from torch.distributed.fsdp import MixedPrecision, ShardingStrategy +from transformers import AutoModelForCausalLM, AutoTokenizer, Qwen2Config + +from verl.utils.activation_offload import enable_activation_offloading +from verl.utils.checkpoint.fsdp_checkpoint_manager import FSDPCheckpointManager +from verl.utils.device import get_device_name, get_nccl_backend, get_torch_device +from verl.utils.fsdp_utils import MixedPrecisionPolicy, apply_fsdp2, get_fsdp_wrap_policy + + +def create_random_input_ids(batch_size, seq_len, vocab_size): + if get_device_name() == "cuda": + from flash_attn.bert_padding import unpad_input + elif get_device_name() == "npu": + from verl.utils.attention_utils import unpad_input + from verl.utils.model import compute_position_id_with_mask, create_random_mask + + input_ids = torch.randint(0, vocab_size, (batch_size, seq_len), device=get_device_name()) + + attention_mask = create_random_mask( + input_ids, max_ratio_of_left_padding=0.1, min_ratio_of_valid_token=0.5, max_ratio_of_valid_token=0.7 + ) + position_ids = compute_position_id_with_mask(attention_mask) + + input_ids = unpad_input(input_ids.unsqueeze(-1), attention_mask)[0].transpose(0, 1) + position_ids = unpad_input(position_ids.unsqueeze(-1), attention_mask)[0].transpose(0, 1) + return input_ids, position_ids + + +def _fsdp_activation_offloading_test(rank, world_size, rendezvous_file, strategy="fsdp"): + get_torch_device().set_device(rank) + torch.distributed.init_process_group( + backend=get_nccl_backend(), + init_method=f"file://{rendezvous_file}", + rank=rank, + world_size=world_size, + ) + device_mesh = init_device_mesh(get_device_name(), mesh_shape=(world_size,), mesh_dim_names=("dp",)) + + model_name = os.path.expanduser("~/models/Qwen/Qwen2.5-0.5B-Instruct") + config = Qwen2Config(num_hidden_layers=4) + + with torch.device(get_device_name()): + model = AutoModelForCausalLM.from_config( + config=config, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2" + ) + model = model.to(device=get_device_name()) + + # Wrap model with FSDP + mixed_precision = MixedPrecision(param_dtype=torch.bfloat16, reduce_dtype=torch.float32, buffer_dtype=torch.float32) + + if strategy == "fsdp": + model = FSDP( + model, + use_orig_params=False, + device_id=get_torch_device().current_device(), + sharding_strategy=ShardingStrategy.FULL_SHARD, + mixed_precision=mixed_precision, + device_mesh=device_mesh, + auto_wrap_policy=get_fsdp_wrap_policy(module=model), + ) + else: + mp_policy = MixedPrecisionPolicy( + param_dtype=torch.bfloat16, reduce_dtype=torch.float32, cast_forward_inputs=True + ) + fsdp_kwargs = { + "mesh": device_mesh, + "mp_policy": mp_policy, + } + apply_fsdp2(model, fsdp_kwargs, {}) + + optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4) + lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.9) + + # Create checkpoint manager + tokenizer = AutoTokenizer.from_pretrained(model_name) + checkpoint_manager = FSDPCheckpointManager( + model=model, optimizer=optimizer, lr_scheduler=lr_scheduler, tokenizer=tokenizer + ) + + # Generate sample input + batch_size = 2 + seq_len = 32 + vocab_size = 32000 + # First input for initial update + input_ids1, position_ids1 = create_random_input_ids(batch_size, seq_len, vocab_size) + + # Second input for verification + input_ids2, position_ids2 = create_random_input_ids(batch_size, seq_len, vocab_size) + + # Step 1: Initial update and save checkpoint + outputs1 = model(input_ids=input_ids1, position_ids=position_ids1) + loss1 = outputs1.logits.mean() + loss1.backward() + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + # Save checkpoint after first update + temp_dir = tempfile.mkdtemp() + checkpoint_path = os.path.join(temp_dir, "checkpoint") + checkpoint_manager.save_checkpoint(local_path=checkpoint_path, hdfs_path=None, global_step=0) + + # Step 2: Second update and forward pass + outputs2 = model(input_ids=input_ids2, position_ids=position_ids2) + loss2 = outputs2.logits.mean() + loss2.backward() + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + # Record logits after second update + with torch.no_grad(): + logits_without_offloading = model(input_ids=input_ids2, position_ids=position_ids2).logits + + # Step 3: wrap module with activation offloading and load checkpoint + enable_activation_offloading(model, strategy=strategy) + checkpoint_manager.load_checkpoint(checkpoint_path) + + # Step 4: Repeat the second update with same input + outputs3 = model(input_ids=input_ids2, position_ids=position_ids2) + loss3 = outputs3.logits.mean() + loss3.backward() + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + # Record logits after loaded checkpoint and update + with torch.no_grad(): + logits_with_offloading = model(input_ids=input_ids2, position_ids=position_ids2).logits + + # Step 4: Verify outputs match + torch.testing.assert_close(logits_without_offloading, logits_with_offloading, atol=0.0, rtol=0.0) + print(f"Activaiton offloading for {strategy} test passed on {world_size} GPUs!") + + # Cleanup + shutil.rmtree(temp_dir) + torch.distributed.barrier() + torch.distributed.destroy_process_group() + + +@pytest.mark.parametrize("world_size", (2, 4)) +@pytest.mark.parametrize("strategy", ("fsdp", "fsdp2")) +def test_activation_offloading(world_size, strategy, tmp_path): + rendezvous_file = str(tmp_path / "rdzv_file") + os.makedirs(os.path.dirname(rendezvous_file), exist_ok=True) + + mp.spawn( + fn=_fsdp_activation_offloading_test, + args=(world_size, rendezvous_file, strategy), + nprocs=world_size, + join=True, + ) diff --git a/verl/tests/utils/test_audio_input_support_on_cpu.py b/verl/tests/utils/test_audio_input_support_on_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..202474f0b40c8246bd8bc855c72ddede1b8ab474 --- /dev/null +++ b/verl/tests/utils/test_audio_input_support_on_cpu.py @@ -0,0 +1,112 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import types + +import pytest +import torch + +from verl.utils.model import extract_multi_modal_inputs +from verl.utils.tokenizer import build_multimodal_processor_inputs + + +def test_build_messages_replaces_audio_placeholder() -> None: + pytest.importorskip("datasets") + from verl.utils.dataset.rl_dataset import RLHFDataset + + dataset = RLHFDataset.__new__(RLHFDataset) + dataset.prompt_key = "prompt" + dataset.image_key = "images" + dataset.video_key = "videos" + dataset.audio_key = "audios" + dataset.processor = object() + + example = { + "prompt": [ + {"role": "user", "content": "Listen to this: + score = score / 4 + return score + return score + else: + return format_score + + +def compute_score_subem(solution_str, ground_truth, method="strict", format_score=0.0, score=1.0): + """The scoring function for substring exact match (EM). + + Args: + solution_str: the solution text + ground_truth: the ground truth + method: the method to extract the solution, choices are 'strict' and 'flexible' + format_score: the score for the format + score: the score for the correct answer + """ + answer = extract_solution(solution_str=solution_str) + do_print = random.randint(1, 64) == 1 + + if do_print: + print("--------------------------------") + print(f"Golden answers: {ground_truth['target']}") + print(f"Extracted answer: {answer}") + print(f"Solution string: {solution_str}") + + if answer is None: + return 0 + else: + if subem_check(answer, ground_truth["target"]): + return score + else: + return format_score diff --git a/verl/verl/utils/rollout_skip.py b/verl/verl/utils/rollout_skip.py new file mode 100644 index 0000000000000000000000000000000000000000..04414265ea7b09af954305af4f880b06e7ca1665 --- /dev/null +++ b/verl/verl/utils/rollout_skip.py @@ -0,0 +1,453 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import json +from enum import Enum +from pathlib import Path +from typing import Any, Callable + +from verl.protocol import DataProto +from verl.workers.config.rollout import RolloutConfig + + +def _get_skip_attr(skip_config, key: str, default): + """Get attribute from skip config, supporting both dict and SkipConfig dataclass.""" + if isinstance(skip_config, dict): + return skip_config.get(key, default) + return getattr(skip_config, key, default) + + +def _find_last_gen_step_for_train_step(step_file: Path, target_train_step: int) -> tuple[int, int] | None: + """ + Find the last `(train_step, gen_step)` pair for a given train_step without loading the + entire file into memory. + + This scans the file line-by-line (O(n) time, O(1) memory) and keeps the last match. + It also stops early once `train_step` exceeds `target_train_step` (assuming chronological logs). + """ + step_file = Path(step_file) + if not step_file.is_file(): + return None + + last_match: tuple[int, int] | None = None + with step_file.open("r", encoding="utf-8", errors="ignore") as f: + for raw in f: + line = raw.strip() + if not line: + continue + parts = line.split() + if len(parts) < 2: + continue + try: + train_step = int(parts[0]) + gen_step = int(parts[1]) + except Exception: + continue + + if train_step < target_train_step: + continue + if train_step == target_train_step: + last_match = (train_step, gen_step) + continue + # train_step > target_train_step: no more matches expected + break + + return last_match + + +class SkipAction(Enum): + CACHE = "cache" # cache the sample. If dump_date is found, use it. If not found, dump it. + REPEAT = "repeat" # Repeat the sample when gen_step reach skip.max_dump_step + REPEAT_LAST = "repeat_last" # Repeat the last sample when gen_step reach skip.max_dump_step + + +class RolloutSkip: + """ + RolloutSkip skips sequence generation during rollout by attempting to load previously dumped data. + If no dumped data is found, it generates new sequences and saves them to disk. + + Args: + config: The configuration object containing rollout settings. + rollout_wg: The worker group that handles the rollout process. + + Note: + Whenever any of the following parameters differ from previous runs—trainer.experiment_name, + trainer.project_name, rollout.n, or rollout.gen_batch_size—new sequences will be generated + and saved under different filenames. + + + """ + + print_mark = "[RolloutSkip()] " + + def __init__(self, config, rollout_wg) -> None: + self.rollout_config: RolloutConfig = config.actor_rollout_ref.rollout + self.skip_config = self.rollout_config.skip + self.is_enable = _get_skip_attr(self.skip_config, "enable", False) + self._rollout_wg = rollout_wg + + if not self.is_enable: + return + + self.exp_name = config.trainer.get("experiment_name", "") + self.project_name = config.trainer.get("project_name", "") + self.n = int(getattr(self.rollout_config, "n", 0)) + self.gbs = int(config.data.get("gen_batch_size", config.data.get("train_batch_size", 0))) + self.response_length = config.data.get("max_response_length", 0) + self.prompt_length = config.data.get("max_prompt_length", 0) + + self._new_batch = None + self.curr_gen_step: int = 0 # mark the index of rollout result, start from 1 + self.curr_train_step: int = 0 + + self.record_global_steps = None # Given from xxx_ray_tainer.py, start from 1 + self.record_gen_steps = None # Given from xxx_ray_tainer.py, start from 1 + self.__gen_offset_step = 0 + + self.max_dump_step = max(0, _get_skip_attr(self.skip_config, "max_dump_step", 1)) # at least dump once + self.action = _get_skip_attr(self.skip_config, "action", SkipAction.REPEAT) + self.action = SkipAction(self.action) + + if self.max_dump_step <= 0: + assert self.action in [SkipAction.CACHE] + + self._create_dump_path() + self._flag_record = False + self.list_dumped_steps = [] + + @property + def is_active(self) -> bool: + """Whether RolloutSkip is enabled and has a rollout worker group.""" + return self.is_enable and self._rollout_wg is not None + + @property + def is_dump_step(self) -> bool: + """ + Determine if the current step is a dump step based on the configured dump interval. + If train_step is given, it follows the train_step, otherwise it follows the gen_step. + """ + return self.is_active and self.curr_train_step <= self.max_dump_step + + @property + def num_dumped_step(self) -> int: + return len(self.list_dumped_steps) + + def _get_path_dump(self, gen_step: int | None = None) -> Path: + """Return the directory path for a given gen_step (one dir per step, no .pkl).""" + if gen_step is None: + gen_step = self.curr_gen_step + return self.specify_dumped_dir.joinpath(f"genstep_{gen_step:06d}").absolute() + + def _get_path_step_record(self) -> Path: + return self.specify_dumped_dir.joinpath("train_step__gen_step.txt").absolute() + + def step(self) -> None: + if self.record_global_steps is None: + self.curr_train_step += 1 + else: + self.curr_train_step = self.record_global_steps + + if self.record_gen_steps is None: + self.curr_gen_step = self.curr_train_step + else: + self.curr_gen_step = self.record_gen_steps + + def _create_dump_path(self) -> None: + """ + Create the directory for dumping rollout data if it doesn't exist. + Warn if the directory is within Ray's temporary session directory. + Relative dump_dir is resolved against cwd; use an absolute path under Ray/multi-process. + """ + + raw = _get_skip_attr(self.skip_config, "dump_dir", "~/.verl/rollout_dump") + dumped_dir = Path(raw).expanduser().resolve() + sub_dir = ( + f"{self.exp_name}_{self.project_name}" + + f"/GBS{self.gbs}_N{self.n}_in{self.prompt_length}_out{self.response_length}" + ) + + self.specify_dumped_dir = dumped_dir.joinpath(sub_dir) + self.specify_dumped_dir.mkdir(parents=True, exist_ok=True) + + tmp_ray = "/tmp/ray/session" + + # Check if path is in Ray temporary directory + if str(self.specify_dumped_dir.absolute()).startswith(tmp_ray): + print( + f"{self.print_mark}\033[33mWarning: \nUsing dump path ", + f"'{self.specify_dumped_dir.absolute()}' is not recommended ", + f"as it's located in {tmp_ray}*\033[0m", + flush=True, + ) + print( + f"{self.print_mark}Rollout skip dump path set to: ", + str(self.specify_dumped_dir.absolute()), + flush=True, + ) + + def record( + self, + new_batch: DataProto, + global_steps: int | None = None, + gen_steps: int | None = None, + *args: Any, + **kwargs: Any, + ) -> None: + """Record the current training step based on the new batch. + + Args: + new_batch (DataProto): The new batch of data being processed. + """ + if self._rollout_wg is None: + return + if self._flag_record is False: + # make sure one record only corresponds to one skip + self._flag_record = True + self._new_batch = new_batch + else: + print( + f"{self.print_mark}Warning, duplicate record new_batch, " + "it was not a problem if acc/reward is not cared.", + flush=True, + ) + + if gen_steps is None: + gen_steps = global_steps + + # Check if train_step not start from 1 + if global_steps is not None: + if self.record_global_steps is None and global_steps > 1: + print(f"{self.print_mark}\033[32mResume Mode.\033[0m", flush=True) + last_train_step = global_steps - 1 # default when step file missing + last_gen_step = 0 + try: + found = _find_last_gen_step_for_train_step( + self._get_path_step_record(), + target_train_step=global_steps - 1, + ) + if found is not None: + last_train_step, last_gen_step = found + if last_train_step + 1 != global_steps: + print(f"{self.print_mark}\033[31mWarning: Train step not continues.\033[0m") + self.__gen_offset_step = last_gen_step + except Exception as e: + print( + f"{self.print_mark}\033[31mFailed to read step describe file. {e.__repr__()}\033[0m", + flush=True, + ) + print( + f"{self.print_mark}\033[32mResume from train_step: {last_train_step}, " + f"gen_step: {last_gen_step}.\033[0m", + flush=True, + ) + + if global_steps is not None: + self.record_global_steps = global_steps + if gen_steps is not None: + #! it is not right since dapo_trainer reset `gen_steps` when resume + self.record_gen_steps = gen_steps + self.__gen_offset_step + + def wrap_generate_sequences(self) -> None: + # if self.is_enable: + # self._rollout_wg = rollout_wg + + try: + self._rollout_wg.generate_sequences = wrap_generate_sequences(self, self._rollout_wg) + print( + f"{self.print_mark}\033[32mSuccessfully patched `actor_rollout_wg.generate_sequences()`.\033[0m", + flush=True, + ) + except Exception as e: + raise RuntimeError( + f"{self.print_mark}\033[31mFailed to patch `actor_rollout_wg.generate_sequences()`.\033[0m", + flush=True, + ) from e + + def try_load(self, step: int | None = None) -> tuple[DataProto | None, DataProto | None]: + dumped_gen_batch = None + dumped_new_batch = None + if step is None: + step = self.curr_gen_step + + step_dir = self._get_path_dump(step) + if not step_dir.exists() or not step_dir.is_dir(): + print( + f"{self.print_mark}\033[33mNo dumped data found at gen_step {step} " + f"from {step_dir}. The trainer will generate and dump the data for this gen_step.\033[0m", + flush=True, + ) + return dumped_new_batch, dumped_gen_batch + + new_batch_path = step_dir / "new_batch.dp" + gen_batch_path = step_dir / "gen_batch.dp" + if not (new_batch_path.is_file() and gen_batch_path.is_file()): + print( + f"{self.print_mark}\033[33mNo dumped data found at gen_step {step} " + f"(missing new_batch.dp or gen_batch.dp in {step_dir}).\033[0m", + flush=True, + ) + return dumped_new_batch, dumped_gen_batch + + try: + dumped_new_batch = DataProto.load_from_disk(new_batch_path) + dumped_gen_batch = DataProto.load_from_disk(gen_batch_path) + print( + f"{self.print_mark}\033[32mSuccessfully load pre-generated data from {step_dir}.\033[0m", + flush=True, + ) + if step not in self.list_dumped_steps: + self.list_dumped_steps.append(step) + except Exception as e: + print( + f"{self.print_mark}\033[31mFailed to load pre-generated data from {step_dir}: {e}\033[0m", + flush=True, + ) + + return dumped_new_batch, dumped_gen_batch + + def dump(self, outputs: DataProto) -> None: + if self._flag_record is False or self._new_batch is None: + raise AssertionError( + f"{self.print_mark}\033[33mError: \n" + + "The new_batch record is required." + + "Please record the new_batch using `RolloutSkip.record(new_batch)` in trainer.fit().\033[0m" + ) + self._flag_record = False + + train_step = self.record_global_steps if self.record_global_steps is not None else self.curr_train_step + gen_step = self.record_gen_steps if self.record_gen_steps is not None else self.curr_gen_step + step_dir = self._get_path_dump(gen_step) + step_dir.mkdir(parents=True, exist_ok=True) + + try: + self._new_batch.save_to_disk(step_dir / "new_batch.dp") + outputs.save_to_disk(step_dir / "gen_batch.dp") + meta_path = step_dir / "meta.json" + meta_path.write_text(json.dumps({"global_steps": train_step, "gen_steps": gen_step})) + + with open(str(self._get_path_step_record()), "a") as f: + f.write(f"{train_step} {gen_step}\n") + + print( + f"{self.print_mark}\033[32mSuccessfully dump data in {step_dir}\033[0m", + flush=True, + ) + if self.curr_gen_step not in self.list_dumped_steps: + self.list_dumped_steps.append(self.curr_gen_step) + + except Exception as e: + print( + f"{self.print_mark}\033[31mFailed to dump data in {step_dir}: {e}\033[0m", + flush=True, + ) + + def replace_curr_new_batch(self, dumped_new_batch: DataProto) -> None: + """Replace the current new_batch's content with that from the dumped_new_batch. + In case of [Answer] mismatch. + """ + + if self._flag_record is False: + raise AssertionError( + f"{self.print_mark}\033[33mError: \n" + + "The new_batch is not recorded. Please record the new_batch" + + "using `RolloutSkip.record(new_batch)`. \033[0m" + ) + self._flag_record = False + + self._new_batch.batch = dumped_new_batch.batch + self._new_batch.non_tensor_batch = dumped_new_batch.non_tensor_batch + self._new_batch.meta_info = dumped_new_batch.meta_info + + +def wrap_generate_sequences(rolloutskip: RolloutSkip, rollout_wg: Any) -> Callable[..., DataProto]: + generate_sequences = rollout_wg.generate_sequences + + def rollout_skip_wrap_fn(batch: DataProto, **kwargs: Any) -> DataProto: + rolloutskip.step() + # Record input batch as new_batch so dump() / replace_curr_new_batch() have it + rolloutskip.record(batch) + return_batch = None + + if rolloutskip.is_dump_step: + # * try load + dumped_new_batch, return_batch = rolloutskip.try_load() + + if return_batch is None: + # 1. Generation + return_batch = generate_sequences(batch, **kwargs) + # 2. Dump + rolloutskip.dump(return_batch) + else: + rolloutskip.replace_curr_new_batch(dumped_new_batch) + + elif rolloutskip.action == SkipAction.CACHE: + return_batch = generate_sequences(batch, **kwargs) + + elif rolloutskip.action == SkipAction.REPEAT: + if rolloutskip.num_dumped_step == 0: + return_batch = generate_sequences(batch, **kwargs) + rolloutskip.dump(return_batch) + else: + target_step = rolloutskip.list_dumped_steps[ + (rolloutskip.curr_gen_step - 1) % rolloutskip.num_dumped_step + ] + dumped_new_batch, return_batch = rolloutskip.try_load(step=target_step) + if return_batch is None: + return_batch = generate_sequences(batch, **kwargs) + rolloutskip.dump(return_batch) + else: + rolloutskip.replace_curr_new_batch(dumped_new_batch) + + elif rolloutskip.action == SkipAction.REPEAT_LAST: + target_step = rolloutskip.list_dumped_steps[-1] + dumped_new_batch, return_batch = rolloutskip.try_load(step=target_step) + if return_batch is None: + return_batch = generate_sequences(batch, **kwargs) + rolloutskip.dump(return_batch) + else: + rolloutskip.replace_curr_new_batch(dumped_new_batch) + + # clean + return return_batch + + return rollout_skip_wrap_fn + + +def read_dumped_data(path_dump: Path | str) -> dict[str, DataProto]: + """ + Read dumped rollout data from a step directory (DataProto.save_to_disk format). + + path_dump should point to a step directory containing new_batch.dp and gen_batch.dp, + e.g. .../GBS8_N16_in1024_out10240/genstep_000001/ + + ``` + from verl.utils.rollout_skip import read_dumped_data + + dumped_data = read_dumped_data("path/to/rollout_dump/.../genstep_000001") + print(dumped_data["new_batch"]) + print(dumped_data["gen_batch"]) + ``` + """ + path_dump = Path(path_dump) + if not path_dump.is_dir(): + raise FileNotFoundError(f"Directory {path_dump} does not exist.") + + new_batch_path = path_dump / "new_batch.dp" + gen_batch_path = path_dump / "gen_batch.dp" + if not (new_batch_path.is_file() and gen_batch_path.is_file()): + raise FileNotFoundError(f"Missing new_batch.dp or gen_batch.dp under {path_dump}.") + + return { + "new_batch": DataProto.load_from_disk(new_batch_path), + "gen_batch": DataProto.load_from_disk(gen_batch_path), + } diff --git a/verl/verl/utils/rollout_trace.py b/verl/verl/utils/rollout_trace.py new file mode 100644 index 0000000000000000000000000000000000000000..45a3f3461017dff121d8545bff2725141ba4e57a --- /dev/null +++ b/verl/verl/utils/rollout_trace.py @@ -0,0 +1,291 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import contextlib +import functools +import inspect +import os +from contextvars import ContextVar +from typing import Optional + +from pydantic import BaseModel + +from verl.utils.ray_utils import get_event_loop + +_trace_enabled: ContextVar[bool] = ContextVar("_trace_enabled", default=True) + + +class RolloutTraceConfig: + """Configuration for rollout tracing with various backends. + + Singleton configuration class for managing rollout trace settings across different + tracing backends like Weave and MLflow. + + Args: + backend (Optional[str]): Tracing backend to use ('weave', 'mlflow', or None). + client (Optional[object]): Client instance for the selected backend. + token2text (bool): Whether to convert tokens to text in traces. Defaults to False. + project_name (str): Name of the project for tracing. + experiment_name (str): Name of the experiment for tracing. + max_samples_per_step_per_worker (Optional[int]): Maximum number of unique samples to trace + per worker per step. If None, all samples are traced. If set, each worker will randomly + select up to this many unique samples to trace (including all their rollouts for GRPO). + Total traces = max_samples_per_step_per_worker * num_workers * n_rollouts_per_sample. + """ + + _instance: Optional["RolloutTraceConfig"] = None + backend: Optional[str] = None + client: Optional[object] = None + token2text: bool = False + _initialized: bool = False + project_name: str = None + experiment_name: str = None + max_samples_per_step_per_worker: Optional[int] = None + + def __new__(cls, *args, **kwargs): + if cls._instance is None: + cls._instance = super().__new__(cls) + cls._instance._initialized = False + return cls._instance + + @classmethod + def get_instance(cls) -> "RolloutTraceConfig": + if cls._instance is None: + cls._instance = cls() + return cls._instance + + @classmethod + def init( + cls, + project_name: str, + experiment_name: str, + backend: str, + token2text: bool = False, + max_samples_per_step_per_worker: Optional[int] = None, + ): + config = cls.get_instance() + if config._initialized: + return + + config.backend = backend + config.token2text = token2text + config.project_name = project_name + config.experiment_name = experiment_name + config.max_samples_per_step_per_worker = max_samples_per_step_per_worker + + if backend == "weave": + import weave + + config.client = weave.init(project_name) + elif backend == "mlflow": + import mlflow + + mlflow.config.enable_async_logging() + config.client = mlflow + + MLFLOW_TRACKING_URI = os.environ.get("MLFLOW_TRACKING_URI", "sqlite:////tmp/mlruns.db") + mlflow.set_tracking_uri(MLFLOW_TRACKING_URI) + + mlflow.set_experiment(project_name) + else: + config.client = None + + config._initialized = True + + @classmethod + def get_backend(cls) -> Optional[str]: + return cls.get_instance().backend + + @classmethod + def get_client(cls) -> Optional[object]: + return cls.get_instance().client + + @classmethod + def enable_token2text(cls) -> Optional[bool]: + return cls.get_instance().token2text + + @classmethod + def reset(cls): + cls._instance = None + + +@contextlib.contextmanager +def rollout_trace_attr( + sample_index=None, step=None, rollout_n=None, name="rollout_trace", validate=False, trace: bool = True +): + """A context manager to add attributes to a trace for the configured backend. + + Args: + sample_index: Sample index for the trace. + step: Training step number. + rollout_n: Rollout number (for GRPO with multiple rollouts per sample). + name: Name for the trace span (used by mlflow backend). + validate: Whether this is a validation run. + trace: If False, disables tracing for the duration of the context. + """ + backend = RolloutTraceConfig.get_backend() + + should_skip = backend is not None and not trace + + if should_skip: + token = _trace_enabled.set(False) + try: + yield + finally: + _trace_enabled.reset(token) + return + + # Build attributes for the trace + attributes = {} + if backend: + if sample_index is not None: + attributes["sample_index"] = sample_index + if step is not None: + attributes["step"] = step + if rollout_n is not None: + attributes["rollout_n"] = rollout_n + attributes["validate"] = validate + attributes["experiment_name"] = RolloutTraceConfig.get_instance().experiment_name + + if not attributes or backend is None: + yield + return + + if backend == "weave": + import weave + + with weave.attributes(attributes): + yield + elif backend == "mlflow": + import mlflow + + with mlflow.start_span(name=name) as span: + trace_id = span.trace_id + for key, value in attributes.items(): + mlflow.set_trace_tag(trace_id, str(key), str(value)) + yield + else: + yield + + +def rollout_trace_op(func): + @functools.wraps(func) + async def async_wrapper(self, *args, **kwargs): + if not _trace_enabled.get(): + return await func(self, *args, **kwargs) + + backend = RolloutTraceConfig.get_backend() + enable_token2text = RolloutTraceConfig.enable_token2text() + if backend is None: + return await func(self, *args, **kwargs) + + sig = inspect.signature(func) + bound_args = sig.bind(self, *args, **kwargs) + bound_args.apply_defaults() + inputs = dict(bound_args.arguments) + del inputs["self"] + + async def add_token2text(self, result): + if hasattr(result, "prompt_ids") and hasattr(self, "tokenizer") and hasattr(self.tokenizer, "decode"): + # Use model_dump() for Pydantic models to get a proper copy, + # otherwise vars() returns a reference to internal __dict__ which + # can cause serialization issues with MLflow + if isinstance(result, BaseModel): + _result = result.model_dump() + else: + _result = dict(vars(result)) + loop = get_event_loop() + if hasattr(result, "prompt_ids"): + prompt_text = await loop.run_in_executor(None, self.tokenizer.decode, result.prompt_ids) + _result["prompt_text"] = prompt_text + + if hasattr(result, "response_ids"): + response_text = await loop.run_in_executor(None, self.tokenizer.decode, result.response_ids) + _result["response_text"] = response_text + return _result + return result + + if backend == "weave": + tracer = RolloutTraceConfig.get_client() + from weave.trace.context import call_context + + cur_attributes = {**call_context.call_attributes.get()} + call = tracer.create_call(op=func.__qualname__, inputs=inputs, attributes=cur_attributes) + try: + result = await func(self, *args, **kwargs) + + if enable_token2text: + _result = await add_token2text(self, result) + tracer.finish_call(call, output=_result) + else: + tracer.finish_call(call, output=result) + + return result + + except Exception as e: + tracer.finish_call(call, exception=e) + raise e + elif backend == "mlflow": + import mlflow + + with mlflow.start_span(name=func.__qualname__) as span: + span.set_inputs(inputs) + result = await func(self, *args, **kwargs) + if enable_token2text: + _result = await add_token2text(self, result) + span.set_outputs(_result) + else: + span.set_outputs(result) + + return result + + else: + return await func(self, *args, **kwargs) + + @functools.wraps(func) + def wrapper(self, *args, **kwargs): + if not _trace_enabled.get(): + return func(self, *args, **kwargs) + + backend = RolloutTraceConfig.get_backend() + if backend is None: + return func(self, *args, **kwargs) + + sig = inspect.signature(func) + bound_args = sig.bind(self, *args, **kwargs) + bound_args.apply_defaults() + inputs = dict(bound_args.arguments) + del inputs["self"] + + if backend == "weave": + tracer = RolloutTraceConfig.get_client() + from weave.trace.context import call_context + + cur_attributes = {**call_context.call_attributes.get()} + call = tracer.create_call(op=func.__qualname__, inputs=inputs, attributes=cur_attributes) + try: + result = func(self, *args, **kwargs) + tracer.finish_call(call, output=result) + return result + except Exception as e: + tracer.finish_call(call, exception=e) + raise e + elif backend == "mlflow": + import mlflow + + return mlflow.trace(func)(self, *args, **kwargs) + else: + return func(self, *args, **kwargs) + + return async_wrapper if inspect.iscoroutinefunction(func) else wrapper diff --git a/verl/verl/utils/seqlen_balancing.py b/verl/verl/utils/seqlen_balancing.py new file mode 100644 index 0000000000000000000000000000000000000000..4e172722a896ced6296143ed3d428bc0eaaa6187 --- /dev/null +++ b/verl/verl/utils/seqlen_balancing.py @@ -0,0 +1,626 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy +import heapq +from itertools import chain + +import torch +from torch import distributed as dist + +from verl.protocol import DataProto +from verl.utils import tensordict_utils as tu +from verl.utils.device import get_device_name + + +def calculate_workload(seqlen_list: torch.Tensor) -> torch.Tensor: + """Calculate approximate computational workload for transformer attention. + + Estimates FLOPs for dense transformer blocks based on sequence length using + the formula: FLOPs ≈ 12 * hidden_size² * seqlen + 2 * hidden_size * seqlen² + + The constants are calibrated for a 7B model (hidden_size=4096), yielding: + workload ∝ 24576 * seqlen + seqlen² + + Args: + seqlen_list: Sequence lengths as a tensor. + + Returns: + torch.Tensor: Estimated workload values proportional to actual FLOPs. + + Note: + The returned values are relative workloads, not actual FLOP counts. + Useful for balancing computation across data parallel ranks. + """ + return 24576 * seqlen_list + seqlen_list**2 + + +def karmarkar_karp(seqlen_list: list[int], k_partitions: int, equal_size: bool) -> list[list[int]]: + """Partition items into k groups using the Karmarkar-Karp differencing method. + + Implements the Largest Differencing Method (LDM) algorithm for balanced + multi-way number partitioning. This heuristic produces near-optimal partitions + by iteratively combining the sets with the largest difference. + + Args: + seqlen_list: Values to partition (typically sequence lengths or workloads). + k_partitions: Number of partitions to create. + equal_size: If True, each partition will have exactly len(seqlen_list) / k_partitions + items. If False, partitions may have different sizes. + + Returns: + list[list[int]]: List of k partitions, each containing indices into seqlen_list. + + See Also: + https://en.wikipedia.org/wiki/Largest_differencing_method + + Note: + When equal_size=True, len(seqlen_list) must be divisible by k_partitions. + """ + + # see: https://en.wikipedia.org/wiki/Largest_differencing_method + class Set: + def __init__(self) -> None: + self.sum = 0 + self.items = [] + + def add(self, idx: int, val: int): + self.items.append((idx, val)) + self.sum += val + + def merge(self, other): + for idx, val in other.items: + self.items.append((idx, val)) + self.sum += val + + def __lt__(self, other): + if self.sum != other.sum: + return self.sum < other.sum + if len(self.items) != len(other.items): + return len(self.items) < len(other.items) + return self.items < other.items + + class State: + def __init__(self, items: list[tuple[int, int]], k: int) -> None: + self.k = k + # sets should always be decreasing order + self.sets = [Set() for _ in range(k)] + assert len(items) in [1, k], f"{len(items)} not in [1, {k}]" + for i, (idx, seqlen) in enumerate(items): + self.sets[i].add(idx=idx, val=seqlen) + self.sets = sorted(self.sets, reverse=True) + + def get_partitions(self): + partitions = [] + for i in range(len(self.sets)): + cur_partition = [] + for idx, _ in self.sets[i].items: + cur_partition.append(idx) + partitions.append(cur_partition) + return partitions + + def merge(self, other): + for i in range(self.k): + self.sets[i].merge(other.sets[self.k - 1 - i]) + self.sets = sorted(self.sets, reverse=True) + + @property + def spread(self) -> int: + return self.sets[0].sum - self.sets[-1].sum + + def __lt__(self, other): + # least heap, let the state with largest spread to be popped first, + # if the spread is the same, let the state who has the largest set + # to be popped first. + if self.spread != other.spread: + return self.spread > other.spread + return self.sets[0] > other.sets[0] + + def __repr__(self) -> str: + repr_str = "[" + for i in range(self.k): + if i > 0: + repr_str += "," + repr_str += "{" + for j, (_, seqlen) in enumerate(self.sets[i].items): + if j > 0: + repr_str += "," + repr_str += str(seqlen) + repr_str += "}" + repr_str += "]" + return repr_str + + sorted_seqlen_list = sorted([(seqlen, i) for i, seqlen in enumerate(seqlen_list)]) + states_pq = [] + if equal_size: + assert len(seqlen_list) % k_partitions == 0, f"{len(seqlen_list)} % {k_partitions} != 0" + for offset in range(0, len(sorted_seqlen_list), k_partitions): + items = [] + for i in range(k_partitions): + seqlen, idx = sorted_seqlen_list[offset + i] + items.append((idx, seqlen)) + heapq.heappush(states_pq, State(items=items, k=k_partitions)) + else: + for seqlen, idx in sorted_seqlen_list: + heapq.heappush(states_pq, State(items=[(idx, seqlen)], k=k_partitions)) + + while len(states_pq) > 1: + state0 = heapq.heappop(states_pq) + state1 = heapq.heappop(states_pq) + # merge states + state0.merge(state1) + heapq.heappush(states_pq, state0) + + final_state = states_pq[0] + partitions = final_state.get_partitions() + if equal_size: + for i, partition in enumerate(partitions): + assert len(partition) * k_partitions == len(seqlen_list), ( + f"{len(partition)} * {k_partitions} != {len(seqlen_list)}" + ) + return partitions + + +def greedy_partition(seqlen_list: list[int], k_partitions: int, equal_size: bool) -> list[list[int]]: + """Partition items into k groups using a greedy assignment strategy. + + Assigns each item to the partition with the smallest current sum, iterating + through items in order. Simpler but typically less optimal than Karmarkar-Karp. + + Args: + seqlen_list: Values to partition (typically sequence lengths or workloads). + k_partitions: Number of partitions to create. + equal_size: If True, adds a bias to ensure equal partition sizes. + Requires len(seqlen_list) to be divisible by k_partitions. + + Returns: + list[list[int]]: List of k partitions, each containing indices into seqlen_list. + + Note: + When equal_size=True, a large bias is added to encourage equal distribution + of items before considering the actual values. + """ + bias = sum(seqlen_list) + 1 if equal_size else 0 + sorted_seqlen = [(seqlen + bias, i) for i, seqlen in enumerate(seqlen_list)] + partitions = [[] for _ in range(k_partitions)] + partition_sums = [0 for _ in range(k_partitions)] + for seqlen, i in sorted_seqlen: + min_idx = None + for j in range(k_partitions): + if min_idx is None or partition_sums[j] < partition_sums[min_idx]: + min_idx = j + partitions[min_idx].append(i) + partition_sums[min_idx] += seqlen + if equal_size: + for i, partition in enumerate(partitions): + assert len(partition) * k_partitions == len(seqlen_list), ( + f"{len(partition)} * {k_partitions} != {len(seqlen_list)}" + ) + return partitions + + +def get_seqlen_balanced_partitions(seqlen_list: list[int], k_partitions: int, equal_size: bool): + """ + Calculates partitions of indices from seqlen_list such that the sum of sequence lengths + in each partition is balanced. Uses the Karmarkar-Karp differencing method. + + This is useful for balancing workload across devices or batches, especially when + dealing with variable sequence lengths. + + Args: + seqlen_list (List[int]): A list of sequence lengths for each item. + k_partitions (int): The desired number of partitions. + equal_size (bool): If True, ensures that each partition has the same number of items. + Requires len(seqlen_list) to be divisible by k_partitions. + If False, partitions can have varying numbers of items, focusing + only on balancing the sum of sequence lengths. + + Returns: + List[List[int]]: A list containing k_partitions lists. Each inner list contains the + original indices of the items assigned to that partition. The indices + within each partition list are sorted. + + Raises: + AssertionError: If len(seqlen_list) < k_partitions. + AssertionError: If equal_size is True and len(seqlen_list) is not divisible by k_partitions. + AssertionError: If any resulting partition is empty. + """ + assert len(seqlen_list) >= k_partitions, f"number of items:[{len(seqlen_list)}] < k_partitions:[{k_partitions}]" + + def _check_and_sort_partitions(partitions): + assert len(partitions) == k_partitions, f"{len(partitions)} != {k_partitions}" + seen_idx = set() + sorted_partitions = [None] * k_partitions + for i, partition in enumerate(partitions): + assert len(partition) > 0, f"the {i}-th partition is empty" + for idx in partition: + seen_idx.add(idx) + sorted_partitions[i] = sorted(partition) + assert seen_idx == set(range(len(seqlen_list))) + return sorted_partitions + + partitions = karmarkar_karp(seqlen_list=seqlen_list, k_partitions=k_partitions, equal_size=equal_size) + return _check_and_sort_partitions(partitions) + + +def log_seqlen_unbalance(seqlen_list: list[int], partitions: list[list[int]], prefix): + """ + Calculate and log metrics related to sequence length imbalance before and after partitioning. + + Args: + seqlen_list (List[int]): A list of sequence lengths for each item. + partitions (List[List[int]]): A list of partitions, where each inner list contains indices + from seqlen_list assigned to that partition. + prefix (str): A prefix to be added to each metric key in the returned dictionary. + + Returns: + dict: A dictionary containing metrics related to sequence length imbalance. + """ + # Get the number of partitions + k_partition = len(partitions) + # assert len(seqlen_list) % k_partition == 0 + batch_size = len(seqlen_list) // k_partition + min_sum_seqlen = None + max_sum_seqlen = None + total_sum_seqlen = 0 + + # Iterate over each batch of sequence lengths + for offset in range(0, len(seqlen_list), batch_size): + cur_sum_seqlen = sum(seqlen_list[offset : offset + batch_size]) + if min_sum_seqlen is None or cur_sum_seqlen < min_sum_seqlen: + min_sum_seqlen = cur_sum_seqlen + if max_sum_seqlen is None or cur_sum_seqlen > max_sum_seqlen: + max_sum_seqlen = cur_sum_seqlen + total_sum_seqlen += cur_sum_seqlen + + balanced_sum_seqlen_list = [] + for partition in partitions: + cur_sum_seqlen_balanced = sum([seqlen_list[i] for i in partition]) + balanced_sum_seqlen_list.append(cur_sum_seqlen_balanced) + # print("balanced_sum_seqlen_list: ", balanced_sum_seqlen_list) + min_sum_seqlen_balanced = min(balanced_sum_seqlen_list) + max_sum_seqlen_balanced = max(balanced_sum_seqlen_list) + + return { + f"{prefix}/min": min_sum_seqlen, + f"{prefix}/max": max_sum_seqlen, + f"{prefix}/minmax_diff": max_sum_seqlen - min_sum_seqlen, + f"{prefix}/balanced_min": min_sum_seqlen_balanced, + f"{prefix}/balanced_max": max_sum_seqlen_balanced, + f"{prefix}/mean": total_sum_seqlen / len(partitions), + } + + +def ceildiv(a: int, b: int) -> int: + """Compute ceiling division of a by b. + + Returns the smallest integer greater than or equal to a/b. + Uses the identity: ceil(a/b) = floor((a + b - 1) / b) = -(-a // b) + + Args: + a: Dividend (numerator). + b: Divisor (denominator), must be non-zero. + + Returns: + int: Ceiling of a divided by b. + + Example: + >>> ceildiv(7, 3) # ceil(7/3) = ceil(2.33) = 3 + 3 + >>> ceildiv(6, 3) # ceil(6/3) = ceil(2.0) = 2 + 2 + """ + return -(a // -b) + + +def roundup_divisible(a: int, b: int) -> int: + """Round up a to the nearest multiple of b. + + Returns the smallest multiple of b that is >= a. + + Args: + a: Value to round up. + b: Divisor to round to (must be positive). + + Returns: + int: Smallest multiple of b that is >= a. + + Example: + >>> roundup_divisible(7, 4) # nearest multiple of 4 >= 7 is 8 + 8 + >>> roundup_divisible(8, 4) # 8 is already a multiple of 4 + 8 + """ + return ((a + b - 1) // b) * b + + +def rearrange_micro_batches( + batch, + max_token_len, + dp_group=None, + num_batches_divided_by=None, + same_micro_num_in_dp=True, + min_num_micro_batch=None, + use_dynamic_bsz_balance=True, + force_group_size=1, +): + """ + Split a batch into micro-batches by total token count, with optional DP sync and padding. + + Args: + batch (TensorDict): must include "attention_mask" (B*S); other fields are sliced similarly. + max_token_len (int): max sum of attention_mask per micro-batch. + dp_group (optional): torch.distributed group for data-parallel sync. + num_batches_divided_by (optional): virtual pipeline parallel size, for megatron. + same_micro_num_in_dp (bool): if True and dp_group set, pad all ranks to the same count. + min_num_micro_batch (int, optional): force at least this many splits (pads empty ones). + use_dynamic_bsz_balance (bool, optional): balance the computational workload between micro-batches + force_group_size (int, optional): force consecutive samples to be in the same micro-batch (for RM training). + + Returns: + List[TensorDict]: the micro-batches. + List[List[int]]: index lists mapping each micro-batch back to original positions. + """ + # this is per local micro_bsz + input_ids = batch["input_ids"] + if input_ids.is_nested: + seq_len_effective: torch.Tensor = input_ids.offsets().diff() + max_seq_len = max(seq_len_effective) + else: + max_seq_len = batch["attention_mask"].shape[-1] + seq_len_effective: torch.Tensor = batch["attention_mask"].sum(dim=1) + + assert max_token_len >= max_seq_len, ( + f"max_token_len must be greater than the sequence length. Got {max_token_len=} and {max_seq_len=}" + ) + + # Validate force_group_size + batch_size = len(seq_len_effective) + assert batch_size % force_group_size == 0, ( + f"Batch size {batch_size} must be divisible by force_group_size {force_group_size}" + ) + + total_seqlen = seq_len_effective.sum().item() + # NOTE: num_microbatches <= batch_size, so take the min of this two. + # When force_group_size > 1, we work with groups instead of individual samples + num_groups = batch_size // force_group_size + num_micro_batches = min(num_groups, ceildiv(total_seqlen, max_token_len)) + if min_num_micro_batch is not None: + # used to support pp + num_micro_batches = max(min_num_micro_batch, num_micro_batches) + if dist.is_initialized() and same_micro_num_in_dp and dp_group is not None: + num_micro_batches = torch.tensor([num_micro_batches], device=get_device_name()) + dist.all_reduce(num_micro_batches, op=dist.ReduceOp.MAX, group=dp_group) + num_micro_batches = num_micro_batches.cpu().item() + if num_batches_divided_by is not None: + num_micro_batches = roundup_divisible(num_micro_batches, num_batches_divided_by) + + assert num_micro_batches <= num_groups + + # upcast to int64 to avoid potential overflow im `calculate_workload` computation. + seq_len_effective = seq_len_effective.long() + + # When force_group_size > 1, aggregate workloads by groups + if force_group_size > 1: + # Calculate workload for each group (sum of workloads of samples in the group) + workloads_per_sample = calculate_workload(seq_len_effective) + workloads_per_sample_grouped = workloads_per_sample.view(num_groups, force_group_size) + group_workloads = workloads_per_sample_grouped.sum(dim=1).cpu().tolist() + + # Partition groups instead of individual samples + micro_bsz_group_idx = get_seqlen_balanced_partitions(group_workloads, num_micro_batches, equal_size=False) + + # Convert group indices back to sample indices + micro_bsz_idx = [] + for group_partition in micro_bsz_group_idx: + sample_partition = [] + for group_idx in group_partition: + start_idx = group_idx * force_group_size + sample_partition.extend(range(start_idx, start_idx + force_group_size)) + micro_bsz_idx.append(sample_partition) + + workloads = group_workloads + else: + # Original logic for force_group_size == 1 + # note that seq_len_effective is a GPU tensor. We need to make it a list to avoid D2H! + workloads = calculate_workload(seq_len_effective).cpu().tolist() + micro_bsz_idx = get_seqlen_balanced_partitions(workloads, num_micro_batches, equal_size=False) + + if use_dynamic_bsz_balance: + # Use the sum of squared sequence lengths to approximate attention computation workload + if force_group_size > 1: + # For grouped samples, use group workloads for sorting + micro_bsz_idx.sort( + key=lambda partition: ( + sum(workloads[idx // force_group_size] for idx in partition[::force_group_size]), + partition[0] if partition else 0, + ), + reverse=True, + ) + else: + micro_bsz_idx.sort( + key=lambda partition: ( + sum(workloads[idx] for idx in partition), + partition[0] if partition else 0, + ), + reverse=True, + ) + # Place smaller micro-batches at both ends to reduce the bubbles exposed during the warm-up and cool-down. + micro_bsz_idx = micro_bsz_idx[::2][::-1] + micro_bsz_idx[1::2] + + micro_batches = [] + + for partition in micro_bsz_idx: + curr_micro_batch = tu.index_select_tensor_dict(batch, partition) + micro_batches.append(curr_micro_batch) + + return micro_batches, micro_bsz_idx + + +def get_reverse_idx(idx_map): + """ + Build the inverse of an index mapping. + + Args: + idx_map (Sequence[int]): Sequence where idx_map[i] = j. + + Returns: + List[int]: Inverse mapping list such that output[j] = i for each i. + """ + reverse_idx_map = copy.deepcopy(idx_map) + + for i, idx in enumerate(idx_map): + reverse_idx_map[idx] = i + + return reverse_idx_map + + +def prepare_dynamic_batch( + data: DataProto, + max_token_len: int, + dp_group=None, + num_batches_divided_by=None, + same_micro_num_in_dp=True, + min_num_micro_batch=None, + use_dynamic_bsz_balance=True, +) -> tuple[list[DataProto], list[list[int]]]: + """ + Prepare a batch for dynamic batching. + + Args: + data (DataProto): The input data. + max_token_len (int): The maximum token length for dynamic batching. + + Returns: + Tuple[List[DataProto], List[List[int]]]: A tuple containing a list of DataProto objects + and a list of index lists. + """ + batch, batch_idx_list = rearrange_micro_batches( + data.batch, + max_token_len=max_token_len, + dp_group=dp_group, + num_batches_divided_by=num_batches_divided_by, + same_micro_num_in_dp=same_micro_num_in_dp, + min_num_micro_batch=min_num_micro_batch, + use_dynamic_bsz_balance=use_dynamic_bsz_balance, + ) + micro_batches = [] + for i, batch_idx in enumerate(batch_idx_list): + tensors = dict(batch[i]) + non_tensors = {key: value[batch_idx] for key, value in data.non_tensor_batch.items()} + meta_info = copy.deepcopy(data.meta_info) + micro_batches.append(DataProto.from_dict(tensors, non_tensors, meta_info=meta_info)) + + return micro_batches, batch_idx_list + + +def restore_dynamic_batch(data: torch.Tensor, batch_idx_list: list[list[int]]) -> torch.Tensor: + """ + Restore a batch from dynamic batching. + + Args: + data (torch.Tensor): The input data. + batch_idx_list (List[List[int]]): The list of index lists. + + Returns: + torch.Tensor: The restored data. + """ + indices = list(chain.from_iterable(batch_idx_list)) + batch_size = data.shape[0] + assert len(indices) == batch_size, f"{len(indices)} vs. {batch_size}" + revert_indices = torch.tensor(get_reverse_idx(indices), dtype=torch.long) + + if data.is_nested: + data_lst = data.unbind() + tensors = [data_lst[i] for i in revert_indices] + reverted_data = torch.nested.as_nested_tensor(tensors, layout=torch.jagged) + else: + reverted_data = data[revert_indices] + + return reverted_data + + +def get_group_balanced_partitions( + seqlen_list: list[int], + uid_list: list, + k_partitions: int, +) -> list[list[int]]: + """ + Partition samples into k groups while keeping samples with the same uid together. + + Args: + seqlen_list: List of sequence lengths for each sample. + uid_list: List of uids identifying which samples share the same prefix. + Samples with the same uid will be kept together. + k_partitions: Number of partitions (typically world_size). + + Returns: + List of k lists, each containing sample indices assigned to that partition. + Samples with the same uid are guaranteed to be in the same partition. + """ + assert len(seqlen_list) == len(uid_list), "seqlen_list and uid_list must have same length" + + # Build groups: each group contains indices of samples with the same uid + # Assumes samples with same uid are contiguous + groups = [] # List of (group_indices, group_total_seqlen) + current_uid = None + current_indices = [] + current_seqlen = 0 + + for i, (seqlen, uid) in enumerate(zip(seqlen_list, uid_list, strict=False)): + if uid != current_uid: + if current_indices: + groups.append((current_indices, current_seqlen)) + current_uid = uid + current_indices = [i] + current_seqlen = seqlen + else: + current_indices.append(i) + current_seqlen += seqlen + + # Don't forget the last group + if current_indices: + groups.append((current_indices, current_seqlen)) + + num_groups = len(groups) + assert num_groups >= k_partitions, ( + f"Number of uid groups ({num_groups}) must be >= k_partitions ({k_partitions}). " + f"Consider reducing world_size or increasing batch_size." + ) + + # Calculate workload for each group (as integers for partitioning) + group_workloads = [] + for indices, total_seqlen in groups: + # Use sum of individual workloads for more accurate estimation + workload = sum(int(calculate_workload(torch.tensor([seqlen_list[i]])).item()) for i in indices) + group_workloads.append(workload) + + # Use Karmarkar-Karp to partition groups + # equal_size=True ensures each partition gets the same number of groups, + # which is required when each group has the same number of samples (rollout.n) + group_partitions = get_seqlen_balanced_partitions( + seqlen_list=group_workloads, + k_partitions=k_partitions, + equal_size=True, + ) + + # Convert group partitions to sample partitions + sample_partitions = [] + for group_partition in group_partitions: + sample_indices = [] + for group_idx in group_partition: + sample_indices.extend(groups[group_idx][0]) + sample_partitions.append(sorted(sample_indices)) + + return sample_partitions diff --git a/verl/verl/utils/sglang/__init__.py b/verl/verl/utils/sglang/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d5aee6bdcdd86a19e72d4a8e96dd3f6b9abc337 --- /dev/null +++ b/verl/verl/utils/sglang/__init__.py @@ -0,0 +1,14 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/verl/verl/utils/sglang/sglang_fp8_utils.py b/verl/verl/utils/sglang/sglang_fp8_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ac3a8ff0699fab957fb126cde942d97f6fada5f4 --- /dev/null +++ b/verl/verl/utils/sglang/sglang_fp8_utils.py @@ -0,0 +1,21 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from verl.utils.fp8_utils import FP8QuantizerHelper + + +class SGLangFP8QuantizerHelper(FP8QuantizerHelper): + def __init__(self, quant_config): + super().__init__(quant_config) diff --git a/verl/verl/utils/tensordict_utils.py b/verl/verl/utils/tensordict_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..91d82b15ffd7990c5e50107f3882d83f1558c604 --- /dev/null +++ b/verl/verl/utils/tensordict_utils.py @@ -0,0 +1,950 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +from typing import Any, Iterable + +import numpy as np +import tensordict +import torch +from packaging.version import parse as parse_version +from tensordict import TensorDict +from tensordict.tensorclass import NonTensorData, NonTensorStack + + +def assign_non_tensor_data(tensor_dict: TensorDict, key, val): + """Assign a single non-tensor value to a TensorDict. + + Wraps the value in NonTensorData so it can be stored alongside tensors + in the TensorDict. Use this for scalar metadata or simple non-tensor values. + + Args: + tensor_dict: The TensorDict to assign to. + key: The key under which to store the value. + val: Any non-tensor value to store (e.g., string, int, dict). + + Raises: + AssertionError: If tensor_dict is not a TensorDict. + + Example: + >>> td = TensorDict({"obs": torch.randn(3, 4)}, batch_size=[3]) + >>> assign_non_tensor_data(td, "experiment_name", "run_001") + """ + assert isinstance(tensor_dict, TensorDict), "input dict must be a TensorDict" + tensor_dict[key] = NonTensorData(val) + + +def assign_non_tensor_stack(tensor_dict: TensorDict, key, val: list): + """Assign a list with potentially nested structures (lists, dicts, etc.) to TensorDict. + + This function handles complex nested data structures like: + - Lists of lists: [[], [0.5, 0.8], [0.9]] + - Lists of dicts: [{"acc": 1.0}, {"acc": 0.0}] + - Lists of lists of dicts: [[{"content": "...", "role": "user"}]] + + These structures are wrapped in NonTensorStack so TensorDict can handle them correctly. + + Args: + tensor_dict: The TensorDict to assign to + key: The key to assign the value under + val: A list containing potentially nested structures + + Example: + >>> td = TensorDict({}, batch_size=[]) + >>> turn_scores = [[], [0.5, 0.8], [0.9]] + >>> assign_non_tensor_stack(td, "turn_scores", turn_scores) + >>> # Now td["turn_scores"] contains the nested data + """ + # Convert list to NonTensorStack to handle nested structures + # This wraps each item in NonTensorData to preserve complex objects + # TODO(petersh6): can convert back to val directly if we are not accessing .data from the NonTensorStack + assert isinstance(tensor_dict, TensorDict), "input dict must be a TensorDict" + tensor_dict[key] = NonTensorStack.from_list([NonTensorData(item) for item in val]) + + +def assign_non_tensor(tensor_dict: TensorDict, **kwargs): + """Assign non-tensor data to a TensorDict. + + Automatically detects if the value is a list with nested structures and uses + the appropriate assignment method (NonTensorData for simple values, + NonTensorStack for lists with nested structures). + + Args: + tensor_dict: The TensorDict to assign to + **kwargs: Key-value pairs where values can be: + - Simple values (stored as NonTensorData) + - Lists with nested structures (stored as NonTensorStack) + + Example: + >>> td = TensorDict({"obs": torch.randn(3, 4)}, batch_size=[3]) + >>> assign_non_tensor( + ... tensor_dict=td, + ... metadata="experiment_1", # Simple value + ... turn_scores=[[], [0.5, 0.8], [0.9]] # Nested list + ... ) + """ + assert isinstance(tensor_dict, TensorDict), "input dict must be a TensorDict" + for key, val in kwargs.items(): + if isinstance(val, (NonTensorData | NonTensorStack)): + tensor_dict[key] = val + elif isinstance(val, list): + # For lists, use NonTensorStack + assign_non_tensor_stack(tensor_dict=tensor_dict, key=key, val=val) + else: + # For non-list values, use NonTensorData + assign_non_tensor_data(tensor_dict=tensor_dict, key=key, val=val) + return tensor_dict + + +def unwrap_non_tensor_data(data): + """Unwrap a NonTensorData object to get the underlying value. + + If the input is a NonTensorData wrapper, extracts and returns the + underlying data. Otherwise, returns the input unchanged. + + Args: + data: Either a NonTensorData object or any other value. + + Returns: + The unwrapped data if input was NonTensorData, otherwise the + original input unchanged. + + Example: + >>> wrapped = NonTensorData("hello") + >>> unwrap_non_tensor_data(wrapped) + 'hello' + >>> unwrap_non_tensor_data(42) # Non-wrapped value + 42 + """ + if isinstance(data, NonTensorData): + return data.data + return data + + +def get_non_tensor_data(data: TensorDict, key: str, default): + """Retrieve and unwrap non-tensor data from a TensorDict. + + Fetches the value for the given key from the TensorDict and automatically + unwraps it if it's stored as NonTensorData. + + Args: + data: The TensorDict to retrieve from. + key: The key to look up. + default: Value to return if the key is not found. + + Returns: + The unwrapped value if the key exists and was wrapped in NonTensorData, + the raw value if it wasn't wrapped, or the default if key not found. + + Example: + >>> td = TensorDict({}, batch_size=[]) + >>> assign_non_tensor_data(td, "config", {"lr": 0.01}) + >>> get_non_tensor_data(td, "config", None) + {'lr': 0.01} + >>> get_non_tensor_data(td, "missing", "default_value") + 'default_value' + """ + output = data.get(key, default) + return unwrap_non_tensor_data(output) + + +def nested_tensor_from_tensor_list(tensors: list[torch.Tensor], ragged_idx: int | None = None) -> torch.Tensor: + assert len(tensors) > 0, "Must provide at least one tensor" + sample_dim = tensors[0].dim() + if ragged_idx is None: + ragged_idx = sample_dim + assert 1 <= ragged_idx <= sample_dim, ( + f"ragged_idx must be in [1, {sample_dim}]. Got {ragged_idx=} and {sample_dim=}" + ) + + if sample_dim == 1: + return torch.nested.as_nested_tensor(tensors, layout=torch.jagged) + + cat_dim = ragged_idx - 1 + values = torch.cat(tensors, dim=cat_dim) + lengths = torch.tensor([tensor.shape[cat_dim] for tensor in tensors], dtype=torch.long, device=values.device) + offsets = torch.zeros(len(tensors) + 1, dtype=torch.long, device=values.device) + torch.cumsum(lengths, dim=0, out=offsets[1:]) + + nested_tensor = torch.nested.nested_tensor_from_jagged(values=values, offsets=offsets) + nested_tensor._ragged_idx = ragged_idx + return nested_tensor + + +def concat_nested_tensors(tensors: list[torch.Tensor]) -> torch.Tensor: + """Concatenate multiple nested tensors along the batch dimension. + + Takes a list of nested tensors with jagged layout and concatenates them + into a single nested tensor. Each input tensor must have 2 or more dimensions and be contiguous. + + Args: + tensors: List of nested tensors to concatenate. All tensors must + be nested, contiguous, and have 2 or more dimensions. + + Returns: + A new nested tensor with jagged layout containing all rows from + the input tensors concatenated along dimension 0. + + Raises: + AssertionError: If any tensor is not nested, not contiguous, or + doesn't have 2 or more dimensions. + + Example: + >>> t1 = torch.nested.as_nested_tensor([torch.randn(3), torch.randn(5)], layout=torch.jagged) + >>> t2 = torch.nested.as_nested_tensor([torch.randn(2), torch.randn(4)], layout=torch.jagged) + >>> result = concat_nested_tensors([t1, t2]) + >>> # result contains 4 rows: lengths [3, 5, 2, 4] + """ + for tensor in tensors: + assert tensor.is_nested and tensor.is_contiguous() + unbind_tensors = [] + for tensor in tensors: + assert len(tensor.shape) >= 2, f"nested tensor must have 2 or more dimensions. Got {tensor.shape}" + unbind_tensor = tensor.unbind(0) + unbind_tensors.extend(list(unbind_tensor)) + + ragged_idx = getattr(tensors[0], "_ragged_idx", tensors[0].dim() - 1) + return nested_tensor_from_tensor_list(unbind_tensors, ragged_idx=ragged_idx) + + +def concat_tensordict_with_none_bsz(data: list[TensorDict]): + """Handle concatenation of TensorDicts with empty batch size. + + For TensorDicts that contain only metadata (NonTensorData) with no batch + dimension, returns the first TensorDict as the concatenation result. + + Args: + data: List of TensorDicts, each with empty batch_size (batch_size=[]). + + Returns: + The first TensorDict from the list, as metadata concatenation + simply preserves the first instance. + + Raises: + AssertionError: If any TensorDict has a non-empty batch_size. + + Note: + This is used internally by concat_tensordict when handling + TensorDicts that contain only non-tensor metadata. + """ + for d in data: + assert len(d.batch_size) == 0 + # directly return the first meta info + return data[0] + + +def concat_tensordict(data: list[TensorDict]) -> TensorDict: + """Concatenate multiple TensorDicts along dimension zero. + + Combines a list of TensorDicts into a single TensorDict by concatenating + all tensors along the batch dimension (dim=0). Handles nested tensors + specially by unbinding and rebinding them. + + Args: + data: List of TensorDicts to concatenate. All TensorDicts must have + the same keys and the same set of nested tensor keys. + + Returns: + A new TensorDict containing concatenated tensors from all inputs. + + Raises: + AssertionError: If data is empty or if TensorDicts have inconsistent + nested tensor keys. + + Note: + - For TensorDicts with empty batch_size, returns the first one + - Nested tensors are handled specially via concat_nested_tensors + - Regular tensors use TensorDict.cat for efficient concatenation + """ + assert len(data) > 0, "Must have at least one tensordict" + + # Find nested tensor keys from the first tensordict + nested_tensor_keys = {key for key, value in data[0].items() if isinstance(value, torch.Tensor) and value.is_nested} + + if not nested_tensor_keys: + if len(data[0].batch_size) == 0: + return concat_tensordict_with_none_bsz(data) + # if batch size is None (only contain NonTensorData) + return TensorDict.cat(data, dim=0) + + # Create a list of tensordicts containing only non-nested tensors for concatenation + regular_tds = [] + for td in data: + current_nested_keys = {k for k, v in td.items() if isinstance(v, torch.Tensor) and v.is_nested} + assert current_nested_keys == nested_tensor_keys, "All tensordicts must have the same set of nested tensors." + + # Create a new TensorDict with non-nested items without modifying the original + regular_items = {k: v for k, v in td.items() if k not in nested_tensor_keys} + regular_tds.append(TensorDict(regular_items, batch_size=td.batch_size, device=td.device)) + + # Concatenate the regular tensordicts + output = TensorDict.cat(regular_tds, dim=0) + + # Concatenate and add nested tensors to the output + for key in nested_tensor_keys: + nested_tensors_to_concat = [td[key] for td in data] + output[key] = concat_nested_tensors(nested_tensors_to_concat) + + return output + + +def chunk_tensordict(td: TensorDict, chunks: int) -> list[TensorDict]: + """Split a TensorDict into equal-sized chunks with special nested tensor handling. + + Divides a TensorDict into the specified number of chunks along the batch + dimension. Handles NestedTensors specially since TensorDict.chunk() doesn't + support jagged tensors. + + Args: + td: The TensorDict to split. + chunks: Number of chunks to create. Must evenly divide len(td). + + Returns: + List of TensorDicts, each containing a portion of the original data. + + Raises: + AssertionError: If td is not a TensorDict or if its length is not + evenly divisible by chunks. + + Note: + PyTorch ``unbind(dim=0)`` on 3D+ jagged NestedTensors has a bug where + ``split_with_sizes`` is applied to the wrong dimension of the internal + ``_values`` tensor. For example, mRoPE ``position_ids`` with per-sample + shape ``(4, seq_len)`` becomes a 3D jagged NestedTensor + ``[B, *(ragged=4), seq_len]``; ``_values`` is ``[B*4, seq_len]`` and + ``unbind`` erroneously splits dimension 1 (``seq_len``) instead of + dimension 0, causing:: + + RuntimeError: split_with_sizes expects split_sizes to sum exactly + to , but got split_sizes=[4, 4, ...] + + 2D jagged NestedTensors (e.g. ``input_ids``, ``loss_mask``) are + unaffected — ``unbind(dim=0)`` works correctly for them. + + The workaround: try ``unbind`` first (fast path for 2D); on failure, + fall back to ``to_padded_tensor`` → ``chunk`` → reconstruct per-chunk + NestedTensors using the original ragged lengths from ``offsets``. + + See https://github.com/pytorch/pytorch/issues/153238 + """ + assert isinstance(td, TensorDict) and len(td) % chunks == 0, ( + f"expecting td with length divisible by chunks, but got {len(td)} and {chunks}" + ) + chunk_size = len(td) // chunks + nested_keys = {key for key, val in td.items() if isinstance(val, torch.Tensor) and val.is_nested} + new_td = TensorDict( + {k: v for k, v in td.items() if k not in nested_keys}, batch_size=td.batch_size, device=td.device + ) + + tds = new_td.chunk(chunks=chunks) + for key in nested_keys: + nt = td[key] + try: + tensors = nt.unbind(dim=0) + except RuntimeError: + padded = nt.to_padded_tensor(0) + padded_chunks = padded.chunk(chunks, dim=0) + offsets = nt.offsets() + lengths = offsets.diff().tolist() + for i, chunk_td in enumerate(tds): + chunk_lengths = lengths[i * chunk_size : (i + 1) * chunk_size] + chunk_tensors = [padded_chunks[i][j, :seq_len] for j, seq_len in enumerate(chunk_lengths)] + chunk_td[key] = nested_tensor_from_tensor_list( + chunk_tensors, ragged_idx=getattr(nt, "_ragged_idx", nt.dim() - 1) + ) + continue + + for i, chunk_td in enumerate(tds): + chunk_td[key] = nested_tensor_from_tensor_list( + list(tensors[i * chunk_size : (i + 1) * chunk_size]), + ragged_idx=getattr(nt, "_ragged_idx", nt.dim() - 1), + ) + + return tds + + +def get_tensordict(tensor_dict: dict[str, torch.Tensor | list], non_tensor_dict: dict = None) -> TensorDict: + """Create a TensorDict from tensors and non-tensor data. + + Automatically handles nested structures in lists by converting them to NonTensorStack. + This enables support for: + - Lists of lists: [[], [0.5, 0.8], [0.9]] + - Lists of dicts: [{"acc": 1.0}, {"acc": 0.0}] + - Lists of lists of dicts: [[{"content": "...", "role": "user"}]] + + Args: + tensor_dict: Dictionary of tensors and lists to include in the TensorDict + non_tensor_dict: Dictionary of metadata to store as NonTensorData + + Returns: + TensorDict with proper handling of nested structures + + Example: + >>> td = get_tensordict( + ... tensor_dict={ + ... "obs": torch.randn(3, 4), + ... "turn_scores": [[], [0.5, 0.8], [0.9]] # Nested list + ... }, + ... non_tensor_dict={"experiment": "test"} + ... ) + """ + tensor_dict = tensor_dict.copy() + if non_tensor_dict is None: + non_tensor_dict = {} + + batch_size = None + + for key, val in tensor_dict.items(): + if isinstance(val, torch.Tensor) and val.is_nested: + assert val.is_contiguous(), "Nested tensors must be contiguous. Try setting layout=torch.jagged" + assert val.layout == torch.jagged, "Nested tensors must be jagged." + + # Skip validation for NonTensorStack as it's already properly formatted + if isinstance(val, NonTensorStack): + if batch_size is None: + batch_size = len(val) + else: + assert len(val) == batch_size, ( + f"Batch size of NonTensorStack {key} is not consistent with other tensors. " + f"Expected {batch_size}, got {len(val)}" + ) + continue + + if isinstance(val, list | np.ndarray): + for v in val: + assert not isinstance(v, torch.Tensor), ( + "Passing a list makes the data NonTensorStack, " + "which doesn't support torch.Tensor. Please convert to numpy first" + ) + # Convert to NonTensorStack to handle nested structures + tensor_dict[key] = NonTensorStack.from_list([NonTensorData(item) for item in val]) + + assert isinstance(val, torch.Tensor | list | np.ndarray), ( + f"{key} -> {type(val)} isn't of 'torch.Tensor | list | np.ndarray' type" + ) + + if batch_size is None: + batch_size = val.size(0) if isinstance(val, torch.Tensor) else len(val) + else: + val_batch_size = val.size(0) if isinstance(val, torch.Tensor) else len(val) + assert val_batch_size == batch_size, ( + f"Batch size of tensor {key} is not consistent with other tensors. " + f"Expected {batch_size}, got {val_batch_size}" + ) + + if batch_size is None: + batch_size = [] + else: + batch_size = [batch_size] + + for key, val in non_tensor_dict.items(): + assert key not in tensor_dict + tensor_dict[key] = NonTensorData(val) + + return TensorDict(source=tensor_dict, batch_size=batch_size) + + +def index_select_tensor_dict(batch: TensorDict, indices: torch.Tensor | list[int]) -> TensorDict: + """Select rows from a TensorDict using indices. + + Creates a new TensorDict containing only the rows specified by indices. + Handles regular tensors, nested tensors, NonTensorStack, and NonTensorData + appropriately. + + Args: + batch: The TensorDict to index into. Can be None. + indices: 1D tensor or list of integers specifying which rows to select. + + Returns: + A new TensorDict containing only the selected rows, or None if + batch was None. + + Raises: + AssertionError: If indices is not 1-dimensional. + + Note: + - Regular tensors are indexed directly + - Nested tensors are unbound, indexed, and rebound + - NonTensorStack is indexed by batch dimension + - NonTensorData (scalar metadata) is preserved unchanged + """ + if isinstance(indices, list): + indices = torch.tensor(indices) + + assert indices.dim() == 1, "indices must be a 1D tensor" + + data_dict = {} + batch_size = indices.shape[0] + + if batch is not None: + for key, tensor in batch.items(): + if isinstance(tensor, torch.Tensor) and not tensor.is_nested: + data_dict[key] = tensor[indices] + elif isinstance(tensor, torch.Tensor) and tensor.is_nested: + tensor_lst = tensor.unbind() # for performance + selected_tensors = [tensor_lst[idx] for idx in indices] + data_dict[key] = nested_tensor_from_tensor_list( + selected_tensors, ragged_idx=getattr(tensor, "_ragged_idx", tensor.dim() - 1) + ) + else: + # This handles NonTensorStack (indexable by batch dim) and NonTensorData (scalar metadata). + if tensor.shape: + data_dict[key] = tensor[indices] + else: + data_dict[key] = tensor + selected_batch = TensorDict(source=data_dict, batch_size=batch_size) + else: + selected_batch = None + + return selected_batch + + +def union_tensor_dict(tensor_dict1: TensorDict, tensor_dict2: TensorDict) -> TensorDict: + """Merge two TensorDicts, adding keys from the second to the first. + + Performs an in-place union of two TensorDicts. Keys from tensor_dict2 + that don't exist in tensor_dict1 are added. Keys that exist in both + must have identical values. + + Args: + tensor_dict1: The base TensorDict to merge into (modified in-place). + tensor_dict2: The TensorDict whose keys will be added to tensor_dict1. + + Returns: + The modified tensor_dict1 containing the union of both TensorDicts. + + Raises: + AssertionError: If batch sizes don't match, or if a key exists in + both TensorDicts with different values. + + Example: + >>> td1 = TensorDict({"a": torch.tensor([1, 2])}, batch_size=[2]) + >>> td2 = TensorDict({"b": torch.tensor([3, 4])}, batch_size=[2]) + >>> result = union_tensor_dict(td1, td2) + >>> list(result.keys()) + ['a', 'b'] + """ + assert tensor_dict1.batch_size == tensor_dict2.batch_size, ( + f"Two tensor dict must have identical batch size. Got {tensor_dict1.batch_size} and {tensor_dict2.batch_size}" + ) + for key in tensor_dict2.keys(): + if key not in tensor_dict1.keys(): + # Note that there is a difference between tensor_dict2[key] and tensor_dict2.get(key) + tensor_dict1[key] = tensor_dict2.get(key) + else: + if isinstance(tensor_dict2[key], torch.Tensor): + assert tensor_dict1[key].equal(tensor_dict2[key]), ( + f"{key} in tensor_dict1 and tensor_dict2 are not the same object" + ) + else: + # non-tensor + assert tensor_dict1[key] == tensor_dict2[key], ( + f"{key} in tensor_dict1 and tensor_dict2 are not the same object" + ) + + return tensor_dict1 + + +def make_iterator(tensordict: TensorDict, mini_batch_size, epochs, seed=None, dataloader_kwargs=None): + """Create an iterator that yields mini-batches from a TensorDict. + + Wraps a TensorDict in a DataLoader-style iterator that yields mini-batches + for the specified number of epochs. Useful for training loops. + + Args: + tensordict: The TensorDict to iterate over. + mini_batch_size: Size of each mini-batch. Must evenly divide the + TensorDict's batch size. + epochs: Number of times to iterate through the entire dataset. + seed: Optional random seed for reproducible shuffling. + dataloader_kwargs: Optional dict of additional kwargs to pass to + the underlying DataLoader (e.g., shuffle=True, num_workers=4). + + Returns: + An iterator that yields TensorDict mini-batches. + + Raises: + AssertionError: If batch size is not divisible by mini_batch_size. + + Example: + >>> td = TensorDict({"obs": torch.randn(100, 4)}, batch_size=[100]) + >>> for batch in make_iterator(td, mini_batch_size=10, epochs=2): + ... # batch is a TensorDict with batch_size=[10] + ... pass + """ + from torch.utils.data import DataLoader + + assert tensordict.batch_size[0] % mini_batch_size == 0, f"{tensordict.batch_size[0]} % {mini_batch_size} != 0" + # we can directly create a dataloader from TensorDict + if dataloader_kwargs is None: + dataloader_kwargs = {} + + if seed is not None: + generator = torch.Generator() + generator.manual_seed(seed) + else: + generator = None + + assert isinstance(dataloader_kwargs, dict) + + idx_lst = torch.arange(tensordict.shape[0]) + + train_dataloader = DataLoader( + dataset=idx_lst, batch_size=mini_batch_size, collate_fn=lambda x: x, generator=generator, **dataloader_kwargs + ) + + def get_data(): + for _ in range(epochs): + for idx in train_dataloader: + yield index_select_tensor_dict(tensordict, idx) + + return iter(get_data()) + + +def assert_tensordict_eq(tensordict1: TensorDict, tensordict2: TensorDict): + """Assert that two TensorDicts are equal. + + Performs a deep equality check between two TensorDicts, verifying that + they have the same keys with identical values. Handles nested tensors + by comparing their unbound components. + + Args: + tensordict1: First TensorDict to compare. + tensordict2: Second TensorDict to compare. + + Raises: + AssertionError: If the TensorDicts differ in keys, value types, or + value contents. The error message indicates what differs. + + Note: + - Regular tensors are compared element-wise + - Nested tensors are unbound and compared component by component + - Non-tensor values are compared with standard equality + """ + tensordict1_key_set = set(tensordict1.keys()) + tensordict2_key_set = set(tensordict2.keys()) + assert tensordict1_key_set == tensordict2_key_set, ( + f"key set diffs. Got {tensordict2_key_set=} vs {tensordict1_key_set=}" + ) + + for key in tensordict1.keys(): + val = tensordict1[key] + val2 = tensordict2[key] + + assert type(val) is type(val2), f"The type of {key} must be the same. Got {type(val)} vs {type(val2)}" + + if isinstance(val, torch.Tensor): + if val.is_nested: + assert val.is_nested and val2.is_nested, ( + f"Both tensors must be nested tensors. {val.is_nested=}, {val2.is_nested=}" + ) + t1, t2 = val.unbind(), val2.unbind() + assert len(t1) == len(t2), f"Nested tensor should have the same lengths. {len(t1)=} vs {len(t2)=}" + for c1, c2 in zip(t1, t2, strict=True): + assert torch.equal(c1, c2), f"Nested tensor components have different values. {c1=} vs {c2=}" + else: + assert torch.all(torch.eq(val, val2)).item() + else: + assert val == val2 + + +def get(tensordict: TensorDict, key: str, default=None) -> Any: + """Get a value from a TensorDict with automatic unwrapping. + + Retrieves a value from the TensorDict and automatically converts it + to a Python-native format: + - Tensors are returned as-is + - NonTensorStack is converted to a Python list + - NonTensorData is unwrapped to its underlying value + + Args: + tensordict: The TensorDict to retrieve from. + key: The key to look up. + default: Value to return if the key doesn't exist. Defaults to None. + + Returns: + The value for the key in its native format, or default if not found. + + Example: + >>> td = get_tensordict({"obs": torch.randn(3, 4), "labels": ["a", "b", "c"]}) + >>> get(td, "obs") # Returns torch.Tensor + >>> get(td, "labels") # Returns ["a", "b", "c"] as a list + >>> get(td, "missing", "default") # Returns "default" + """ + if key not in tensordict: + return default + + output = tensordict.get(key) + if isinstance(output, torch.Tensor): + return output + elif isinstance(output, NonTensorStack): + return output.tolist() + else: + assert isinstance(output, NonTensorData) + return output.data + + +def get_keys(tensordict: TensorDict, keys: Iterable[str]) -> TensorDict: + """Extract a subset of keys from a TensorDict into a new TensorDict. + + Creates a new TensorDict containing only the specified keys. Values + are properly categorized as tensor or non-tensor data. + + Args: + tensordict: The source TensorDict. + keys: Iterable of key names to extract. + + Returns: + A new TensorDict containing only the specified keys with their values. + + Raises: + KeyError: If any key in keys doesn't exist in the tensordict. + + Example: + >>> td = get_tensordict({"a": torch.randn(3), "b": torch.randn(3), "c": torch.randn(3)}) + >>> subset = get_keys(td, ["a", "c"]) + >>> list(subset.keys()) + ['a', 'c'] + """ + tensor_output = {} + non_tensor_output = {} + for key in keys: + if key not in tensordict.keys(): + raise KeyError(f"key {key} not in tensordict") + output = tensordict.get(key) + if isinstance(output, torch.Tensor): + tensor_output[key] = output + elif isinstance(output, NonTensorStack): + tensor_output[key] = output.tolist() + else: + assert isinstance(output, NonTensorData) + non_tensor_output[key] = output.data + + return get_tensordict(tensor_output, non_tensor_output) + + +def pop(tensordict: TensorDict, key: str, default=None) -> Any: + """Remove and return a value from a TensorDict with automatic unwrapping. + + Removes the specified key from the TensorDict and returns its value, + automatically converting to Python-native format (same as get()). + + Args: + tensordict: The TensorDict to pop from. + key: The key to remove and return. + default: Value to return if the key doesn't exist. Defaults to None. + + Returns: + The value for the key in its native format, or default if not found. + The key is removed from the TensorDict. + + Example: + >>> td = get_tensordict({"obs": torch.randn(3, 4), "labels": ["a", "b", "c"]}) + >>> labels = pop(td, "labels") # Returns ["a", "b", "c"], removes from td + >>> "labels" in td.keys() + False + """ + _sentinel = object() + output = tensordict.pop(key, _sentinel) + if output is _sentinel: + return default + + if isinstance(output, torch.Tensor): + return output + elif isinstance(output, NonTensorStack): + return output.tolist() + else: + assert isinstance(output, NonTensorData) + return output.data + + +def pop_keys(tensordict: TensorDict, keys: Iterable[str]) -> TensorDict: + """Remove multiple keys from a TensorDict and return them as a new TensorDict. + + Removes the specified keys from the source TensorDict and creates a new + TensorDict containing those keys and their values. + + Args: + tensordict: The source TensorDict to pop from (modified in-place). + keys: Iterable of key names to remove and return. + + Returns: + A new TensorDict containing the popped keys and their values. + + Raises: + KeyError: If any key in keys doesn't exist in the tensordict. + + Example: + >>> td = get_tensordict({"a": torch.randn(3), "b": torch.randn(3), "c": torch.randn(3)}) + >>> popped = pop_keys(td, ["a", "c"]) + >>> list(td.keys()) # Only 'b' remains + ['b'] + >>> list(popped.keys()) + ['a', 'c'] + """ + tensor_output = {} + non_tensor_output = {} + for key in keys: + if key not in tensordict.keys(): + raise KeyError(f"key {key} not in tensordict") + output = tensordict.get(key) + if isinstance(output, torch.Tensor): + tensor_output[key] = tensordict.pop(key) + elif isinstance(output, NonTensorStack): + tensor_output[key] = tensordict.pop(key).tolist() + else: + assert isinstance(output, NonTensorData) + non_tensor_output[key] = tensordict.pop(key) + + return get_tensordict(tensor_output, non_tensor_output) + + +def pad_to_divisor(data: TensorDict, size_divisor: int): + """Pad a TensorDict's batch dimension to be divisible by a given divisor. + + If the TensorDict's length is not evenly divisible by size_divisor, + pads the batch dimension by repeating elements from the beginning. + Useful for ensuring even distribution across workers in distributed training. + + Args: + data: The TensorDict to pad. + size_divisor: The divisor that the padded length must be divisible by. + + Returns: + tuple: A tuple containing: + - data (TensorDict): The padded TensorDict (or original if no padding needed) + - pad_size (int): Number of elements added as padding (0 if none) + + Raises: + AssertionError: If data is not a TensorDict. + + Example: + >>> td = TensorDict({"obs": torch.randn(10, 4)}, batch_size=[10]) + >>> padded, pad_size = pad_to_divisor(td, 4) + >>> len(padded) # 12 (next multiple of 4 after 10) + 12 + >>> pad_size + 2 + """ + assert isinstance(data, TensorDict), "data must be a TensorDict" + if len(data) % size_divisor != 0: + pad_size = size_divisor - len(data) % size_divisor + padding_protos = [] + remaining_pad = pad_size + while remaining_pad > 0: + take_size = min(remaining_pad, len(data)) + padding_protos.append(data[:take_size]) + remaining_pad -= take_size + data_padded = torch.cat([data] + padding_protos) + else: + if len(data) == 0: + logging.warning("padding a DataProto with no item, no changed made") + pad_size = 0 + data_padded = data + return data_padded, pad_size + + +def unpad(data: TensorDict, pad_size): + """Remove padding from a TensorDict. + + Reverses the effect of pad_to_divisor by removing the specified number + of elements from the end of the TensorDict. + + Args: + data: The padded TensorDict. + pad_size: Number of padding elements to remove. If 0, returns + data unchanged. + + Returns: + The TensorDict with padding removed, equivalent to data[:-pad_size]. + + Example: + >>> td = TensorDict({"obs": torch.randn(12, 4)}, batch_size=[12]) + >>> unpadded = unpad(td, pad_size=2) + >>> len(unpadded) + 10 + """ + if pad_size != 0: + data = data[:-pad_size] + return data + + +def contiguous(data: TensorDict) -> TensorDict: + """Call contiguous on a tensor dict. The contiguous function of tensordict lib will make NonTensorStack. + This function will always return a new tensordict + + Args: + data: The input tensordict + + Returns: + a tensordict that is contiguous + + """ + tensor_dict = {} + non_tensor_dict = {} + + for key in data.keys(): + val = data.get(key) + if isinstance(val, NonTensorData): + non_tensor_dict[key] = val + elif isinstance(val, NonTensorStack): + tensor_dict[key] = val + else: + assert isinstance(val, torch.Tensor), f"Expect val to be a torch.Tensor. Got {type(val)}" + tensor_dict[key] = val.contiguous() + + return get_tensordict(tensor_dict=tensor_dict, non_tensor_dict=non_tensor_dict) + + +def maybe_fix_3d_position_ids(data: TensorDict): + # note for tensordict with pickle/unpickle. nested tensor in tensordict after consolidate and pickle/unpickle + # will incur indexing error for ragged tensor. This only happens when using 3D position ids in VLMs. + # This is likely a bug in tensordict. As a workaround, we manually set _ragged_index. + if "position_ids" in data.keys() and data["position_ids"].dim() == 3 and data["position_ids"].is_nested: + data["position_ids"]._ragged_idx = 2 + + +def list_of_dict_to_tensordict(list_of_dicts: list[dict[str, Any]]) -> TensorDict: + """ + Create a TensorDict from a list of dict of tensors and non_tensors. + Note that this requires tensordict version at least 0.10 + """ + assert parse_version(tensordict.__version__) >= parse_version("0.10"), ( + "Storing non-tensor data in TensorDict at least requires tensordict version 0.10" + ) + + assert len(list_of_dicts) > 0 + + keys = list_of_dicts[0].keys() + dict_of_lists = {key: [d[key] for d in list_of_dicts] for key in keys} + batch_size = len(list_of_dicts) + + final_data = { + key: ( + torch.stack(val_list) + if val_list + and all(isinstance(item, torch.Tensor) for item in val_list) + and all(item.shape == val_list[0].shape for item in val_list) + else ( + torch.nested.as_nested_tensor(val_list, layout=torch.jagged) + if val_list and all(isinstance(item, torch.Tensor) for item in val_list) + else NonTensorStack(*val_list) + ) + ) + for key, val_list in dict_of_lists.items() + } + + td = TensorDict(final_data, batch_size=[batch_size]) + + return td diff --git a/verl/verl/utils/tokenizer.py b/verl/verl/utils/tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..812b1f0e8e540e580137c8b9068d2adbe1fe5b7e --- /dev/null +++ b/verl/verl/utils/tokenizer.py @@ -0,0 +1,244 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Utils for tokenization.""" + +import types +import warnings + +__all__ = [ + "hf_tokenizer", + "hf_processor", + "normalize_token_ids", + "build_multimodal_processor_inputs", + "get_processor_token_id", +] + + +def normalize_token_ids(tokenized_output) -> list[int]: + """Normalize tokenizer outputs into a flat ``list[int]``. + + This handles Transformers 4/5 differences where ``apply_chat_template(tokenize=True)`` + may return either ``list[int]`` or a ``BatchEncoding``/mapping with ``input_ids``. + """ + + token_ids = tokenized_output + if isinstance(tokenized_output, dict): + if "input_ids" in tokenized_output: + token_ids = tokenized_output["input_ids"] + elif hasattr(tokenized_output, "input_ids"): + token_ids = tokenized_output.input_ids + + if hasattr(token_ids, "tolist"): + token_ids = token_ids.tolist() + + if isinstance(token_ids, tuple): + token_ids = list(token_ids) + + if isinstance(token_ids, list) and len(token_ids) == 1 and isinstance(token_ids[0], list | tuple): + token_ids = list(token_ids[0]) + + if not isinstance(token_ids, list): + raise TypeError(f"token_ids must be list-like token ids, got {type(token_ids).__name__}: {token_ids!r}") + + normalized_ids = [] + for idx, token_id in enumerate(token_ids): + if hasattr(token_id, "item"): + token_id = token_id.item() + try: + normalized_ids.append(int(token_id)) + except (TypeError, ValueError) as e: + raise TypeError(f"token_id must be int-convertible, got {type(token_id).__name__}: {token_id!r}") from e + return normalized_ids + + +def get_processor_token_id(processor, token_name: str) -> int | None: + """Resolve a multimodal special token id from a processor. + + Newer processors may expose ``image_token``/``video_token`` strings instead + of ``image_token_id``/``video_token_id`` integers. Fall back to tokenizer + conversion so rollout code can stay processor-agnostic. + """ + + if processor is None: + return None + + token_id_attr = f"{token_name}_token_id" + token_id = getattr(processor, token_id_attr, None) + if token_id is not None: + return int(token_id) + + token_attr = f"{token_name}_token" + token = getattr(processor, token_attr, None) + tokenizer = getattr(processor, "tokenizer", None) + if token is not None and tokenizer is not None: + converted = tokenizer.convert_tokens_to_ids(token) + if converted is not None: + return int(converted) + + return None + + +def _split_videos_and_metadata(videos): + if videos is None: + return None, None + videos = list(videos) + if len(videos) > 0 and isinstance(videos[0], tuple): + video_values, video_metadata = zip(*videos, strict=False) + return list(video_values), list(video_metadata) + return videos, None + + +def build_multimodal_processor_inputs( + processor, + *, + text, + images=None, + videos=None, + audio=None, + mm_processor_kwargs=None, + return_tensors: str = "pt", +): + """Build kwargs for multimodal processor calls. + + This keeps the existing VL flow intact while extending it with audio-aware + paths for processors that accept audio inputs. + """ + processor_kwargs = dict(mm_processor_kwargs or {}) + if audio is not None and "sampling_rate" not in processor_kwargs: + sampling_rate = getattr(getattr(processor, "feature_extractor", None), "sampling_rate", None) + if sampling_rate is not None: + processor_kwargs["sampling_rate"] = int(sampling_rate) + + videos, video_metadata = _split_videos_and_metadata(videos) + processor_kwargs.setdefault("return_tensors", return_tensors) + + if video_metadata is not None: + processor_kwargs.setdefault("video_metadata", video_metadata) + processor_kwargs.setdefault("do_sample_frames", False) + + processor_inputs = {"text": text, "images": images, "videos": videos, **processor_kwargs} + if audio is not None: + processor_inputs["audio"] = audio + + return processor(**processor_inputs) + + +def set_pad_token_id(tokenizer): + """Set pad_token_id to eos_token_id if it is None. + + Args: + tokenizer (transformers.PreTrainedTokenizer): The tokenizer to be set. + + """ + if tokenizer.pad_token_id is None: + tokenizer.pad_token_id = tokenizer.eos_token_id + warnings.warn(f"tokenizer.pad_token_id is None. Now set to {tokenizer.eos_token_id}", stacklevel=1) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + warnings.warn(f"tokenizer.pad_token is None. Now set to {tokenizer.eos_token}", stacklevel=1) + + +def hf_tokenizer(name_or_path, correct_pad_token=True, correct_gemma2=True, **kwargs): + """Create a huggingface pretrained tokenizer which correctness handles eos and pad tokens. + + Args: + + name (str): The name of the tokenizer. + correct_pad_token (bool): Whether to correct the pad token id. + correct_gemma2 (bool): Whether to correct the gemma2 tokenizer. + + Returns: + + transformers.PreTrainedTokenizer: The pretrained tokenizer. + + """ + from transformers import AutoTokenizer + + if correct_gemma2 and isinstance(name_or_path, str) and "gemma-2-2b-it" in name_or_path: + # the EOS token in gemma2 is ambiguious, which may worsen RL performance. + # https://huggingface.co/google/gemma-2-2b-it/commit/17a01657f5c87135bcdd0ec7abb4b2dece04408a + warnings.warn( + "Found gemma-2-2b-it tokenizer. Set eos_token and eos_token_id to and 107.", stacklevel=1 + ) + kwargs["eos_token"] = "" + kwargs["eos_token_id"] = 107 + tokenizer = AutoTokenizer.from_pretrained(name_or_path, **kwargs) + if correct_pad_token: + set_pad_token_id(tokenizer) + return tokenizer + + +def hf_processor(name_or_path, **kwargs): + """Create a huggingface processor to process multimodal data. + + Args: + name_or_path (str): The name of the processor. + + Returns: + Optional[transformers.ProcessorMixin]: The pretrained multimodal processor. + Returns ``None`` for text-only models (including AutoProcessor fallbacks to + tokenizer backends such as ``TokenizersBackend``). + """ + from transformers import AutoConfig, AutoProcessor, PreTrainedTokenizerBase + + try: + processor = AutoProcessor.from_pretrained(name_or_path, **kwargs) + # In newer transformers, AutoProcessor may legitimately fall back to a + # tokenizer backend (e.g. TokenizersBackend) for text-only models. + # Treat it as "no multimodal processor" and let callers use hf_tokenizer. + if isinstance(processor, PreTrainedTokenizerBase): + return None + + config = AutoConfig.from_pretrained(name_or_path, **kwargs) + + # Bind vlm model's get_rope_index method to processor. + processor.config = config + model_class = None + match processor.__class__.__name__: + case "Qwen2VLProcessor": + from transformers.models.qwen2_vl import Qwen2VLModel + + model_class = Qwen2VLModel + case "Qwen2_5_VLProcessor": + from transformers.models.qwen2_5_vl import Qwen2_5_VLModel + + model_class = Qwen2_5_VLModel + case "Qwen3VLProcessor": + from transformers.models.qwen3_vl import Qwen3VLModel + + model_class = Qwen3VLModel + case "Glm4vImageProcessor": + from transformers.models.glm4v import Glm4vModel + + model_class = Glm4vModel + case "MllamaProcessor": + pass # MllamaProcessor and MllamaModel doesn't have get_rope_index property + case _: + raise ValueError(f"Unsupported processor type: {processor.__class__.__name__}") + + if model_class is not None: + processor.get_rope_index = types.MethodType(model_class.get_rope_index, processor) + if hasattr(model_class, "get_vision_position_ids"): + processor.get_vision_position_ids = types.MethodType(model_class.get_vision_position_ids, processor) + except Exception as e: + processor = None + # TODO(haibin.lin): try-catch should be removed after adding transformer version req to setup.py to avoid + # silent failure + warnings.warn(f"Failed to create processor: {e}. This may affect multimodal processing", stacklevel=1) + # Avoid load tokenizer, see: + # https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/models/auto/processing_auto.py#L344 + if processor is not None and "Processor" not in processor.__class__.__name__: + processor = None + + return processor diff --git a/verl/verl/utils/torch_dtypes.py b/verl/verl/utils/torch_dtypes.py new file mode 100644 index 0000000000000000000000000000000000000000..f2f445c26140ceeec25c1d3cf5b3df249c6dffb1 --- /dev/null +++ b/verl/verl/utils/torch_dtypes.py @@ -0,0 +1,80 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Adapted from Cruise. +""" + +import torch + +HALF_LIST = [16, "16", "fp16", "float16", torch.float16] +FLOAT_LIST = [32, "32", "fp32", "float32", torch.float32] +BFLOAT_LIST = ["bf16", "bfloat16", torch.bfloat16] + + +class PrecisionType: + """Type of precision used. + + >>> PrecisionType.HALF == 16 + True + >>> PrecisionType.HALF in (16, "16") + True + """ + + HALF = "16" + FLOAT = "32" + FULL = "64" + BFLOAT = "bf16" + MIXED = "mixed" + + @staticmethod + def supported_type(precision: str | int) -> bool: + return any(x == precision for x in PrecisionType) + + @staticmethod + def supported_types() -> list[str]: + return [x.value for x in PrecisionType] + + @staticmethod + def is_fp16(precision): + return precision in HALF_LIST + + @staticmethod + def is_fp32(precision): + return precision in FLOAT_LIST + + @staticmethod + def is_bf16(precision): + return precision in BFLOAT_LIST + + @staticmethod + def to_dtype(precision): + if precision in HALF_LIST: + return torch.float16 + elif precision in FLOAT_LIST: + return torch.float32 + elif precision in BFLOAT_LIST: + return torch.bfloat16 + else: + raise RuntimeError(f"unexpected precision: {precision}") + + @staticmethod + def to_str(precision): + if precision == torch.float16: + return "fp16" + elif precision == torch.float32: + return "fp32" + elif precision == torch.bfloat16: + return "bf16" + else: + raise RuntimeError(f"unexpected precision: {precision}") diff --git a/verl/verl/utils/torch_functional.py b/verl/verl/utils/torch_functional.py new file mode 100644 index 0000000000000000000000000000000000000000..8666bec2d1607c1586c759a1f895ef3c0f0c069f --- /dev/null +++ b/verl/verl/utils/torch_functional.py @@ -0,0 +1,1028 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Contain small torch utilities +""" + +import math +from contextlib import contextmanager +from typing import Optional + +import torch +import torch.distributed +import torch.nn.functional as F +from tensordict import TensorDict +from torch import nn +from torch.optim import Optimizer +from torch.optim.lr_scheduler import LambdaLR +from transformers import PreTrainedTokenizer + +from verl.utils.device import get_device_name, get_torch_device + +try: + from flash_attn.ops.triton.cross_entropy import cross_entropy_loss + + FLAH_ATTN_CROSS_ENTROPY_LOSS_AVAILABLE = True +except ImportError: + FLAH_ATTN_CROSS_ENTROPY_LOSS_AVAILABLE = False + + +try: + import torch_npu + + NPU_CROSS_ENTROPY_LOSS_AVAILABLE = hasattr(torch_npu, "npu_cross_entropy_loss") +except ImportError: + NPU_CROSS_ENTROPY_LOSS_AVAILABLE = False + + +def gather_from_labels(data: torch.Tensor, label: torch.Tensor) -> torch.Tensor: + """Gather values from data tensor at positions specified by label indices. + + Selects elements from the last dimension of `data` based on indices in `label`. + Commonly used to extract log-probabilities for specific token IDs from a + vocabulary distribution. + + Args: + data: Input tensor of shape (..., vocab_size) containing values to gather from. + label: Index tensor of shape (...,) with values in range [0, vocab_size). + + Returns: + torch.Tensor: Gathered values with shape (...,), same as label shape. + + Example: + >>> logits = torch.randn(2, 3, 100) # [batch, seq, vocab] + >>> labels = torch.randint(0, 100, (2, 3)) # [batch, seq] + >>> gathered = gather_from_labels(logits, labels) # [batch, seq] + """ + output = torch.gather(data, -1, label.unsqueeze(-1)).squeeze(-1) + return output + + +def logprobs_from_logits(logits, labels, inplace_backward=True): + """ + Compute per-token log-probabilities for the given labels. + + Uses a Flash-Attention–based cross-entropy (if available) for efficient backward, + otherwise falls back to a standard log-softmax+gather approach. + + See: https://github.com/pytorch/pytorch/issues/563#issuecomment-330103591 + + Args: + logits (Tensor): Model outputs of shape (..., vocab_size). + labels (LongTensor): True class indices of shape matching logits[..., :-1]. + inplace_backward (bool): If True and Flash-Attn is available, perform backward in-place. + + Returns: + Tensor: Log-probabilities of the target labels, shape logits.shape[:-1]. + """ + if FLAH_ATTN_CROSS_ENTROPY_LOSS_AVAILABLE: + batch_dim = logits.shape[:-1] + last_dim = logits.shape[-1] + logits = logits.reshape(-1, last_dim) + labels = labels.reshape(-1) + output = logprobs_from_logits_flash_attn(logits, labels, inplace_backward=inplace_backward) + output = output.view(*batch_dim) + elif NPU_CROSS_ENTROPY_LOSS_AVAILABLE: + output = logprobs_from_logits_torch_npu(logits, labels) + else: + output = logprobs_from_logits_v2(logits, labels) + return output + + +def logprobs_from_logits_flash_attn( + logits: torch.Tensor, labels: torch.Tensor, inplace_backward: bool = True +) -> torch.Tensor: + """Compute log-probabilities using Flash Attention's optimized cross-entropy. + + Uses the Flash Attention library's Triton-based cross-entropy implementation + for efficient computation on NVIDIA GPUs. + + Args: + logits: Model output logits of shape (batch_size, vocab_size). + labels: Target token indices of shape (batch_size,). + inplace_backward: If True, perform backward pass in-place for memory efficiency. + + Returns: + torch.Tensor: Log-probabilities for target labels, shape (batch_size,). + + Raises: + AssertionError: If flash-attn version < 2.4.3 (different return format). + """ + output = cross_entropy_loss(logits, labels, inplace_backward=inplace_backward) + assert isinstance(output, tuple), ( + "please make sure flash-attn>=2.4.3 where cross_entropy_loss returns Tuple[losses, z_losses]." + ) + return -output[0] + + +def logprobs_from_logits_torch_npu(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: + """Compute log-probabilities using Ascend NPU's optimized cross-entropy. + + Uses torch_npu's native cross-entropy implementation for efficient + computation on Huawei Ascend NPU devices. + + Args: + logits: Model output logits of shape (..., vocab_size). + labels: Target token indices of shape (...,). + + Returns: + torch.Tensor: Log-probabilities for target labels, same shape as labels. + """ + batch_dim = logits.shape[:-1] + logits = logits.reshape(-1, logits.shape[-1]) + loss, _, _, _ = torch_npu.npu_cross_entropy_loss(logits, labels.reshape(-1), reduction="none") + return -loss.view(*batch_dim) + + +def logprobs_from_logits_naive(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: + """Compute log-probabilities using standard log-softmax approach. + + Simple implementation using PyTorch's log_softmax followed by gathering. + Less memory-efficient than specialized implementations but works on all devices. + + Args: + logits: Model output logits of shape (..., vocab_size). + labels: Target token indices of shape (...,). + + Returns: + torch.Tensor: Log-probabilities for target labels, same shape as labels. + """ + logp = F.log_softmax(logits, dim=-1) + logpy = gather_from_labels(logp, labels) + return logpy + + +def logprobs_from_logits_v2(logits: torch.FloatTensor, labels: torch.Tensor) -> torch.Tensor: + """Memory-efficient log-probability computation using row-wise processing. + + Computes log-probabilities by processing one row at a time to reduce peak + memory consumption. Uses logsumexp for float32/float64, falls back to + log_softmax for bfloat16 due to numerical stability concerns. + + The mathematical identity used is: log_softmax(x_i) = x_i - logsumexp(x) + + Args: + logits: Model output logits of shape (batch_size, seq_len, vocab_size) + or (batch_size, vocab_size). + labels: Target token indices matching logits shape without vocab dimension. + + Returns: + torch.Tensor: Log-probabilities for target labels. + + Note: + This implementation trades compute for memory by iterating over batch + dimension, making it suitable for large vocabulary sizes. + """ + if logits.dtype in [torch.float32, torch.float64]: + logits_labels = torch.gather(logits, dim=-1, index=labels.unsqueeze(-1)).squeeze(-1) + # loop to reduce peak mem consumption + logsumexp_values = torch.stack([torch.logsumexp(logit, dim=-1) for logit in logits]) + logprobs_labels = logits_labels - logsumexp_values # log_softmax(x_i) = x_i - logsumexp(x) + else: + # logsumexp approach is unstable with bfloat16, fall back to slightly less efficent approach + logprobs_labels = [] + for row_logits, row_labels in zip(logits, labels, strict=True): # loop to reduce peak mem consumption + row_logprobs = F.log_softmax(row_logits, dim=-1) + row_logprobs_labels = row_logprobs.gather(dim=-1, index=row_labels.unsqueeze(-1)).squeeze(-1) + logprobs_labels.append(row_logprobs_labels) + logprobs_labels = torch.stack(logprobs_labels) + return logprobs_labels + + +def clip_by_value(x: torch.Tensor, tensor_min: torch.Tensor, tensor_max: torch.Tensor) -> torch.Tensor: + """Clip tensor values to a range defined by tensor bounds. + + Extension of torch.clamp that supports tensor-valued min/max bounds + instead of only scalar bounds. + + Args: + x: Input tensor to clip. + tensor_min: Minimum bound tensor (broadcastable to x). + tensor_max: Maximum bound tensor (broadcastable to x). + + Returns: + torch.Tensor: Clipped tensor with values in [tensor_min, tensor_max]. + + See Also: + https://github.com/pytorch/pytorch/issues/2793#issuecomment-428784713 + """ + clipped = torch.max(torch.min(x, tensor_max), tensor_min) + return clipped + + +def entropy_from_logits(logits: torch.Tensor) -> torch.Tensor: + """Calculate Shannon entropy from unnormalized logits. + + Computes H(p) = -sum(p * log(p)) using the numerically stable formula: + entropy = logsumexp(logits) - sum(softmax(logits) * logits) + + Args: + logits: Unnormalized log-probabilities of shape (..., vocab_size). + + Returns: + torch.Tensor: Entropy values with shape (...,), one per distribution. + """ + pd = torch.nn.functional.softmax(logits, dim=-1) + entropy = torch.logsumexp(logits, dim=-1) - torch.sum(pd * logits, dim=-1) + return entropy + + +def entropy_from_logits_with_chunking(logits: torch.Tensor, chunk_size: int = 2048) -> torch.Tensor: + """Memory-efficient entropy calculation using chunked processing. + + Computes entropy by processing the batch in chunks to reduce peak memory + usage. Useful for large batch sizes or when memory is constrained. + + Args: + logits: Unnormalized log-probabilities of shape (batch_size, vocab_size). + chunk_size: Number of samples to process at once. Defaults to 2048. + + Returns: + torch.Tensor: Entropy values with shape (batch_size,). + + Note: + Converts chunks to float32 for numerical stability during computation. + """ + entropy = torch.zeros(logits.shape[0], device=logits.device) + for i in range(0, logits.shape[0], chunk_size): + logits_chunk = logits[i : i + chunk_size].float() + pd_chunk = torch.nn.functional.softmax(logits_chunk, dim=-1) + entropy_chunk = torch.logsumexp(logits_chunk, dim=-1) - torch.sum(pd_chunk * logits_chunk, dim=-1) + entropy[i : i + chunk_size] = entropy_chunk + return entropy + + +def masked_sum(values: torch.Tensor, mask: torch.Tensor, axis: int | tuple[int, ...] | None = None) -> torch.Tensor: + """Compute sum of tensor values where mask is True. + + NaN values outside the mask are replaced with zeros to prevent + contaminating the sum. + + Args: + values: Input tensor containing values to sum. + mask: Boolean or numeric mask tensor (same shape as values). + Non-zero values indicate elements to include. + axis: Dimension(s) along which to sum. None sums all elements. + + Returns: + torch.Tensor: Sum of masked values, reduced along specified axis. + """ + # If NaNs exist out of mask, replace NaNs in values with a value that + # won't affect the sum (e.g., 0 for masked regions) + valid_values = torch.where(mask.bool(), values, 0.0) + return (valid_values * mask).sum(axis=axis) + + +def masked_mean(values, mask, axis=None): + """ + Compute the mean of `values` over elements selected by `mask`. + + Args: + values (Tensor): Input tensor. + mask (Tensor): Boolean or numeric mask of the same shape as `values`. + axis (int or tuple of int, optional): Dimension(s) along which to compute the mean. + Defaults to None (over all elements). + + Returns: + Tensor: Masked mean, with shape equal to `values` reduced over `axis`. + """ + s = masked_sum(values, mask, axis) + return s / (mask.sum(axis=axis) + 1e-8) + + +def masked_var(values, mask, unbiased=True): + """Compute variance of tensor with masked values.""" + mean = masked_mean(values, mask) + centered_values = values - mean + variance = masked_mean(centered_values**2, mask) + if unbiased: + mask_sum = mask.sum() + if mask_sum == 0: + raise ValueError("At least one element in the mask has to be 1.") + # note that if mask_sum == 1, then there is a division by zero issue + # to avoid it you just need to use a larger minibatch_size + if mask_sum == 1: + raise ValueError("The sum of the mask is one, which can cause a division by zero.") + bessel_correction = mask_sum / (mask_sum - 1) + variance = variance * bessel_correction + return variance + + +def masked_whiten(values, mask, shift_mean=True): + """ + Whiten `values` by normalizing with mean and variance computed over `mask`. + + Args: + values (torch.Tensor): Input tensor. + mask (torch.Tensor): Boolean tensor of same shape, selects elements for stats. + shift_mean (bool): If True (default), output is zero-mean; + if False, the original mean is re-added after scaling. + + Returns: + torch.Tensor: Whitened tensor of same shape as `values`. + """ + mean, var = masked_mean(values, mask), masked_var(values, mask) + whitened = (values - mean) * torch.rsqrt(var + 1e-8) + if not shift_mean: + whitened += mean + return whitened + + +def get_response_mask(response_id: torch.Tensor, eos_token: int | list[int] = 2, dtype=torch.int64): + """ + end of sentence token can be int or list: 1 or [1, 2] + e.g. + response_id = torch.tensor([[20, 10, 34, 1, 0, 0, 0], + [78, 0, 76, 2, 1, 0, 0], + [23, 98, 1, 0, 0, 0, 0], + [33, 3, 98, 45, 1, 0, 0]]) + #eos_token=1 + response_mask: tensor([[1, 1, 1, 1, 0, 0, 0], + [1, 1, 1, 1, 1, 0, 0], + [1, 1, 1, 0, 0, 0, 0], + [1, 1, 1, 1, 1, 0, 0]]) + #eos_token=[1,2] + response_mask: tensor([[1, 1, 1, 1, 0, 0, 0], + [1, 1, 1, 1, 0, 0, 0], + [1, 1, 1, 0, 0, 0, 0], + [1, 1, 1, 1, 1, 0, 0]]) + """ + eos_mask = torch.isin(response_id, torch.tensor(eos_token, device=response_id.device)).int() + return (eos_mask.cumsum(dim=1) - eos_mask).eq(0).to(dtype) + + +def compute_grad_norm(model: nn.Module) -> float: + """Compute the squared L2 norm of all gradients in a model. + + Sums the squared values of all gradient tensors across all parameters. + Useful for monitoring gradient magnitudes during training. + + Args: + model: PyTorch model with computed gradients. + + Returns: + float: Sum of squared gradient values (not the square root). + + Note: + Returns the squared norm, not the norm itself. To get the actual + L2 norm, take the square root of the returned value. + """ + total_grad_square = 0 + for param in model.parameters(): + if param.grad is not None: + total_grad_square += torch.sum(torch.square(param.grad.detach())).item() + return total_grad_square + + +def broadcast_dict_tensor(tensors: dict[str, torch.Tensor] | TensorDict, src: int, group) -> None: + """Broadcast all tensors in a dictionary from source rank to all ranks. + + Iterates over all tensors in the dictionary and broadcasts each one + from the source rank to all other ranks in the process group. + + Args: + tensors: Dictionary or TensorDict containing tensors to broadcast. + src: Source rank from which to broadcast. + group: Process group for the broadcast operation. + + Note: + This implementation broadcasts tensors one at a time. Could be optimized + to use a single broadcast with packed tensors. + """ + for key in tensors.sorted_keys: + torch.distributed.broadcast(tensors[key], src=src, group=group, async_op=False) + + +def allgather_dict_tensors( + tensors: dict[str, torch.Tensor] | TensorDict, size: int, group, dim: int = 0 +) -> dict[str, torch.Tensor] | TensorDict: + """Gather tensors from all ranks and concatenate them. + + Performs all_gather on each tensor in the dictionary and concatenates + the results along the specified dimension. + + Args: + tensors: Dictionary or TensorDict containing tensors to gather. + size: Number of ranks in the process group. + group: Process group for the all_gather operation. + dim: Dimension along which to concatenate gathered tensors. Defaults to 0. + + Returns: + Dictionary or TensorDict (matching input type) with gathered and + concatenated tensors. Each tensor's size along `dim` is multiplied by `size`. + + Note: + This implementation gathers tensors one at a time synchronously. + Could be optimized using async ops or packed all_gather. + """ + if isinstance(tensors, TensorDict): + is_tensor_dict = True + tensors_as_dict = tensors.to_dict() + else: + tensors_as_dict = tensors + is_tensor_dict = False + + output = {} + sorted_keys = sorted(tensors_as_dict.keys()) + for key in sorted_keys: + val = tensors_as_dict[key] + output[key] = [torch.empty_like(val) for _ in range(size)] + torch.distributed.all_gather(output[key], val, group=group, async_op=False) + output[key] = torch.cat(output[key], dim=dim) + + if is_tensor_dict: + output = TensorDict(source=output, batch_size=tensors.batch_size[0] * size) + + return output + + +def allgather_dict_into_dict(data: dict, group=None) -> dict: + """allgather a dict into a dict of list + + Args: + data: a dict + group: the process group to allgather + + Returns: dict containing a list of the results from allgather + + """ + assert isinstance(data, dict), f"Expect data to be a dictionary, Got {type(data)}" + + group_size = torch.distributed.get_world_size(group=group) + + final_metrics = {} + all_metrics_lst = [None for _ in range(group_size)] + torch.distributed.all_gather_object(all_metrics_lst, data, group=group) + + for all_metrics in all_metrics_lst: + for key, val in all_metrics.items(): + if key not in final_metrics: + final_metrics[key] = [] + final_metrics[key].append(val) + return final_metrics + + +def split_dict_tensor_into_batches(tensors: TensorDict, batch_size) -> list[TensorDict]: + assert tensors.batch_size[0] % batch_size == 0, ( + f"input data batch size: {tensors.batch_size[0]}, split batch size: {batch_size}" + ) + return tensors.split(batch_size) + + +def pad_2d_list_to_length(response, pad_token_id, max_length=None): + """ + pad a 2D list (e.g. responses, logprobs) to a 2D tensor. + """ + response_length = max(len(sub_list) for sub_list in response) + target_length = max_length if max_length is not None and max_length > response_length else response_length + padded_response = [tuple(sub_list) + (pad_token_id,) * (target_length - len(sub_list)) for sub_list in response] + tensor = torch.tensor(padded_response) + return tensor + + +def pad_sequence_to_length(tensors, max_seq_len, pad_token_id, left_pad=False): + """ + pad a 2D tensors (e.g. responses, logprobs) in the last dim to max_seq_length. + input shape: [bs, seq_length] + output shape: [bs, max_seq_length] + """ + if tensors.shape[-1] >= max_seq_len: + return tensors + # (0, max_seq_len - tensors.shape[-1]) means right pad to max_seq_length and no left pad + pad_tuple = (max_seq_len - tensors.shape[-1], 0) if left_pad else (0, max_seq_len - tensors.shape[-1]) + return F.pad(tensors, pad_tuple, "constant", pad_token_id) + + +def postprocess_data( + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + max_length: int, + pad_token_id: int, + left_pad=True, + truncation="error", +): + """Process tokenizer outputs to consistent shapes via padding/truncation. + + Args: + input_ids: Token indices [batch_size, seq_len] + attention_mask: Mask [batch_size, seq_len] + max_length: Target sequence length + pad_token_id: Padding token ID + left_pad: Pad left if True + truncation: "left", "right", "middle" or "error" + + Returns: + (input_ids, attention_mask) padded/truncated to max_length + """ + assert truncation in ["left", "right", "middle", "error"] + assert input_ids.ndim == 2 + + sequence_length = input_ids.shape[-1] + if sequence_length < max_length: + input_ids = pad_sequence_to_length( + input_ids, max_seq_len=max_length, pad_token_id=pad_token_id, left_pad=left_pad + ) + attention_mask = pad_sequence_to_length( + attention_mask, max_seq_len=max_length, pad_token_id=0, left_pad=left_pad + ) + elif sequence_length > max_length: + if truncation == "left": + # actually, left truncation may not be reasonable + input_ids = input_ids[:, -max_length:] + attention_mask = attention_mask[:, -max_length:] + elif truncation == "right": + input_ids = input_ids[:, :max_length] + attention_mask = attention_mask[:, :max_length] + elif truncation == "middle": + left_half = max_length // 2 + right_half = max_length - left_half + input_ids = torch.cat([input_ids[:, :left_half], input_ids[:, -right_half:]], dim=-1) + attention_mask = torch.cat([attention_mask[:, :left_half], attention_mask[:, -right_half:]], dim=-1) + elif truncation == "error": + raise NotImplementedError(f"{sequence_length=} is larger than {max_length=}") + else: + raise NotImplementedError(f"Unknown truncation method {truncation}") + + return input_ids, attention_mask + + +def tokenize_and_postprocess_data( + prompt: str, tokenizer: PreTrainedTokenizer, max_length: int, pad_token_id: int, left_pad=True, truncation="error" +): + """Tokenize text and process outputs to consistent tensor shapes. + + Args: + prompt: Input text to tokenize + tokenizer: HuggingFace tokenizer instance + max_length: Target sequence length + pad_token_id: Padding token ID + left_pad: Pad left if True + truncation: Truncation strategy ("left"/"right"/"error") + + Returns: + Tuple of (input_ids, attention_mask) from postprocess_data + """ + input_data = tokenizer(prompt, return_tensors="pt", add_special_tokens=False) + input_ids = input_data["input_ids"] + attention_mask = input_data["attention_mask"] + + return postprocess_data(input_ids, attention_mask, max_length, pad_token_id, left_pad, truncation) + + +def remove_pad_token(input_ids: torch.Tensor, attention_mask: torch.Tensor): + """Remove the pad token. + + Args: + input_ids shape: [bs, seq_length] + attention_mask shape: [bs, seq_length] + Returns: + no_padding_batch(List[List[int]]): contains the rmpad token ids per query. + """ + no_padding_batch = [] + for ids, mask in zip(input_ids, attention_mask, strict=True): + no_padding_batch.append((ids[len(ids) - mask.sum() :]).cpu().numpy().tolist()) + return no_padding_batch + + +def log_probs_from_logits_response(input_ids, logits, response_length): + """Compute the response log_probs from full logits. Note that logits = model(input_ids) + + Args: + input_ids: [batch_size, seqlen] + logits: [batch_size, seqlen, vocab_size] + + Returns: + response_log_prob: + """ + response_logits = logits[:, -response_length - 1 : -1] + response = input_ids[:, -response_length:] + response_log_prob = logprobs_from_logits(logits=response_logits, labels=response) + return response_log_prob + + +def log_probs_from_logits_response_rmpad(input_ids, attention_mask, logits_rmpad, response_length): + """Compute the log_probs from logits with rmpad logits and pad input. Note that + logits_rmpad = model(input_ids_rmpad). For each sentences, there is a shift between + logits and input_ids. + The reason for this function to is to compute logprobs_from_logits in rmpad mode because it is memory-intensive + for large vocab_size + + Args: + input_ids: [batch_size, seqlen] + attention_mask: [batch_size, seqlen] + logits_rmpad: [total_nnz, vocab_size] + response_length: int + """ + from flash_attn.bert_padding import pad_input, unpad_input + + batch_size, seqlen = input_ids.shape + input_ids_rmpad, indices, *_ = unpad_input(input_ids.unsqueeze(-1), attention_mask=attention_mask) + input_ids_rmpad = input_ids_rmpad.squeeze(-1) + input_ids_rmpad_rolled = torch.roll(input_ids_rmpad, shifts=-1, dims=0) + full_log_probs_rmpad = logprobs_from_logits(logits=logits_rmpad, labels=input_ids_rmpad_rolled) # (total_nnz,) + full_output = pad_input( + hidden_states=full_log_probs_rmpad.unsqueeze(-1), indices=indices, batch=batch_size, seqlen=seqlen + ) + output = full_output.squeeze(-1)[:, -response_length - 1 : -1] # [batch_size, response_length] + return output + + +def log_probs_from_logits_all_rmpad(input_ids_rmpad, logits_rmpad, indices, batch_size, seqlen, response_length): + """Compute the log_probs from logits with rmpad input_ids and logits. Note that + logits_rmpad = model(input_ids_rmpad). For each sentences, there is a shift between + logits and input_ids. + The reason for this function to is to compute logprobs_from_logits in rmpad mode because it is memory-intensive + for large vocab_size + + Args: + input_ids_rmpad: [1, total_nnz] + logits_rmpad: [total_nnz, vocab_size] + indices: [total_nnz] + batch_size: int + seqlen: int + response_length: int + """ + if get_device_name() == "cuda": + from flash_attn.bert_padding import pad_input + elif get_device_name() == "npu": + from verl.utils.attention_utils import pad_input + + input_ids_rmpad = input_ids_rmpad.transpose(0, 1) # transpose back to [total_nnz, 1] + input_ids_rmpad = input_ids_rmpad.squeeze(-1) + input_ids_rmpad_rolled = torch.roll(input_ids_rmpad, shifts=-1, dims=0) + full_log_probs_rmpad = logprobs_from_logits(logits=logits_rmpad, labels=input_ids_rmpad_rolled) # (total_nnz,) + full_output = pad_input( + hidden_states=full_log_probs_rmpad.unsqueeze(-1), indices=indices, batch=batch_size, seqlen=seqlen + ) + output = full_output.squeeze(-1)[:, -response_length - 1 : -1] # [batch_size, response_length] + return output + + +def post_process_logits(input_ids, logits, temperature, top_k, top_p): + if temperature != 1.0: + logits = logits.div_(temperature) # inplace operation to avoid OOM + # TODO: add them back + # if top_k is not None and top_k > 0: + # logits = TopKLogitsWarper(top_k=top_k)(input_ids, logits) + # if top_p is not None and top_p < 1.0 and top_p > 0.0: + # logits = TopPLogitsWarper(top_p=top_p)(input_ids, logits) + return logits + + +def calculate_sum_pi_squared_from_logits(logits: torch.Tensor): + """ + Compute exact sum of squared probabilities from logits. + Formula: Σπ² = exp(logsumexp(2*logits) - 2*logsumexp(logits)) + + Used for optimal baseline variance reduction as described in + "What Matters for Model Merging at Scale?" (arXiv:2410.03617) + + Args: + logits: Logits tensor (..., vocab_size). + + Returns: + Sum of squared probabilities tensor (...). + """ + return torch.exp(torch.logsumexp(2.0 * logits, dim=-1) - 2.0 * torch.logsumexp(logits, dim=-1)) + + +""" +Optimizer related +""" + + +def get_cosine_schedule_with_warmup( + optimizer: Optimizer, + num_warmup_steps: int, + num_training_steps: int, + min_lr_ratio: float = 0.0, + num_cycles: float = 0.5, + last_epoch: int = -1, + init_lr_ratio: float = None, + zero_indexed_step: bool = True, +): + """ + Create a schedule with a learning rate that decreases following the values of the cosine function between the + initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the + initial lr set in the optimizer. + Args: + optimizer (:class:`~torch.optim.Optimizer`): + The optimizer for which to schedule the learning rate. + num_warmup_steps (:obj:`int`): + The number of steps for the warmup phase. + num_training_steps (:obj:`int`): + The total number of training steps. + min_lr_ratio (:obj:`float`, `optional`, defaults to 0.0): + The minimum lr ratio w.r.t the maximum. + num_cycles (:obj:`float`, `optional`, defaults to 0.5): + The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0 + following a half-cosine). + last_epoch (:obj:`int`, `optional`, defaults to -1): + The index of the last epoch when resuming training. + init_lr_ratio (:obj:`float`, `optional`, defaults to None): + The initial lr ratio w.r.t the maximum. + zero_indexed_step (:obj:`bool`, `optional`, defaults to True): + Whether the LR schedule uses 0-indexed steps. If True (default), step counting starts at 0. + If False (used by torchtitan), step counting starts at 1. + Return: + :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + min_lr_ratio = 0.0 if min_lr_ratio is None else min_lr_ratio + assert min_lr_ratio >= 0 and min_lr_ratio <= 1.0 + coef = (1 - min_lr_ratio) * 0.5 + intercept = (1 + min_lr_ratio) * 0.5 + + init_lr_ratio = 0.0 if init_lr_ratio is None else init_lr_ratio + assert init_lr_ratio >= 0 and init_lr_ratio <= 1.0 + + def lr_lambda(current_step): + if not zero_indexed_step: + current_step += 1 + if current_step < num_warmup_steps: + return init_lr_ratio + (1.0 - init_lr_ratio) * (float(current_step) / float(max(1, num_warmup_steps))) + progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps)) + x = math.cos(math.pi * float(num_cycles) * 2.0 * progress) + return max(min_lr_ratio, x * coef + intercept) + + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +def get_constant_schedule_with_warmup( + optimizer: Optimizer, + num_warmup_steps: int, + last_epoch: int = -1, +): + """ + Create a constant LR schedule with a linear warmup phase. + + Args: + optimizer (Optimizer): Wrapped optimizer. + num_warmup_steps (int): Number of steps to ramp up the LR from 0 to initial value. + last_epoch (int, optional): The index of the last epoch when resuming training. Defaults to -1. + + Returns: + LambdaLR: Scheduler that increases LR linearly during warmup, then holds it constant. + """ + + def lr_lambda(current_step): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1.0, num_warmup_steps)) + return 1.0 + + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +def prepare_decoder_attention_mask(attention_mask, input_shape, inputs_embeds): + # create causal mask + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + combined_attention_mask = None + if input_shape[-1] > 1: + combined_attention_mask = _make_causal_mask( + input_shape, + inputs_embeds.dtype, + device=inputs_embeds.device, + ) + + if attention_mask is not None: + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to( + inputs_embeds.device + ) + combined_attention_mask = ( + expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask + ) + + return combined_attention_mask + + +# Copied from transformers.models.bart.modeling_bart._make_causal_mask +def _make_causal_mask(input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device): + """ + Make causal mask used for bi-directional self-attention. + """ + bsz, tgt_len = input_ids_shape + mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device) + mask_cond = torch.arange(mask.size(-1), device=device) + mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0) + mask = mask.to(dtype) + return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len) + + +# Copied from transformers.models.bart.modeling_bart._expand_mask +def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None): + """ + Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`. + """ + bsz, src_len = mask.size() + tgt_len = tgt_len if tgt_len is not None else src_len + + expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype) + + inverted_mask = 1.0 - expanded_mask + + return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min) + + +def get_unpad_data(attention_mask): + seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = seqlens_in_batch.max().item() + cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) + return ( + indices, + cu_seqlens, + max_seqlen_in_batch, + ) + + +def get_wsd_schedule_with_warmup( + optimizer: Optimizer, + num_warmup_steps: int, + num_training_steps: int, + min_lr_ratio: float = 0.0, + num_cycles: float = 0.5, + last_epoch: int = -1, + stable_ratio: float = 0.9, +): + """ + Create a Warmup-Stable-Decay learning rate scheduler. + + The schedule follows three phases: + 1. Warmup: Learning rate increases linearly from 0 to the initial LR + 2. Stable: Learning rate remains constant at the initial LR + 3. Decay: Learning rate decreases following a cosine curve to min_lr_ratio * initial LR + + Args: + optimizer (:class:`~torch.optim.Optimizer`): + The optimizer for which to schedule the learning rate. + num_warmup_steps (:obj:`int`): + The number of steps for the warmup phase. + num_training_steps (:obj:`int`): + The total number of training steps. + min_lr_ratio (:obj:`float`, `optional`, defaults to 0.0): + The minimum learning rate ratio w.r.t the initial learning rate. + num_cycles (:obj:`float`, `optional`, defaults to 0.5): + The number of waves in the cosine schedule during decay phase. + last_epoch (:obj:`int`, `optional`, defaults to -1): + The index of the last epoch when resuming training. + stable_ratio (:obj:`float`, `optional`, defaults to 0.0): + The ratio of non-warmup steps that should maintain a constant learning rate. + Set to 0.0 to behave exactly like cosine schedule. + + Return: + :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + remaining_steps = max(0, num_training_steps - num_warmup_steps) + num_stable_steps = int(remaining_steps * stable_ratio) + num_decay_steps = remaining_steps - num_stable_steps + + def lr_lambda(current_step): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1, num_warmup_steps)) + if current_step < num_warmup_steps + num_stable_steps: + return 1.0 + if current_step < num_training_steps: + progress = float(current_step - num_warmup_steps - num_stable_steps) / float(max(1, num_decay_steps)) + value = max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress))) + return (1.0 - min_lr_ratio) * value + min_lr_ratio + return min_lr_ratio + + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +@contextmanager +def check_device_is_available(): + """ + Some modules must be imported after CUDA is initialized. Such as sglang's sharding manager. + + This context manager checks if CUDA is available and raises an error if it is not. + """ + if not get_torch_device().is_available(): + raise RuntimeError("Device {} must be initialized before importing this module.".format(get_device_name())) + + yield + + +def distributed_mean_max_min_std(local_tensor, compute_max=True, compute_min=True, compute_std=True): + """Compute distributed statistics across all processes. + + Args: + local_tensor: Tensor containing local values + compute_max: Include maximum value calculation + compute_min: Include minimum value calculation + compute_std: Include standard deviation calculation + + Returns: + Tuple containing (mean, max, min, std) in this order. None for disabled metrics. + """ + # Sum the local tensor across all processes + local_sum = torch.sum(local_tensor) + local_num = torch.tensor(torch.numel(local_tensor), device=get_device_name()) + + torch.distributed.all_reduce(local_sum, op=torch.distributed.ReduceOp.SUM) + torch.distributed.all_reduce(local_num, op=torch.distributed.ReduceOp.SUM) + + global_mean = local_sum / local_num + + if compute_max: + local_max = torch.max(local_tensor) + torch.distributed.all_reduce(local_max, op=torch.distributed.ReduceOp.MAX) + else: + local_max = None + + if compute_min: + local_min = torch.min(local_tensor) + torch.distributed.all_reduce(local_min, op=torch.distributed.ReduceOp.MIN) + else: + local_min = None + + if compute_std: + square_diff = torch.sum(torch.pow(local_tensor - global_mean, 2)) + torch.distributed.all_reduce(square_diff, op=torch.distributed.ReduceOp.SUM) + global_std = torch.sqrt(square_diff / (local_num - 1)) + else: + global_std = None + + return global_mean, local_max, local_min, global_std + + +def distributed_masked_mean(local_tensor, local_mask): + """Compute global mean of non-masked elements across distributed processes. + + Args: + local_tensor (torch.Tensor): Input tensor with local values + local_mask (torch.Tensor): Binary mask (1=valid, 0=ignore) matching local_tensor shape + + Returns: + torch.Tensor: Global mean of all valid elements across processes + """ + local_tensor = local_tensor * local_mask + + local_sum = torch.sum(local_tensor) + local_num = torch.sum(local_mask) + + torch.distributed.all_reduce(local_sum, op=torch.distributed.ReduceOp.SUM) + torch.distributed.all_reduce(local_num, op=torch.distributed.ReduceOp.SUM) + + global_mean = local_sum / local_num + return global_mean + + +def expand_as_nested(tensor: torch.Tensor, nested_tensor: torch.Tensor) -> torch.Tensor: + """ + + Args: + tensor: a tensor with shape (bsz,) + nested_tensor: a nested tensor with shape (bsz, xxx) + + Returns: + a tensor with the same shape as nested_tensor + + """ + assert nested_tensor.is_nested, "nested_tensor must be nested" + assert tensor.shape[0] == nested_tensor.shape[0], ( + f"The batch shape must be the same. Got {tensor.shape[0]} vs {nested_tensor.shape[0]}" + ) + assert len(tensor.shape) == 1, "The ndim of tensor must be 1" + assert len(nested_tensor.shape) == 2, "The ndim of nested_tensor must be 2" + + offsets = nested_tensor.offsets() + seqlens = offsets.diff() + output = torch.repeat_interleave(tensor, seqlens, dim=0) + output = torch.nested.nested_tensor_from_jagged(values=output, offsets=offsets) + return output + + +@contextmanager +def use_original_torch_compile(): + """torch.compile might be replaced by mindspeed on NPU, this contextmanager + can revert torch.compile temporarily. + """ + try: + from mindspeed.patch_utils import MindSpeedPatchesManager + + compile_patch = None + for patch in MindSpeedPatchesManager.patches_info.values(): + if patch.orig_module_name == "torch" and patch.orig_func_name == "compile": + if patch.is_applied(): + compile_patch = patch + break + if compile_patch is not None: + compile_patch.remove_patch() + yield + compile_patch.apply_patch() + else: + yield + except Exception: + yield diff --git a/verl/verl/utils/tracking.py b/verl/verl/utils/tracking.py new file mode 100644 index 0000000000000000000000000000000000000000..a2f5d142dcc80c5819189bc1d0d3c35a6e405364 --- /dev/null +++ b/verl/verl/utils/tracking.py @@ -0,0 +1,545 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +A unified tracking interface that supports logging data to different backend +""" + +import dataclasses +import json +import logging +import os +from enum import Enum +from functools import partial +from pathlib import Path +from typing import Any + +import orjson + +logger = logging.getLogger(__name__) + +MLFLOW_MAX_ATTEMPTS = 3 +MLFLOW_SLEEP_SECONDS = 5 + + +class Tracking: + """A unified tracking interface for logging experiment data to multiple backends. + + This class provides a centralized way to log experiment metrics, parameters, and artifacts + to various tracking backends including WandB, MLflow, SwanLab, TensorBoard, and console. + + Attributes: + supported_backend: List of supported tracking backends. + logger: Dictionary of initialized logger instances for each backend. + """ + + supported_backend = [ + "wandb", + "mlflow", + "swanlab", + "vemlp_wandb", + "tensorboard", + "console", + "clearml", + "trackio", + "file", + ] + + def __init__(self, project_name, experiment_name, default_backend: str | list[str] = "console", config=None): + if isinstance(default_backend, str): + default_backend = [default_backend] + for backend in default_backend: + if backend == "tracking": + import warnings + + warnings.warn("`tracking` logger is deprecated. use `wandb` instead.", DeprecationWarning, stacklevel=2) + else: + assert backend in self.supported_backend, f"{backend} is not supported" + + self.logger = {} + + if "tracking" in default_backend or "wandb" in default_backend: + import os + + import wandb + + settings = None + if config and config["trainer"].get("wandb_proxy", None): + settings = wandb.Settings(https_proxy=config["trainer"]["wandb_proxy"]) + entity = os.environ.get("WANDB_ENTITY", None) + wandb.init(project=project_name, name=experiment_name, entity=entity, config=config, settings=settings) + self.logger["wandb"] = wandb + + if "trackio" in default_backend: + import trackio + + trackio.init(project=project_name, name=experiment_name, config=config) + self.logger["trackio"] = trackio + + if "mlflow" in default_backend: + import os + import time + + import mlflow + + for _mlflow_attempt in range(1, MLFLOW_MAX_ATTEMPTS + 1): + try: + MLFLOW_TRACKING_URI = os.environ.get("MLFLOW_TRACKING_URI", "sqlite:////tmp/mlruns.db") + logger.info("Using MLFlow tracking URI: %s", MLFLOW_TRACKING_URI) + mlflow.set_tracking_uri(MLFLOW_TRACKING_URI) + + # Some cloud providers like Azure ML or Databricks automatically set MLFLOW_RUN_ID + # If set, attach to the existing run instead of creating a new one + run_id = os.environ.get("MLFLOW_RUN_ID") + if run_id: + mlflow.start_run(run_id=run_id) + else: + # Project_name is actually experiment_name in MLFlow + # If experiment does not exist, will create a new experiment + experiment = mlflow.set_experiment(project_name) + mlflow.start_run(experiment_id=experiment.experiment_id, run_name=experiment_name) + + mlflow.log_params(_compute_mlflow_params_from_objects(config)) + self.logger["mlflow"] = _MlflowLoggingAdapter() + break # Success + except Exception as e: + logger.warning( + "MLflow initialization attempt %d/%d failed: %s", _mlflow_attempt, MLFLOW_MAX_ATTEMPTS, e + ) + if _mlflow_attempt < MLFLOW_MAX_ATTEMPTS: + time.sleep(MLFLOW_SLEEP_SECONDS) + else: + logger.warning("All MLflow initialization attempts failed. Proceeding without MLflow tracking.") + + if "swanlab" in default_backend: + import os + + import swanlab + + SWANLAB_API_KEY = os.environ.get("SWANLAB_API_KEY", None) + SWANLAB_LOG_DIR = os.environ.get("SWANLAB_LOG_DIR", "swanlog") + SWANLAB_MODE = os.environ.get("SWANLAB_MODE", "cloud") + if SWANLAB_API_KEY: + swanlab.login(SWANLAB_API_KEY) # NOTE: previous login information will be overwritten + + if config is None: + config = {} # make sure config is not None, otherwise **config will raise error + swanlab.init( + project=project_name, + experiment_name=experiment_name, + config={"FRAMEWORK": "verl", **config}, + logdir=SWANLAB_LOG_DIR, + mode=SWANLAB_MODE, + ) + self.logger["swanlab"] = swanlab + + if "vemlp_wandb" in default_backend: + import os + + import volcengine_ml_platform + from volcengine_ml_platform import wandb as vemlp_wandb + + volcengine_ml_platform.init( + ak=os.environ["VOLC_ACCESS_KEY_ID"], + sk=os.environ["VOLC_SECRET_ACCESS_KEY"], + region=os.environ["MLP_TRACKING_REGION"], + ) + + vemlp_wandb.init( + project=project_name, + name=experiment_name, + config=config, + sync_tensorboard=True, + ) + self.logger["vemlp_wandb"] = vemlp_wandb + + if "tensorboard" in default_backend: + self.logger["tensorboard"] = _TensorboardAdapter(project_name, experiment_name) + + if "console" in default_backend: + from verl.utils.logger import LocalLogger + + self.console_logger = LocalLogger(print_to_console=True) + self.logger["console"] = self.console_logger + + if "clearml" in default_backend: + self.logger["clearml"] = ClearMLLogger(project_name, experiment_name, config) + + if "file" in default_backend: + self.logger["file"] = FileLogger(project_name, experiment_name) + + def log(self, data, step, backend=None): + for default_backend, logger_instance in self.logger.items(): + if backend is None or default_backend in backend: + logger_instance.log(data=data, step=step) + + def __del__(self): + if "wandb" in self.logger: + self.logger["wandb"].finish(exit_code=0) + if "swanlab" in self.logger: + self.logger["swanlab"].finish() + if "vemlp_wandb" in self.logger: + self.logger["vemlp_wandb"].finish(exit_code=0) + if "tensorboard" in self.logger: + self.logger["tensorboard"].finish() + if "clearml" in self.logger: + self.logger["clearml"].finish() + if "trackio" in self.logger: + self.logger["trackio"].finish() + if "file" in self.logger: + self.logger["file"].finish() + + +class ClearMLLogger: + def __init__(self, project_name: str, experiment_name: str, config): + self.project_name = project_name + self.experiment_name = experiment_name + + import clearml + + self._task: clearml.Task = clearml.Task.init( + task_name=experiment_name, + project_name=project_name, + continue_last_task=True, + output_uri=False, + ) + + self._task.connect_configuration(config, name="Hyperparameters") + + def _get_logger(self): + return self._task.get_logger() + + def log(self, data, step): + import numpy as np + import pandas as pd + + # logs = self._rewrite_logs(data) + logger = self._get_logger() + for k, v in data.items(): + title, series = k.split("/", 1) + + if isinstance(v, int | float | np.floating | np.integer): + logger.report_scalar( + title=title, + series=series, + value=v, + iteration=step, + ) + elif isinstance(v, pd.DataFrame): + logger.report_table( + title=title, + series=series, + table_plot=v, + iteration=step, + ) + else: + logger.warning( + f'Trainer is attempting to log a value of "{v}" of type {type(v)} for key "{k}". This ' + f"invocation of ClearML logger's function is incorrect so this attribute was dropped. " + ) + + def finish(self): + self._task.close() + + +class FileLogger: + def __init__(self, project_name: str, experiment_name: str): + self.project_name = project_name + self.experiment_name = experiment_name + + self.filepath = os.getenv("VERL_FILE_LOGGER_PATH", None) + if self.filepath is None: + root_path = os.path.expanduser(os.getenv("VERL_FILE_LOGGER_ROOT", ".")) + directory = os.path.join(root_path, self.project_name) + os.makedirs(directory, exist_ok=True) + self.filepath = os.path.join(directory, f"{self.experiment_name}.jsonl") + print(f"Creating file logger at {os.path.abspath(self.filepath)}") + self.fp = open(self.filepath, "wb", buffering=0) + + def log(self, data, step): + data = {"step": step, "data": data} + self.fp.write(orjson.dumps(data, option=orjson.OPT_SERIALIZE_NUMPY) + b"\n") + + def finish(self): + self.fp.close() + + +class _TensorboardAdapter: + def __init__(self, project_name, experiment_name): + import os + + from torch.utils.tensorboard import SummaryWriter + + tensorboard_dir = os.environ.get("TENSORBOARD_DIR", f"tensorboard_log/{project_name}/{experiment_name}") + os.makedirs(tensorboard_dir, exist_ok=True) + print(f"Saving tensorboard log to {tensorboard_dir}.") + self.writer = SummaryWriter(tensorboard_dir) + + def log(self, data, step): + for key in data: + self.writer.add_scalar(key, data[key], step) + + def finish(self): + self.writer.close() + + +class _MlflowLoggingAdapter: + def __init__(self): + import logging + import re + + self.logger = logging.getLogger(__name__) + # Suppress noisy "Found credentials from IAM Role" on every MLflow request + logging.getLogger("botocore.credentials").setLevel(logging.WARNING) + # MLflow metric key validation logic: + # https://github.com/mlflow/mlflow/blob/master/mlflow/utils/validation.py#L157C12-L157C44 + # Only characters allowed: slashes, alphanumerics, underscores, periods, dashes, colons, + # and spaces. + self._invalid_chars_pattern = re.compile( + r"[^/\w.\- :]" + ) # Allowed: slashes, alphanumerics, underscores, periods, dashes, colons, and spaces. + self._consecutive_slashes_pattern = re.compile(r"/+") + self._sanitized_key_cache = {} + + def _sanitize_key(self, key): + if key in self._sanitized_key_cache: + return self._sanitized_key_cache[key] or key + # First replace @ with _at_ for backward compatibility + sanitized = key.replace("@", "_at_") + # Replace consecutive slashes with a single slash (MLflow treats them as file paths) + sanitized = self._consecutive_slashes_pattern.sub("/", sanitized) + # Then replace any other invalid characters with _ + sanitized = self._invalid_chars_pattern.sub("_", sanitized) + if sanitized == key: + self._sanitized_key_cache[key] = None + else: + self.logger.warning("[MLflow] Metric key '%s' sanitized to '%s' due to invalid characters.", key, sanitized) + self._sanitized_key_cache[key] = sanitized + return sanitized + + def log(self, data, step): + import mlflow + + results = {self._sanitize_key(k): v for k, v in data.items()} + for _attempt in range(MLFLOW_MAX_ATTEMPTS): + try: + mlflow.log_metrics(metrics=results, step=step) + return + except Exception as error: + # No sleep between retries — this runs per training step, so we avoid blocking. + msg = "mlflow.log_metrics failed (attempt %d/%d): %s" + args = (_attempt + 1, MLFLOW_MAX_ATTEMPTS, error) + if _attempt < MLFLOW_MAX_ATTEMPTS - 1: + self.logger.info(msg, *args) + else: + self.logger.warning(msg, *args) + + +def _compute_mlflow_params_from_objects(params) -> dict[str, Any]: + if params is None: + return {} + + return _flatten_dict(_transform_params_to_json_serializable(params, convert_list_to_dict=True), sep="/") + + +def _transform_params_to_json_serializable(x, convert_list_to_dict: bool): + _transform = partial(_transform_params_to_json_serializable, convert_list_to_dict=convert_list_to_dict) + + if dataclasses.is_dataclass(x): + return _transform(dataclasses.asdict(x)) + if isinstance(x, dict): + return {k: _transform(v) for k, v in x.items()} + if isinstance(x, list): + if convert_list_to_dict: + return {"list_len": len(x)} | {f"{i}": _transform(v) for i, v in enumerate(x)} + else: + return [_transform(v) for v in x] + if isinstance(x, Path): + return str(x) + if isinstance(x, Enum): + return x.value + + return x + + +def _flatten_dict(raw: dict[str, Any], *, sep: str) -> dict[str, Any]: + import pandas as pd + + ans = pd.json_normalize(raw, sep=sep).to_dict(orient="records")[0] + assert isinstance(ans, dict) + return ans + + +@dataclasses.dataclass +class ValidationGenerationsLogger: + project_name: str = None + experiment_name: str = None + + def log(self, loggers, samples, step): + if "wandb" in loggers: + self.log_generations_to_wandb(samples, step) + if "swanlab" in loggers: + self.log_generations_to_swanlab(samples, step) + if "mlflow" in loggers: + self.log_generations_to_mlflow(samples, step) + + if "clearml" in loggers: + self.log_generations_to_clearml(samples, step) + if "tensorboard" in loggers: + self.log_generations_to_tensorboard(samples, step) + + if "vemlp_wandb" in loggers: + self.log_generations_to_vemlp_wandb(samples, step) + + def log_generations_to_vemlp_wandb(self, samples, step): + from volcengine_ml_platform import wandb as vemlp_wandb + + self._log_generations_to_wandb(samples, step, vemlp_wandb) + + def log_generations_to_wandb(self, samples, step): + import wandb + + self._log_generations_to_wandb(samples, step, wandb) + + def _log_generations_to_wandb(self, samples, step, wandb): + """Log samples to wandb as a table""" + + # Create column names for all samples + columns = ["step"] + sum( + [[f"input_{i + 1}", f"output_{i + 1}", f"score_{i + 1}"] for i in range(len(samples))], [] + ) + + if not hasattr(self, "validation_table"): + # Initialize the table on first call + self.validation_table = wandb.Table(columns=columns) + + # Create a new table with same columns and existing data + # Workaround for https://github.com/wandb/wandb/issues/2981#issuecomment-1997445737 + new_table = wandb.Table(columns=columns, data=self.validation_table.data) + + # Add new row with all data + row_data = [] + row_data.append(step) + for sample in samples: + row_data.extend(sample) + + new_table.add_data(*row_data) + + # Update reference and log + if wandb.run is not None: + wandb.log({"val/generations": new_table}, step=step) + self.validation_table = new_table + + def log_generations_to_swanlab(self, samples, step): + """Log samples to swanlab as text""" + import swanlab + + swanlab_table = swanlab.echarts.Table() + + # Create column names + headers = ["step", "input", "output", "score"] + + swanlab_row_list = [[step, *sample] for sample in samples] + swanlab_table.add(headers=headers, rows=swanlab_row_list) + + # Log to swanlab + swanlab.log({"val/generations": swanlab_table}, step=step) + + def log_generations_to_mlflow(self, samples, step): + """Log validation generation to mlflow as artifacts""" + # https://mlflow.org/docs/latest/api_reference/python_api/mlflow.html?highlight=log_artifact#mlflow.log_artifact + + import tempfile + + import mlflow + + try: + with tempfile.TemporaryDirectory() as tmp_dir: + validation_gen_step_file = Path(tmp_dir, f"val_step{step}.json") + row_data = [] + for sample in samples: + data = {"input": sample[0], "output": sample[1], "score": sample[2]} + row_data.append(data) + with open(validation_gen_step_file, "w") as file: + json.dump(row_data, file) + mlflow.log_artifact(validation_gen_step_file) + except Exception as e: + print(f"WARNING: save validation generation file to mlflow failed with error {e}") + + def log_generations_to_clearml(self, samples, step): + """Log validation generation to clearml as table""" + + import clearml + import pandas as pd + + task: clearml.Task | None = clearml.Task.current_task() + if task is None: + return + + table = [ + { + "step": step, + "input": sample[0], + "output": sample[1], + "score": sample[2], + } + for sample in samples + ] + + logger = task.get_logger() + logger.report_table( + series="Validation generations", + title="Validation", + table_plot=pd.DataFrame.from_records(table), + iteration=step, + ) + + def log_generations_to_tensorboard(self, samples, step): + """Log samples to tensorboard as text""" + # Initialize tensorboard writer if not exists + if not hasattr(self, "writer"): + from torch.utils.tensorboard import SummaryWriter + + # Use the same directory structure as _TensorboardAdapter + if self.project_name and self.experiment_name: + default_dir = os.path.join("tensorboard_log", self.project_name, self.experiment_name) + else: + default_dir = "tensorboard_log" + + tensorboard_dir = os.environ.get("TENSORBOARD_DIR", default_dir) + os.makedirs(tensorboard_dir, exist_ok=True) + self.writer = SummaryWriter(log_dir=tensorboard_dir) + + # Format the samples data into readable text + text_content = f"**Generation Results - Step {step}**\n\n" + + for i, sample in enumerate(samples): + text_content += f"### Sample {i + 1}\n" + + # Assuming sample contains [input, output, score] + if len(sample) >= 3: + input_text, output_text, score = sample[0], sample[1], sample[2] + + text_content += f"**Input:** {input_text}\n\n" + text_content += f"**Output:** {output_text}\n\n" + text_content += f"**Score:** {score}\n\n" + else: + # Handle cases where sample format might be different + text_content += f"**Data:** {sample}\n\n" + + text_content += "---\n\n" + + # Log to tensorboard as text + self.writer.add_text("val/generations", text_content, step) + # Flush to ensure data is written + self.writer.flush() diff --git a/verl/verl/utils/transferqueue_utils.py b/verl/verl/utils/transferqueue_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..6642fd2696699c77eb61a5ba9aff3807b8c1680c --- /dev/null +++ b/verl/verl/utils/transferqueue_utils.py @@ -0,0 +1,430 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +import copy +import functools +import inspect +import logging +import os +import threading +import time +from functools import wraps +from typing import TYPE_CHECKING, Any, Callable + +import torch +from tensordict.tensorclass import NonTensorData, NonTensorStack + +if TYPE_CHECKING: + from verl.single_controller.base.decorator import Dispatch + +from tensordict import TensorDict + +try: + import transfer_queue as tq + from transfer_queue import ( + BatchMeta, + KVBatchMeta, + ) + +except ImportError: + + class BatchMeta: + pass + + class KVBatchMeta: + pass + + # Mock transfer_queue module when not installed + class _MockTQ: + """Mock transfer_queue module that raises RuntimeError on any access.""" + + def __getattr__(self, name: str) -> Any: + def _raise(*args, **kwargs): + raise RuntimeError( + f"transfer_queue is not installed. Cannot use tq.{name}(). " + "Please install it by calling `pip install TransferQueue==0.1.6`" + ) + + return _raise + + tq = _MockTQ() + + +from verl.utils import tensordict_utils as tu + +logger = logging.getLogger(__name__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + +TQ_INITIALIZED = False + + +# TODO (TQ): verl will make all actor async, so this can be cleanup later. +def _run_async_in_temp_loop(async_func: Callable[..., Any], *args, **kwargs) -> Any: + # Use a temporary event loop in a new thread because event + # loop may already exist in server mode + tmp_event_loop = asyncio.new_event_loop() + thread = threading.Thread( + target=tmp_event_loop.run_forever, + name="batchmeta tensordict converter", + daemon=True, + ) + + def run_coroutine(coroutine): + if not thread.is_alive(): + thread.start() + future = asyncio.run_coroutine_threadsafe(coroutine, tmp_event_loop) + return future.result() + + async def stop_loop(): + tmp_event_loop.stop() + + try: + return run_coroutine(async_func(*args, **kwargs)) + finally: + if thread.is_alive(): + asyncio.run_coroutine_threadsafe(stop_loop(), tmp_event_loop) + thread.join() + + +def _find_meta(*args, **kwargs): + for arg in args: + if isinstance(arg, BatchMeta | KVBatchMeta): + return arg + for v in kwargs.values(): + if isinstance(v, BatchMeta | KVBatchMeta): + return v + return None + + +async def _async_meta_to_realdata(meta: BatchMeta | KVBatchMeta) -> TensorDict: + if isinstance(meta, KVBatchMeta): + meta = await async_kv_batch_meta2batch_meta(meta) + meta_info = copy.deepcopy(meta.extra_info) + if meta.size == 0: + empty_td = TensorDict({}, batch_size=(0,)) + tu.assign_non_tensor(empty_td, **meta_info) + return empty_td + + tq_client = tq.get_client() + tensordict = await tq_client.async_get_data(meta) + + for key, val in meta_info.items(): + if isinstance(val, (NonTensorData | NonTensorStack)): + tensordict[key] = val + else: + tu.assign_non_tensor_data(tensor_dict=tensordict, key=key, val=val) + return tensordict + + +def _meta_to_realdata(meta: BatchMeta) -> TensorDict: + return _run_async_in_temp_loop(_async_meta_to_realdata, meta) + + +async def _async_update_meta_with_output(output: TensorDict, meta: BatchMeta, func_name=None) -> BatchMeta: + fields, meta_data = [], {} + for k, v in output.items(): + if isinstance(v, torch.Tensor | NonTensorStack): + fields.append(k) + elif isinstance(v, NonTensorData): + meta_data[k] = v.data + else: + raise ValueError(f"Unsupported type {type(v)} for key {k} in output TensorDict.") + + if fields: + t1 = time.time() + tq_client = tq.get_client() + meta = await tq_client.async_put(data=output.select(*fields), metadata=meta) + t2 = time.time() + + logger.info(f"Task {func_name} (pid={os.getpid()}) is writing to TransferQueue, cost time: {t2 - t1}s)") + meta.extra_info = meta_data + return meta + + +def _update_meta_with_output(output: TensorDict, meta: BatchMeta, func_name=None) -> BatchMeta: + updated_meta = _run_async_in_temp_loop(_async_update_meta_with_output, output, meta, func_name) + return updated_meta + + +def _compute_need_collect(dispatch_mode: "dict | Dispatch", args: list) -> bool: + """Compute whether data collection is needed for the current worker. + + This function determines whether the current worker should collect data based on + the dispatch mode configuration and worker parameters. It's used to optimize + distributed data collection by ensuring only the appropriate rank collects data. + + Args: + dispatch_mode: Controls data collection logic for the current worker. Can be None, + a Dispatch instance, or a dict with 'collect_fn' key. If None or Dispatch, + always returns True (current worker should collect). If dict, checks + collect_fn for lazy compute optimization. + args: List of arguments passed to the function. Should contain a Worker instance + as the first argument when using lazy compute mode. + + Returns: + bool: True if data collection is needed, False otherwise. + + Note: + Only checks worker attributes when dispatch_mode is a dict with 'collect_fn', + the collect_fn is 'collect_lazy_compute_data_proto', and args[0] is a Worker. + Otherwise, returns True. For the lazy compute case, checks the worker's + data parallel rank for the mesh specified in collect_fn.args[0] to determine + if this worker should collect data. + """ + from verl.single_controller.base.decorator import Dispatch + from verl.single_controller.base.worker import Worker + + if dispatch_mode is None or isinstance(dispatch_mode, Dispatch): + return True + + assert "collect_fn" in dispatch_mode.keys(), "collect_fn should be in dispatch_mode." + + collect_fn = dispatch_mode["collect_fn"] + + # Check if collect_fn is a functools.partial and handle gracefully + if isinstance(collect_fn, functools.partial): + collect_fn_name = collect_fn.func.__name__ + if collect_fn_name != "collect_lazy_compute_data_proto" or len(args) < 1 or not isinstance(args[0], Worker): + return True + + collect_mesh_name = collect_fn.args[0] if collect_fn.args else None + if collect_mesh_name is None: + return True + + return args[0].query_collect_info(collect_mesh_name) + else: + # If collect_fn is not a partial, we can't extract mesh_name information + # Fall back to default behavior (collect data) + return True + + +def _postprocess_common(output, put_data, need_collect): + """Common post-processing logic for function outputs in TransferQueue bridge. + + This function handles the final return value based on whether data should be + put into storage (put_data) and whether collection is needed (need_collect). + It ensures proper return types based on the execution context. + + Args: + output: The original output from the decorated function. Can be any type. + put_data: bool, indicating whether the output should be put into TransferQueue. + If True, output will be put to TQ and return the corresponding BatchMeta; + if False, output will not be put into TQ. + need_collect: bool, indicating whether this process needs to collect data. + If False, the output will be replaced by an empty BatchMeta or DataProto + to avoid redundant communication. + + Returns: + - BatchMeta.empty(): When put_data=True but need_collect=False, indicating + no data should be stored but BatchMeta structure is expected. + - DataProto(): When put_data=False, need_collect=False, and output is DataProto, + returning an empty DataProto. + - output: In all other cases, returns the original output unchanged. + + Note: + This function is used in the tqbridge decorator to normalize return values + across different execution paths and avoid redundant data operations in + distributed scenarios. + """ + from verl.protocol import DataProto + + if put_data and not need_collect: + return BatchMeta() + elif not put_data and not need_collect and isinstance(output, DataProto): + return DataProto() + elif not put_data and not need_collect and isinstance(output, TensorDict): + return TensorDict({}, batch_size=(0,)) + else: + return output + + +async def async_kv_batch_meta2batch_meta(meta: KVBatchMeta) -> BatchMeta: + global TQ_INITIALIZED + if not TQ_INITIALIZED: + tq.init() + TQ_INITIALIZED = True + tq_client = tq.get_client() + batch_meta = await tq_client.async_kv_retrieve_meta(keys=meta.keys, partition_id=meta.partition_id, create=False) + fields = meta.fields + if fields is not None: + if isinstance(fields, str): + fields = [fields] + batch_meta = batch_meta.select_fields(fields) + + batch_meta.extra_info = meta.extra_info + return batch_meta + + +def kv_batch_meta2batch_meta(meta: KVBatchMeta): + return _run_async_in_temp_loop(async_kv_batch_meta2batch_meta, meta) + + +async def async_batch_meta2kv_batch_meta(meta: BatchMeta) -> KVBatchMeta: + global TQ_INITIALIZED + if not TQ_INITIALIZED: + tq.init() + TQ_INITIALIZED = True + tq_client = tq.get_client() + partition_id = meta.partition_ids[0] + assert all([partition_id == pid for pid in meta.partition_ids]) + keys = await tq_client.async_kv_retrieve_keys(global_indexes=meta.global_indexes, partition_id=partition_id) + + kv_batch_meta = KVBatchMeta( + keys=keys, + tags=[{}] * meta.size, + partition_id=partition_id, + fields=meta.field_names, + extra_info=meta.extra_info, + ) + return kv_batch_meta + + +def batch_meta2kv_batch_meta(meta: BatchMeta): + return _run_async_in_temp_loop(async_batch_meta2kv_batch_meta, meta) + + +def tqbridge(dispatch_mode: "dict | Dispatch" = None): + """Creates a decorator for bridging KVBatchMeta and TensorDict. + + This decorator automatically handles conversions between `KVBatchMeta` + and `TensorDict` in function parameters, and decides whether to sync function + output back to `KVBatchMeta` based on configuration(`put_data`). It supports + both synchronous and asynchronous functions (async def). When TQ is not enabled, it + simply calls the original function as-is. + + Args: + dispatch_mode: Controls data collection behavior for the current worker. Passed to + _compute_need_collect to determine if current worker should collect data. + If None, _compute_need_collect will return True to fallback default logics. + + + Returns: + A decorator function used to decorate target functions (synchronous or asynchronous). + """ + # TODO: move to the top + from verl.single_controller.base.decorator import _check_dispatch_mode + + if dispatch_mode is not None: + _check_dispatch_mode(dispatch_mode) + + def decorator(func): + pid = os.getpid() + + @wraps(func) + def inner(*args, **kwargs): + batch_meta = _find_meta(*args, **kwargs) + if batch_meta is None: + return func(*args, **kwargs) + else: + global TQ_INITIALIZED + if not TQ_INITIALIZED: + tq.init() + TQ_INITIALIZED = True + + is_kv_batch_meta = isinstance(batch_meta, KVBatchMeta) + if is_kv_batch_meta: + tags = batch_meta.tags + batch_meta = kv_batch_meta2batch_meta(batch_meta) + t1 = time.time() + args = [_meta_to_realdata(arg) if isinstance(arg, BatchMeta | KVBatchMeta) else arg for arg in args] + kwargs = { + k: _meta_to_realdata(v) if isinstance(v, BatchMeta | KVBatchMeta) else v for k, v in kwargs.items() + } + t2 = time.time() + logger.info( + f"Task {func.__name__} (pid={pid}) is getting len_samples={batch_meta.size}, cost time: {t2 - t1}" + ) + + output = func(*args, **kwargs) + + put_data = False + if isinstance(output, TensorDict): + if output.batch_size: + assert output.batch_size[0] == batch_meta.size, ( + f"output batch size {output.batch_size} != meta size {batch_meta.size}" + ) + put_data = True + + if dispatch_mode is not None: + need_collect = _compute_need_collect(dispatch_mode, args) + else: + need_collect = True + if put_data and need_collect: + updated_meta = _update_meta_with_output(output, batch_meta, func.__name__) + if is_kv_batch_meta: + updated_meta = batch_meta2kv_batch_meta(updated_meta) + updated_meta.tags = tags + return updated_meta + return _postprocess_common(output, put_data, need_collect) + + @wraps(func) + async def async_inner(*args, **kwargs): + batch_meta = _find_meta(*args, **kwargs) + if batch_meta is None: + return await func(*args, **kwargs) + else: + global TQ_INITIALIZED + if not TQ_INITIALIZED: + tq.init() + TQ_INITIALIZED = True + + is_kv_batch_meta = isinstance(batch_meta, KVBatchMeta) + if is_kv_batch_meta: + tags = batch_meta.tags + batch_meta = await async_kv_batch_meta2batch_meta(batch_meta) + + t1 = time.time() + args = [ + await _async_meta_to_realdata(arg) if isinstance(arg, BatchMeta | KVBatchMeta) else arg + for arg in args + ] + kwargs = { + k: await _async_meta_to_realdata(v) if isinstance(v, BatchMeta | KVBatchMeta) else v + for k, v in kwargs.items() + } + t2 = time.time() + logger.info( + f"Task {func.__name__} (pid={pid}) is getting len_samples={batch_meta.size}, cost time: {t2 - t1}" + ) + + output = await func(*args, **kwargs) + + put_data = False + if isinstance(output, TensorDict): + if output.batch_size: + assert output.batch_size[0] == batch_meta.size, ( + f"output batch size {output.batch_size} != meta size {batch_meta.size}" + ) + put_data = True + + if dispatch_mode is not None: + need_collect = _compute_need_collect(dispatch_mode, args) + else: + need_collect = True + if put_data and need_collect: + updated_meta = await _async_update_meta_with_output(output, batch_meta, func.__name__) + if is_kv_batch_meta: + updated_meta = await async_batch_meta2kv_batch_meta(updated_meta) + updated_meta.tags = tags + return updated_meta + return _postprocess_common(output, put_data, need_collect) + + wrapper_inner = inner + wrapper_async_inner = async_inner + + wrapper = wrapper_async_inner if inspect.iscoroutinefunction(func) else wrapper_inner + return wrapper + + return decorator diff --git a/verl/verl/utils/transformers_compat.py b/verl/verl/utils/transformers_compat.py new file mode 100644 index 0000000000000000000000000000000000000000..e7fb659fddbbdd0239a5006da6ba9b7d291b61bf --- /dev/null +++ b/verl/verl/utils/transformers_compat.py @@ -0,0 +1,121 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Compatibility utilities for different versions of transformers library. +""" + +import importlib.metadata +from functools import lru_cache +from typing import Optional + +from packaging import version + +# Handle version compatibility for flash_attn_supports_top_left_mask +# This function was added in newer versions of transformers +try: + from transformers.modeling_flash_attention_utils import flash_attn_supports_top_left_mask +except ImportError: + # For older versions of transformers that don't have this function + # Default to False as a safe fallback for older versions + def flash_attn_supports_top_left_mask(): + """Fallback implementation for older transformers versions. + Returns False to disable features that require this function. + """ + return False + + +@lru_cache +def is_transformers_version_in_range(min_version: Optional[str] = None, max_version: Optional[str] = None) -> bool: + try: + # Get the installed version of the transformers library + transformers_version_str = importlib.metadata.version("transformers") + except importlib.metadata.PackageNotFoundError as e: + raise ModuleNotFoundError("The `transformers` package is not installed.") from e + + transformers_version = version.parse(transformers_version_str) + + lower_bound_check = True + if min_version is not None: + lower_bound_check = version.parse(min_version) <= transformers_version + + upper_bound_check = True + if max_version is not None: + upper_bound_check = transformers_version <= version.parse(max_version) + + return lower_bound_check and upper_bound_check + + +@lru_cache +def get_auto_model_for_vision2seq(): + """Return the available VL auto model class across transformers versions.""" + + try: + # Prefer the newer class when available. In transformers 4.x this class has + # a broader mapping than AutoModelForVision2Seq, and AutoModelForVision2Seq + # is deprecated for removal in v5. + from transformers import AutoModelForImageTextToText + except ImportError: + from transformers import AutoModelForVision2Seq + + return AutoModelForVision2Seq + + return AutoModelForImageTextToText + + +def unpack_visual_output(visual_output): + """Unpack the output from the visual encoder, handling both tuple and object return types. + + Newer versions of transformers return an object with `pooler_output` and `deepstack_features` + attributes instead of a plain tuple. + """ + if hasattr(visual_output, "pooler_output"): + # For newer versions(>=5.0.0) of transformers, return the pooler_output and deepstack_features + if hasattr(visual_output, "deepstack_features"): + return visual_output.pooler_output, visual_output.deepstack_features + else: + return visual_output.pooler_output, None + if isinstance(visual_output, tuple): + return visual_output + else: + return visual_output, None + + +def drop_tied_target_keys(state_dict, model, model_config) -> None: + """Drop tied alias keys (e.g. ``lm_head.weight``) from ``state_dict``. + + FSDP gather and the model merger materialise one independent CPU tensor + per state_dict key, which defeats ``save_pretrained``'s storage-pointer- + based dedup. When both keys end up in safetensors, ``transformers>=5`` + silently refuses to re-tie on reload. Detects aliases by ``Parameter`` + identity after ``tie_weights()``. + + Only drops the alias when the canonical name has been confirmed to exist + in ``state_dict``; otherwise the alias is promoted to canonical so we + never accidentally remove every entry for a tied parameter. + """ + if not getattr(model_config, "tie_word_embeddings", False): + return + model.tie_weights() + canonical_by_id: dict[int, str] = {} + for name, param in model.named_parameters(remove_duplicate=False): + pid = id(param) + if pid not in canonical_by_id: + canonical_by_id[pid] = name + continue + if canonical_by_id[pid] in state_dict: + state_dict.pop(name, None) + else: + # Canonical name absent: promote the alias instead of erasing it. + canonical_by_id[pid] = name diff --git a/verl/verl/utils/trtllm/__init__.py b/verl/verl/utils/trtllm/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d5aee6bdcdd86a19e72d4a8e96dd3f6b9abc337 --- /dev/null +++ b/verl/verl/utils/trtllm/__init__.py @@ -0,0 +1,14 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/verl/verl/utils/trtllm/trtllm_fp8_utils.py b/verl/verl/utils/trtllm/trtllm_fp8_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..43c15351c06057efbdc82b2ab1b3940a3601b106 --- /dev/null +++ b/verl/verl/utils/trtllm/trtllm_fp8_utils.py @@ -0,0 +1,21 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from verl.utils.fp8_utils import FP8QuantizerHelper + + +class TRTLLMFP8QuantizerHelper(FP8QuantizerHelper): + def __init__(self, quant_config): + super().__init__(quant_config) diff --git a/verl/verl/utils/ulysses.py b/verl/verl/utils/ulysses.py new file mode 100644 index 0000000000000000000000000000000000000000..7ee07dd9090f9c6d1e64a9ccefe326b7e907713c --- /dev/null +++ b/verl/verl/utils/ulysses.py @@ -0,0 +1,404 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Utilities for DeepSpeed Ulysses Sequence Parallelism. +DeepSpeed Ulysses Paper: https://arxiv.org/abs/2309.14509 +Inspired from: https://github.com/deepspeedai/DeepSpeed/blob/master/deepspeed/sequence/layer.py +""" + +from typing import TYPE_CHECKING, Any, Optional + +import torch +import torch.distributed as dist +from torch import Tensor +from torch.distributed import ProcessGroup +from torch.distributed.device_mesh import DeviceMesh + +if TYPE_CHECKING: + from verl import DataProto + +_ULYSSES_SEQUENCE_PARALLEL_GROUP = None + + +def set_ulysses_sequence_parallel_group(group: dist.ProcessGroup): + """ + Set ulysses sequence parallel process group. + """ + global _ULYSSES_SEQUENCE_PARALLEL_GROUP + _ULYSSES_SEQUENCE_PARALLEL_GROUP = group + + +def get_ulysses_sequence_parallel_group() -> Optional[dist.ProcessGroup]: + """ + Get ulysses sequence parallel process group. + """ + global _ULYSSES_SEQUENCE_PARALLEL_GROUP + return _ULYSSES_SEQUENCE_PARALLEL_GROUP + + +def get_ulysses_sequence_parallel_world_size(group: ProcessGroup = None) -> int: + """ + Get ulysses sequence parallel world size. + """ + group = get_ulysses_sequence_parallel_group() if group is None else group + return dist.get_world_size(group) if group else 1 + + +def get_ulysses_sequence_parallel_rank(group: ProcessGroup = None) -> int: + """ + Get ulysses sequence parallel rank. + """ + group = get_ulysses_sequence_parallel_group() if group is None else group + return dist.get_rank(group) if group else 0 + + +def gather_seq_scatter_heads( + x: Tensor, + seq_dim: int, + head_dim: int, + unpadded_dim_size: int = 0, + group: ProcessGroup = None, +) -> Tensor: + """ + A func to sync embedding input with alltoall in sequence parallel + gather sequence dimension and scatter head dim: + e.g. seq_dim: 1, head_dim: 2 + [bsz, seq/n, h, ...] -> [bsz, seq, h/n, ...] + """ + group = get_ulysses_sequence_parallel_group() if group is None else group + if not group: + return x + sp_world = get_ulysses_sequence_parallel_world_size(group) + x = SeqAllToAll.apply(group, x, head_dim, seq_dim) + if unpadded_dim_size and unpadded_dim_size % sp_world != 0: + padding_size = x.size(seq_dim) - unpadded_dim_size + x = _unpad_tensor(x, seq_dim, padding_size) + return x + + +def gather_heads_scatter_seq(x: Tensor, head_dim: int, seq_dim: int, group: ProcessGroup = None) -> Tensor: + """ + A func to sync attention result with alltoall in sequence parallel + gather head dimension and scatter seq dim: + e.g. seq_dim: 1, head_dim: 2 + [bsz, seq, h/n, ...] -> [bsz, seq/n, h, ...] + """ + group = get_ulysses_sequence_parallel_group() if group is None else group + if not group: + return x + dim_size = x.size(seq_dim) + sp_world = get_ulysses_sequence_parallel_world_size(group) + if dim_size % sp_world != 0: + padding_size = sp_world - (dim_size % sp_world) + x = _pad_tensor(x, seq_dim, padding_size) + return SeqAllToAll.apply(group, x, seq_dim, head_dim, False) + + +def _pad_tensor(x: Tensor, dim: int, padding_size: int) -> Tensor: + shape = list(x.shape) + shape[dim] = padding_size + pad = torch.zeros(shape, dtype=x.dtype, device=x.device) + return torch.cat([x, pad], dim=dim) + + +def _unpad_tensor(x: Tensor, dim: int, padding_size: int) -> Tensor: + slc = [slice(None)] * len(x.shape) + slc[dim] = slice(0, -padding_size) + return x[tuple(slc)] + + +def slice_input_tensor(x: Tensor, dim: int, padding: bool = True, group: ProcessGroup = None) -> Tensor: + group = get_ulysses_sequence_parallel_group() if group is None else group + sp_world_size = dist.get_world_size(group) + sp_rank = get_ulysses_sequence_parallel_rank() + dim_size = x.size(dim) + # pad before slice + if padding and dim_size % sp_world_size: + padding_size = sp_world_size - (dim_size % sp_world_size) + x = _pad_tensor(x, dim, padding_size) + # slice the input tensor + parts = x.size(dim) // sp_world_size + slc = [slice(None)] * len(x.shape) + slc[dim] = slice(sp_rank * parts, (sp_rank + 1) * parts) + return x[tuple(slc)].contiguous() + + +def all_to_all_tensor( + local_input: Tensor, + scatter_dim: int, + gather_dim: int, + group: Optional[dist.ProcessGroup] = None, + async_op: bool = False, +): + group = get_ulysses_sequence_parallel_group() if group is None else group + seq_world_size = dist.get_world_size(group) + input_list = [t.contiguous() for t in torch.tensor_split(local_input, seq_world_size, scatter_dim)] + output_list = [torch.empty_like(input_list[0]) for _ in range(seq_world_size)] + comm = dist.all_to_all(output_list, input_list, group=group, async_op=async_op) + if async_op: + + def wait(): + comm.wait() + return torch.cat(output_list, dim=gather_dim).contiguous() + + return wait + return torch.cat(output_list, dim=gather_dim).contiguous() + + +def all_gather_tensor(local_tensor: Tensor, group: Optional[dist.ProcessGroup] = None, async_op: bool = False): + group = get_ulysses_sequence_parallel_group() if group is None else group + sp_world_size = dist.get_world_size(group=group) + output_shape = list(local_tensor.shape) + output_shape[0] = output_shape[0] * sp_world_size + output = torch.empty(output_shape, dtype=local_tensor.dtype, device=local_tensor.device) + dist.all_gather_into_tensor(output, local_tensor, group=group, async_op=async_op) + return output + + +class SeqAllToAll(torch.autograd.Function): + @staticmethod + def forward( + ctx: Any, + group: dist.ProcessGroup, + local_input: Tensor, + scatter_dim: int, + gather_dim: int, + async_op: bool = False, + ) -> Tensor: + ctx.group = group + ctx.scatter_dim = scatter_dim + ctx.gather_dim = gather_dim + ctx.async_op = async_op + return all_to_all_tensor(local_input, scatter_dim, gather_dim, group, async_op) + + @staticmethod + def backward(ctx: Any, *grad_output: Tensor) -> tuple[None, Tensor, None, None]: + input_t = torch.cat(grad_output[1:], dim=ctx.gather_dim).contiguous() if ctx.async_op else grad_output[0] + return ( + None, + all_to_all_tensor(input_t, ctx.gather_dim, ctx.scatter_dim, ctx.group, False), + None, + None, + None, + None, + ) + + +class Gather(torch.autograd.Function): + @staticmethod + def forward( + ctx: Any, + group: dist.ProcessGroup, + local_tensor: Tensor, + gather_dim: int, + grad_scaler: bool = True, + async_op=False, + ) -> Tensor: + ctx.group = group + ctx.gather_dim = gather_dim + ctx.grad_scaler = grad_scaler + ctx.async_op = async_op + + sp_world_size = dist.get_world_size(group=group) + ctx.sp_world_size = sp_world_size + + sp_rank = dist.get_rank(group=group) + ctx.sp_rank = sp_rank + + local_shape = list(local_tensor.size()) + split_size = local_shape[0] + part_size = local_shape[gather_dim] # store original size + ctx.part_size = part_size + + output = all_gather_tensor(local_tensor, group, async_op) + return torch.cat(output.split(split_size, dim=0), dim=gather_dim) + + @staticmethod + def backward(ctx: Any, grad_output: Tensor) -> Any: + if ctx.grad_scaler: + grad_output = grad_output * ctx.sp_world_size + return ( + None, + grad_output.split(ctx.part_size, dim=ctx.gather_dim)[ctx.sp_rank].contiguous(), + None, + None, + None, + None, + ) + + +def gather_outpus_and_unpad(*args, **kwargs): + raise RuntimeError( + "please use verl.utils.ulysses.gather_outputs_and_unpad instead of verl.utils.ulysses.gather_outpus_and_unpad" + ) + + +def gather_outputs_and_unpad( + x: Tensor, + gather_dim: int, + unpad_dim: int = None, + padding_size: int = 0, + grad_scaler: bool = True, + group: Optional[dist.ProcessGroup] = None, +): + """ + Gather a tensor across a process group and optionally unpad its padded elements. + + Args: + x (Tensor): Input tensor to gather. + gather_dim (int): Dimension along which to gather across ranks. + unpad_dim (int, optional): Dimension from which to remove padding. If None, no unpadding. + padding_size (int): Number of padding elements to remove on `unpad_dim`. Defaults to 0. + grad_scaler (bool): Whether to apply gradient scaling during gather. Defaults to True. + group (ProcessGroup, optional): Process group for gathering. If None, uses + `get_ulysses_sequence_parallel_group()`. If still None, returns `x` unchanged. + + Returns: + Tensor: The gathered tensor, with padding removed if requested. + """ + group = get_ulysses_sequence_parallel_group() if group is None else group + if group is None: + return x + x = Gather.apply(group, x, gather_dim, grad_scaler) + if unpad_dim is not None: + assert isinstance(padding_size, int), "padding size is not given or is not an integer" + if padding_size == 0: + return x + x = _unpad_tensor(x, unpad_dim, padding_size) + return x + + +def ulysses_pad( + input_ids_rmpad: torch.Tensor, position_ids_rmpad: Optional[torch.Tensor] = None, sp_size: int = 1, pad_value=0 +): + if position_ids_rmpad is not None: + assert position_ids_rmpad.size(-2) == 1 + assert input_ids_rmpad.size(-1) == position_ids_rmpad.size(-1) + if sp_size <= 1: + return input_ids_rmpad, position_ids_rmpad, 0 + _, total_seq_len = input_ids_rmpad.shape + pad_size = (sp_size - total_seq_len % sp_size) % sp_size + if pad_size > 0: + input_ids_rmpad = torch.nn.functional.pad(input_ids_rmpad, (0, pad_size), value=pad_value) + if position_ids_rmpad is not None: + pad_pos_ids = torch.arange(pad_size, device=position_ids_rmpad.device).unsqueeze(0) + if position_ids_rmpad.dim() == 3: + pad_pos_ids = pad_pos_ids.unsqueeze(0).repeat(position_ids_rmpad.size(0), 1, 1) + position_ids_rmpad = torch.cat((position_ids_rmpad, pad_pos_ids), dim=-1) + return input_ids_rmpad, position_ids_rmpad, pad_size + + +def ulysses_pad_and_slice_inputs( + input_ids_rmpad: torch.Tensor, + position_ids_rmpad: Optional[torch.Tensor] = None, + sp_size: int = 1, + skip_position_ids_rmpad: bool = False, + pad_value=0, +): + """ + Pad and slice input_ids to be divisible by sp_size + Pad position_ids to be divisible by sp_size. + + Note both input_ids_rmpad and position_ids_rmpad will be padded and sliced. + + The is the utility of pre-forward for ulysses sequence parallelism + + Args: + input_ids_rmpad: shape of [bsz, seqlen] + position_ids_rmpad: shape of [bsz, seqlen], where bsz must be 1 + sp_size (int): ulysses sequence parallelism size + skip_position_ids_rmpad: whether to skip position_ids_rmpad for VeOmniEngine + + Returns: + torch.Tensor: padded and sliced input_ids + torch.Tensor: padded and sliced position_ids + int: pad size + """ + input_ids_rmpad, position_ids_rmpad, pad_size = ulysses_pad( + input_ids_rmpad, position_ids_rmpad, sp_size, pad_value=pad_value + ) + input_ids_rmpad = slice_input_tensor(input_ids_rmpad, dim=1, padding=False) + if position_ids_rmpad is not None and not skip_position_ids_rmpad: + position_ids_rmpad = slice_input_tensor(position_ids_rmpad, dim=1, padding=False) + return input_ids_rmpad, position_ids_rmpad, pad_size + + +def validate_ulysses_config(num_heads, ulysses_sequence_size): + if ulysses_sequence_size > 1: + assert num_heads % ulysses_sequence_size == 0, ( + f"num_heads ({num_heads}) must be divisible by ulysses sequence size({ulysses_sequence_size})" + ) + + +class BaseShardingManager: + """Base sharding manager used for resharding weights/data across parallel groups.""" + + def __init__(self): + self.timing = {} + + def __enter__(self): + pass + + def __exit__(self, exc_type, exc_value, traceback): + pass + + def preprocess_data(self, data: "DataProto") -> "DataProto": + return data + + def postprocess_data(self, data: "DataProto") -> "DataProto": + return data + + +class FSDPUlyssesShardingManager(BaseShardingManager): + """ + Sharding manager to support data resharding when using FSDP + Ulysses sequence parallelism. + """ + + def __init__(self, device_mesh: DeviceMesh): + super().__init__() + self.device_mesh = device_mesh + self.seed_offset = 12345 + + def __enter__(self): + if self.device_mesh is not None: + self.prev_sp_group = get_ulysses_sequence_parallel_group() + set_ulysses_sequence_parallel_group(self.device_mesh["sp"].get_group()) + + def __exit__(self, exc_type, exc_value, traceback): + if self.device_mesh is not None: + set_ulysses_sequence_parallel_group(self.prev_sp_group) + + def preprocess_data(self, data: "DataProto") -> "DataProto": + """ + AllGather data from sp region. + + This is because the data is first sharded along the FSDP dimension as we utilize the DP_COMPUTE. + In Ulysses, we need to make sure the same data is used across a SP group. + """ + if self.device_mesh is not None: + from verl.protocol import all_gather_data_proto + + group = self.device_mesh["sp"].get_group() + all_gather_data_proto(data=data, process_group=group) + return data + + def postprocess_data(self, data: "DataProto") -> "DataProto": + """ + Split the data to follow FSDP partition. + """ + if self.device_mesh is not None: + sp_size = self.device_mesh["sp"].size() + sp_rank = self.device_mesh["sp"].get_local_rank() + data = data.chunk(chunks=sp_size)[sp_rank] + return data diff --git a/verl/verl/utils/vllm/__init__.py b/verl/verl/utils/vllm/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2221fa8e2de228d889159c89b83de5738fe206cb --- /dev/null +++ b/verl/verl/utils/vllm/__init__.py @@ -0,0 +1,28 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .npu_vllm_patch import check_vllm_ascend_before_server_launch +from .utils import TensorLoRARequest, VLLMHijack, is_version_ge + +# The contents of vllm/patch.py should not be imported here, because the contents of +# patch.py should be imported after the vllm LLM instance is created. Therefore, +# wait until you actually start using it before importing the contents of +# patch.py separately. + +__all__ = [ + "TensorLoRARequest", + "VLLMHijack", + "is_version_ge", + "check_vllm_ascend_before_server_launch", +] diff --git a/verl/verl/utils/vllm/npu_vllm_patch.py b/verl/verl/utils/vllm/npu_vllm_patch.py new file mode 100644 index 0000000000000000000000000000000000000000..c22647ceac607c24e4bebf1fcab92a800357fb52 --- /dev/null +++ b/verl/verl/utils/vllm/npu_vllm_patch.py @@ -0,0 +1,227 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Copyright 2025 The Qwen Team and The HuggingFace Inc. team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import os +from functools import wraps + +from verl.utils.device import is_torch_npu_available + + +def vllm_ascend_v011_select_moe_comm_method_wrapper(fn): + @wraps(fn) + def wrapper(self, num_tokens, with_prefill): + moe_comm_method = fn(self, num_tokens, with_prefill) + from vllm_ascend.ascend_forward_context import MoECommType + from vllm_ascend.utils import AscendSocVersion, enable_sp, get_ascend_soc_version + + soc_version = get_ascend_soc_version() + + # AscendSocVersion.A2 is not support MC2 in Single-card multi-process scenario now. + if soc_version in {AscendSocVersion.A2} and moe_comm_method == MoECommType.MC2: + quant_type = getattr(self.vllm_config.model_config.hf_config, "moe_quantize", None) + # Currently, w4a8_dynamic does not support allgatherep + if quant_type == "w4a8_dynamic": + moe_comm_method = MoECommType.ALLTOALL + else: + moe_comm_method = MoECommType.ALLGATHER + + if with_prefill: + if enable_sp(): + moe_comm_method = MoECommType.ALLGATHER + else: + moe_comm_method = MoECommType.NAIVE_MULTICAST + + return moe_comm_method + + return wrapper + + +def vllm_ascend_v011_matmul_and_reduce_wrapper(fn): + @wraps(fn) + def wrapper(self, *args, **kwargs): + from vllm_ascend.utils import AscendSocVersion, get_ascend_soc_version + + soc_version = get_ascend_soc_version() + # AscendSocVersion.A2 is not support MC2 in Single-card multi-process scenario now. + if soc_version in {AscendSocVersion.A2}: + from vllm.forward_context import get_forward_context + + try: + forward_context = get_forward_context() + forward_context.mmrs_fusion = False + except AssertionError: + # forward_context.mmrs_fusion will be false in matmul_and_reduce func. + pass + return fn(self, *args, **kwargs) + + return wrapper + + +def check_vllm_ascend_before_server_launch(): + import torch_npu + import vllm + + def _is_ascend_soc_version_A2_v011_local(): + from vllm_ascend.utils import AscendSocVersion + + soc_version = torch_npu.npu.get_soc_version() + if 220 <= soc_version <= 225: + _ascend_soc_version = AscendSocVersion.A2 + elif 250 <= soc_version <= 255: + _ascend_soc_version = AscendSocVersion.A3 + else: + _ascend_soc_version = AscendSocVersion.UNDEFINED + + return _ascend_soc_version == AscendSocVersion.A2 + + def _is_ascend_soc_version_A2_v013_local(): + from vllm_ascend.utils import AscendDeviceType + + soc_version = torch_npu.npu.get_soc_version() + if 220 <= soc_version <= 225: + cur_device_type = AscendDeviceType.A2 + elif 250 <= soc_version <= 255: + cur_device_type = AscendDeviceType.A3 + elif 200 <= soc_version <= 205: + cur_device_type = AscendDeviceType._310P + elif soc_version == 260: + cur_device_type = AscendDeviceType.A5 + else: + raise RuntimeError(f"Can not support soc_version: {soc_version}.") + + return cur_device_type == AscendDeviceType.A2 + + if vllm.__version__ == "0.11.0": + is_A2 = _is_ascend_soc_version_A2_v011_local() + elif vllm.__version__ == "0.13.0": + is_A2 = _is_ascend_soc_version_A2_v013_local() + else: + is_A2 = False + + if is_A2: + VLLM_ASCEND_ENABLE_MATMUL_ALLREDUCE = bool(int(os.getenv("VLLM_ASCEND_ENABLE_MATMUL_ALLREDUCE", "0"))) + if VLLM_ASCEND_ENABLE_MATMUL_ALLREDUCE: + raise AssertionError( + "AscendSocVersion.A2 is not support VLLM_ASCEND_ENABLE_MATMUL_ALLREDUCE \ + in Single-card multi-process scenario now. " + ) + + +def vllm_ascend_v013_select_moe_comm_method_wrapper(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + moe_comm_method = fn(*args, **kwargs) + from vllm_ascend.ascend_forward_context import MoECommType + from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type + + ascend_device_type = get_ascend_device_type() + + # AscendSocVersion.A2 is not support MC2 in Single-card multi-process scenario now. + if ascend_device_type in {AscendDeviceType.A2} and moe_comm_method == MoECommType.MC2: + moe_comm_method = MoECommType.ALLGATHER + + return moe_comm_method + + return wrapper + + +def vllm_ascend_v013_matmul_and_reduce_wrapper(fn): + @wraps(fn) + def wrapper(self, *args, **kwargs): + from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type + + ascend_device_type = get_ascend_device_type() + # AscendSocVersion.A2 is not support MC2 in Single-card multi-process scenario now. + if ascend_device_type in {AscendDeviceType.A2}: + from vllm.forward_context import get_forward_context + + try: + forward_context = get_forward_context() + forward_context.mmrs_fusion = False + except AssertionError: + # forward_context.mmrs_fusion will be false in matmul_and_reduce func. + pass + return fn(self, *args, **kwargs) + + return wrapper + + +def vllm_v013_weight_loader_method_wrapper(fn): + @wraps(fn) + def wrapper(self, param, loaded_weight, weight_name, shard_id, expert_id, return_success=False): + if (shard_id in ("w1", "w3") and param.shape[1] == self.hidden_size) or ( + shard_id == "w2" and param.shape[2] == self.hidden_size + ): + param.data = param.data.transpose(1, 2) + return fn(self, param, loaded_weight, weight_name, shard_id, expert_id, return_success) + + return wrapper + + +def patch_vllm013_rotary_emb(): + from vllm.model_executor.layers.rotary_embedding.common import ApplyRotaryEmb + + def vllm013_npu_rotary_embedding_init_impl( + self, + enforce_enable: bool = False, + is_neox_style: bool = True, + enable_fp32_compute: bool = False, + ) -> None: + super(ApplyRotaryEmb, self).__init__() + self.is_neox_style = is_neox_style + self.enable_fp32_compute = enable_fp32_compute + self.apply_rotary_emb_flash_attn = None + + ApplyRotaryEmb.__init__ = vllm013_npu_rotary_embedding_init_impl + + +if is_torch_npu_available(check_device=False): + import vllm + from packaging import version + + _VLLM_VERSION = version.parse(vllm.__version__) + if _VLLM_VERSION >= version.parse("0.13.0") and _VLLM_VERSION <= version.parse("0.14.0"): + # Disable flash_attn in RotaryEmbedding (NPU) when VLLM >= 0.13 + from vllm.model_executor.layers.fused_moe import FusedMoE + + patch_vllm013_rotary_emb() + FusedMoE.weight_loader = vllm_v013_weight_loader_method_wrapper(FusedMoE.weight_loader) + + VERL_NPU_ENABLE_A2_PATCH_VLLM_ASCEND_MC2 = bool(int(os.getenv("VERL_NPU_ENABLE_A2_PATCH_VLLM_ASCEND_MC2", "1"))) + if VERL_NPU_ENABLE_A2_PATCH_VLLM_ASCEND_MC2: + # only support vllm 0.13 and 0.11 now. + if _VLLM_VERSION >= version.parse("0.13.0") and _VLLM_VERSION <= version.parse("0.14.0"): + from vllm_ascend import ascend_forward_context + from vllm_ascend.ops.linear_op import SequenceRowParallelOp + + ascend_forward_context.select_moe_comm_method = vllm_ascend_v013_select_moe_comm_method_wrapper( + ascend_forward_context.select_moe_comm_method + ) + SequenceRowParallelOp.matmul_and_reduce = vllm_ascend_v013_matmul_and_reduce_wrapper( + SequenceRowParallelOp.matmul_and_reduce + ) + + elif _VLLM_VERSION >= version.parse("0.11.0") and _VLLM_VERSION < version.parse("0.13.0"): + from vllm_ascend.ops.linear_op import SequenceRowParallelOp + from vllm_ascend.worker.model_runner_v1 import NPUModelRunner + + NPUModelRunner._select_moe_comm_method = vllm_ascend_v011_select_moe_comm_method_wrapper( + NPUModelRunner._select_moe_comm_method + ) + SequenceRowParallelOp.matmul_and_reduce = vllm_ascend_v011_matmul_and_reduce_wrapper( + SequenceRowParallelOp.matmul_and_reduce + ) diff --git a/verl/verl/utils/vllm/patch.py b/verl/verl/utils/vllm/patch.py new file mode 100644 index 0000000000000000000000000000000000000000..951c5cad3d616b1c0472d64fee2c232c61973968 --- /dev/null +++ b/verl/verl/utils/vllm/patch.py @@ -0,0 +1,142 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# To support different vLLM versions, we add the model into SUPPORTED_MOE_MODELS separately to avoid triggering +# unsupported issues. +SUPPORTED_MOE_MODELS = [] + +try: + from vllm.model_executor.models.deepseek_v2 import DeepseekV2ForCausalLM, DeepseekV3ForCausalLM + + SUPPORTED_MOE_MODELS.append(DeepseekV2ForCausalLM) + SUPPORTED_MOE_MODELS.append(DeepseekV3ForCausalLM) +except ImportError: + pass + +try: + from vllm.model_executor.models.mixtral import MixtralForCausalLM + + SUPPORTED_MOE_MODELS.append(MixtralForCausalLM) +except ImportError: + pass + +try: + from vllm.model_executor.models.qwen2_moe import Qwen2MoeForCausalLM + + SUPPORTED_MOE_MODELS.append(Qwen2MoeForCausalLM) +except ImportError: + pass + +try: + from vllm.model_executor.models.qwen3_moe import Qwen3MoeForCausalLM + + SUPPORTED_MOE_MODELS.append(Qwen3MoeForCausalLM) +except ImportError: + pass + +try: + from vllm.model_executor.models.qwen3_vl_moe import Qwen3MoeLLMForCausalLM + + SUPPORTED_MOE_MODELS.append(Qwen3MoeLLMForCausalLM) +except ImportError: + pass + +try: + from vllm.model_executor.models.qwen3_next import Qwen3NextForCausalLM + + SUPPORTED_MOE_MODELS.append(Qwen3NextForCausalLM) +except ImportError: + pass + +try: + from vllm.model_executor.models.kimi_vl import KimiVLForConditionalGeneration + + SUPPORTED_MOE_MODELS.append(KimiVLForConditionalGeneration) +except ImportError: + pass + +try: + from vllm.model_executor.models.qwen3_5 import Qwen3_5MoeForCausalLM + + SUPPORTED_MOE_MODELS.append(Qwen3_5MoeForCausalLM) +except ImportError: + pass + + +def patch_vllm_moe_model_weight_loader(model): + # this is a work around to load the weight of vllm fused moe model + # it is from a bug from vllm 0.8.2 + # all the weights are supposed to have a weight_loader, but the moe weights + # do not have a weight_loader, so we need to patch it + # (True, 'model.embed_tokens.weight') + # (True, 'model.layers.0.self_attn.qkv_proj.weight') + # (True, 'model.layers.0.self_attn.qkv_proj.bias') + # (True, 'model.layers.0.self_attn.o_proj.weight') + # (True, 'model.layers.0.mlp.gate.weight') + # (True, 'model.layers.0.mlp.shared_expert.gate_up_proj.weight') + # (True, 'model.layers.0.mlp.shared_expert.down_proj.weight') + # (False, 'model.layers.0.mlp.shared_expert_gate.weight') use default + # (False, 'model.layers.0.input_layernorm.weight') use default + # (False, 'model.layers.0.post_attention_layernorm.weight') use default + # (False, 'model.layers.0.mlp.experts.w13_weight') use mlp.experts.weight_loader + # (False, 'model.layers.0.mlp.experts.w2_weight') use mlp.experts.weight_loader + + # Early return if no MOE models are supported + if not SUPPORTED_MOE_MODELS: + return + + original_model_type = type(model) + if hasattr(model, "runnable") and "ACLGraphWrapper" in str(original_model_type): + model = model.runnable + original_model_type = type(model) + + # Define MLP attribute mapping for different model types + MLP_ATTR_MAPPING = {} + try: + from vllm.model_executor.models.mixtral import MixtralForCausalLM + + MLP_ATTR_MAPPING[MixtralForCausalLM] = "block_sparse_moe" + except ImportError: + pass + + DEFAULT_MLP_ATTR = "mlp" + + # Get inner model (either model.model or model.language_model) + inner_model = getattr(model, "model", None) or getattr(model, "language_model", None) + if inner_model is None: + raise ValueError("The provided model does not have a valid 'model' or 'language_model' attribute.") + + if not isinstance(model, tuple(SUPPORTED_MOE_MODELS)) and not isinstance(inner_model, tuple(SUPPORTED_MOE_MODELS)): + return + + # TODO(@leisuzz): class Qwen3MoeLLMForCausalLM is not available if VLLM version < 0.11.0, + # will update the 'if statement' with 'isinstance' when verl commonly use VLLM version >= 0.11.0 + if type(inner_model).__name__ in ("Qwen3MoeLLMForCausalLM", "Qwen3_5MoeForCausalLM"): + inner_model = inner_model.model # Reassign inner_model in Qwen3-vl + + for layer_idx, layer in enumerate(inner_model.layers): + mlp_attr = MLP_ATTR_MAPPING.get(original_model_type, DEFAULT_MLP_ATTR) + + mlp = getattr(layer, mlp_attr, None) + if not mlp: + continue + + experts = getattr(mlp, "experts", None) + if not experts or not hasattr(experts, "weight_loader"): + continue + + # Patch the weight loaders + for name, param in mlp.named_parameters(): + if "w13_weight" in name or "w2_weight" in name: + param.weight_loader = experts.weight_loader diff --git a/verl/verl/utils/vllm/utils.py b/verl/verl/utils/vllm/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..1ac655fcf603b660a28ed56c93f0fd2d4117f0e6 --- /dev/null +++ b/verl/verl/utils/vllm/utils.py @@ -0,0 +1,128 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from msgspec import field +from packaging import version as vs + +try: + from vllm.lora.lora_model import LoRAModel +except ImportError: + from vllm.lora.models import LoRAModel + +from vllm.lora.request import LoRARequest +from vllm.lora.utils import get_adapter_absolute_path +from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager + +from verl.third_party.vllm import get_version + + +class TensorLoRARequest(LoRARequest): + peft_config: dict = field(default=None) + lora_tensors: dict = field(default=None) + + +class VLLMHijack: + @staticmethod + def hijack(): + def hijack__load_adapter(self, lora_request: TensorLoRARequest) -> LoRAModel: + """ + based on vllm.lora.worker_manager.WorkerLoRAManager._load_adapter, support load adapter with lora tensors + + Reason: + VLLM does not support adding LoRA from tensors directly. It only supports adding LoRA via file paths. + To synchronize the LoRA tensors of the actor model, we need to find a workaround to enable VLLM to + load memory-based LoRA tensors. + """ + try: + supported_lora_modules = self._adapter_manager.supported_lora_modules + packed_modules_mapping = self._adapter_manager.packed_modules_mapping + expected_lora_modules: list[str] = [] + for module in supported_lora_modules: + if module in packed_modules_mapping: + expected_lora_modules.extend(packed_modules_mapping[module]) + else: + expected_lora_modules.append(module) + + expected_lora_modules = list(set(expected_lora_modules)) + + lora_tensors = None + from vllm.lora.peft_helper import PEFTHelper + + if isinstance(lora_request, TensorLoRARequest): + peft_config = lora_request.peft_config + lora_tensors = lora_request.lora_tensors + peft_helper = PEFTHelper.from_dict(peft_config) + else: + lora_path = get_adapter_absolute_path(lora_request.lora_path) + + peft_helper = PEFTHelper.from_local_dir(lora_path, self.max_position_embeddings) + + # Validates the LoRA configuration against requirements before + # loading weights, throwing an exception if validation fails. + peft_helper.validate_legal(self.lora_config) + + # For some models like Qwen2VL, we need to use hf_to_vllm_mapper + # to ensure correct loading of lora weights. + model = self._adapter_manager.model + hf_to_vllm_mapper = None + if hasattr(model, "hf_to_vllm_mapper") and model.hf_to_vllm_mapper is not None: + hf_to_vllm_mapper = model.hf_to_vllm_mapper + + lora_request_kwargs = { + "peft_helper": peft_helper, + "lora_model_id": lora_request.lora_int_id, + "device": "cpu", + "dtype": self.lora_config.lora_dtype, + "weights_mapper": hf_to_vllm_mapper, + } + if hasattr(self, "embedding_padding_modules"): + lora_request_kwargs["embedding_modules"] = self.embedding_modules + lora_request_kwargs["embedding_padding_modules"] = self.embedding_padding_modules + else: + lora_request_kwargs["model_vocab_size"] = self.vocab_size + if hasattr(self.lora_config, "lora_extra_vocab_size"): + lora_request_kwargs["target_embedding_padding"] = ( + self.vocab_size + self.lora_config.lora_extra_vocab_size + ) + if isinstance(lora_request, TensorLoRARequest): + lora = self._lora_model_cls.from_lora_tensors( + tensors=lora_tensors, + **lora_request_kwargs, + ) + else: + lora = self._lora_model_cls.from_local_checkpoint( + lora_path, + expected_lora_modules, + **lora_request_kwargs, + ) + except Exception: + raise + + if getattr(lora, "extra_vocab_size", 0) > getattr(self.lora_config, "lora_extra_vocab_size", 0): + raise ValueError( + f"LoRA added vocab size {lora.extra_vocab_size} is greater than lora_extra_vocab_size " + f"{self.lora_config.lora_extra_vocab_size}." + ) + return lora + + def do_hijack(target_cls, target_method_name, hooking_method): + setattr(target_cls, target_method_name, hooking_method) + + do_hijack(LRUCacheWorkerLoRAManager, "_load_adapter", hijack__load_adapter) + + +def is_version_ge(pkg: str = "vllm", minver: str = "0.7.3"): + """check if the package version is greater than or equal to the minimum version""" + return vs.parse(get_version(pkg)) >= vs.parse(minver) diff --git a/verl/verl/utils/vllm/vllm_fp8_utils.py b/verl/verl/utils/vllm/vllm_fp8_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0abee3f33ea6fd816b5323c488184f5e0aeca7cd --- /dev/null +++ b/verl/verl/utils/vllm/vllm_fp8_utils.py @@ -0,0 +1,664 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +import logging +from dataclasses import dataclass, field +from unittest.mock import patch + +import torch +import vllm +from packaging import version + +try: + from vllm.model_executor.layers.fused_moe.layer import FusedMoE + from vllm.model_executor.layers.linear import LinearBase +except ImportError as e: + raise ImportError("FP8 quantization not available") from e + +from verl.utils.kernel.fp8_kernel import scaled_fp8_blockwise + +logger = logging.getLogger(__name__) + + +# Ref: https://github.com/NVIDIA-NeMo/RL/commit/bc24887c72a6e1b2699a228bc87c588546dfe6b7 +@dataclass() +class FP8State: + # A cache of fp8 parameter names, we can check this cache to see if a + # param name corresponds to a fp8 weight + seen_params: set = field(default_factory=lambda: set()) + fp8_param_names: set = field(default_factory=lambda: set()) + vllm_patches: list = field(default_factory=lambda: []) + + +fp8_state: FP8State = FP8State() + + +def is_fp8_model(vllm_config): + from vllm.model_executor.layers.quantization.fp8 import Fp8Config + + if hasattr(vllm_config, "quant_config"): + if isinstance(vllm_config.quant_config, Fp8Config): + return True + elif is_mxfp8_vllm_ascend(vllm_config.quant_config): + return True + + return False + + +def get_module_from_param_name(model, name: str): + # Split the name into parts (e.g., 'layers', '0', 'self_attn', 'q_proj', 'weight') + # The module path is all but the last part (the parameter's own name) + path_parts = name.split(".") + module_path = path_parts[:-1] + # Replace with the fused model name + packed_modules_mapping = model.packed_modules_mapping + reversed_mapping = { + original_name: fused_name + for fused_name, original_names_list in packed_modules_mapping.items() + for original_name in original_names_list + } + if module_path[-1] in reversed_mapping.keys(): + module_path[-1] = reversed_mapping[module_path[-1]] + + current_module = model + try: + # Traverse the model hierarchy + for part in module_path: + if isinstance(current_module, FusedMoE): + return current_module + elif isinstance(current_module, torch.nn.ModuleList): + current_module = current_module[int(part)] + else: + current_module = getattr(current_module, part) + except (AttributeError, IndexError, ValueError) as e: + print(f"Warning: Could not find module for parameter '{name}'. Error: {e}") + return current_module + + +def is_fp8_weight(name, model): + if name not in fp8_state.seen_params: + fp8_state.seen_params.add(name) + # Filter out bias params + if name.endswith("weight"): + module = get_module_from_param_name(model, name) + # We currently only quantize linear layers + + if (isinstance(module, LinearBase) and module.weight.dtype == torch.float8_e4m3fn) or ( + isinstance(module, FusedMoE) + and module.w13_weight.dtype == torch.float8_e4m3fn + and module.w2_weight.dtype == torch.float8_e4m3fn + ): + fp8_state.fp8_param_names.add(name) + return name in fp8_state.fp8_param_names + + +def is_mxfp8_vllm_ascend(quant_config): + try: + from vllm_ascend.quantization.modelslim_config import AscendModelSlimConfig + from vllm_ascend.quantization.quant_config import AscendQuantConfig + + if isinstance(quant_config, AscendModelSlimConfig) or isinstance(quant_config, AscendQuantConfig): + quant_method = quant_config.quant_description.get("quant_method") + return quant_method in ["ascend"] + return False + except ImportError: + # vllm_ascend not installed, so this can't be an Ascend MXFP8 config + return False + + +def restore_mxfp8_weights_for_loading(model): + for name, module in model.named_modules(): + if ( + hasattr(module, "_mxfp8_transformed") + and hasattr(module, "quant_method") + and hasattr(module.quant_method, "quant_method") + and hasattr(module.quant_method.quant_method, "restore_weights_for_rl_loading") + ): + module.quant_method.quant_method.restore_weights_for_rl_loading(module) + + +def apply_mxfp8_transformation_after_loading(model): + """Re-apply MXFP8 transformations after weight loading. + + This function iterates through all linear modules in the model and applies + the MXFP8 transformations (transpose, reshape) that are required for NPU + inference. + + Must be called AFTER model.load_weights() in RL training loops. + """ + try: + from vllm.model_executor.layers.linear import LinearBase + except ImportError: + logger.warning("Could not import LinearBase, skipping MXFP8 transformation") + return + + for name, module in model.named_modules(): + if (isinstance(module, LinearBase) or isinstance(module, FusedMoE)) and hasattr( + module, "_mxfp8_original_shapes" + ): + if hasattr(module, "quant_method") and hasattr(module.quant_method, "process_weights_after_loading"): + logger.debug(f"Applying MXFP8 transformation for module: {name}") + module.quant_method.process_weights_after_loading(module) + + +def quant_weights(weights, model, quant_config, dtype=torch.bfloat16): + """Quantize weights to FP8 format using a memory-efficient generator. + + + Args: + weights: Generator or iterable of (name, tensor) pairs + model: The model to check for FP8 weight names + quant_config: Quantization configuration with weight_block_size + dtype: Data type for intermediate computation (default: bfloat16) + + Yields: + Tuples of (name, tensor) for each weight and its scale + """ + + is_mxfp8_npu = is_mxfp8_vllm_ascend(quant_config) + if is_mxfp8_npu: + import torch_npu + # vLLM v0.11-v0.12 renamed weight_scale_inv → weight_scale in process_weights_after_loading, + # so load_weights expects "_scale" suffix. v0.14+ keeps weight_scale_inv, so expects "_scale_inv". + _use_scale_not_scale_inv = version.parse("0.11.0") <= version.parse(vllm.__version__) < version.parse("0.14.0") + + for k, v in weights: + if not is_fp8_weight(k, model): + yield (k, v) + continue + + # Cast the weight into fp8 and its scale factor + if torch.distributed.get_rank() == 0: + logger.debug(f"Quantizing to FP8 blockwise: {k}") + if is_mxfp8_npu: + param_lp, param_scale = torch_npu.npu_dynamic_mx_quant( + v.to(dtype), + axis=-1, + dst_type=torch_npu.float8_e4m3fn, + ) + param_scale = param_scale.flatten(-2, -1) + else: + param_lp, param_scale = scaled_fp8_blockwise( + v.to(dtype), + weight_block_size=quant_config.weight_block_size, + ) + param_scale = param_scale.squeeze(-1) + + # Yield the quantized weight + yield (k, param_lp) + + # Yield the scale with appropriate naming based on vLLM version + if is_mxfp8_npu: + yield (k + "_scale", param_scale) + elif _use_scale_not_scale_inv and "expert" not in k: + yield (k + "_scale", param_scale) + else: + yield (k + "_scale_inv", param_scale) + + # Explicitly delete original tensor reference to help GC + del v, param_lp, param_scale + + +def load_quanted_weights(weights, model_runner): + model = model_runner.model + quant_config = model_runner.vllm_config.quant_config + vllm_dtype = model_runner.vllm_config.model_config.dtype + + is_mxfp8_npu = is_mxfp8_vllm_ascend(quant_config) + + if is_mxfp8_npu: + # For MXFP8 on NPU, we need to restore weights to original shapes + # before loading, then re-apply transformation after loading. + # This is because process_weights_after_loading transposes the weights, + # but the weight_loader expects original shapes. + restore_mxfp8_weights_for_loading(model) + + weights_quantized = quant_weights(weights, model, quant_config, dtype=vllm_dtype) + + # Monkey patch the param class to their subclass, as certain models + # will check the param type to call the proper weightloader + for name, param in model.named_parameters(): + if hasattr(param, "subclass_type"): + param.orig_type = param.__class__ + param.__class__ = param.subclass_type + # Finally load the weights into vllm + loaded_params = model.load_weights(weights_quantized) + # Undo the type change above to the original type + for name, param in model.named_parameters(): + if hasattr(param, "subclass_type"): + param.__class__ = param.orig_type + + if is_mxfp8_npu: + # Re-apply MXFP8 transformations after weight loading + apply_mxfp8_transformation_after_loading(model) + + return loaded_params + + +def process_weights_after_loading_for_vllm10(self, layer) -> None: + """This function is used to process the weights after loading for a Linear layer, it is used for vllm v0.10 + + Compared to the original process_weights_after_loading in vllm, we just avoid creation of + new torch.nn.Parameter objects, because that removes the weight_loader attribute which we need for refit. + """ + logger.debug("Applying patch process_weights_after_loading") + try: + from vllm.model_executor.parameter import ( + BlockQuantScaleParameter, + ModelWeightParameter, + ) + except Exception: + print("error") + from torch.nn import Parameter + + def _create_param_from_subclass_attributes(custom_param): + param = Parameter(custom_param.data, requires_grad=False) + base_param_dir = dir(torch.nn.Parameter) + custom_param_dir = dir(custom_param) + # Find the attributes that are unique to the custom parameter + custom_attributes = [ + attr for attr in custom_param_dir if attr not in base_param_dir and not attr.startswith("__") + ] + # Set the custom attributes into the base parameter object + for attr in custom_attributes: + setattr(param, attr, getattr(custom_param, attr)) + + param.subclass_type = type(custom_param) + return param + + assert self.block_quant and self.quant_config.is_checkpoint_fp8_serialized + assert self.quant_config.activation_scheme == "dynamic" + weight = layer.weight.data + weight_scale_inv = layer.weight_scale_inv.data + weight = self._maybe_pad_weight(weight) + + layer.weight = _create_param_from_subclass_attributes( + ModelWeightParameter( + data=weight, + output_dim=0, + input_dim=1, + weight_loader=layer.weight.weight_loader, + ) + ) + layer.weight_scale_inv = _create_param_from_subclass_attributes( + BlockQuantScaleParameter( + data=weight_scale_inv, + output_dim=0, + input_dim=1, + weight_loader=layer.weight_scale_inv.weight_loader, + ) + ) + + +def process_weights_after_loading_for_vllm11(self, layer) -> None: + """This function is used to process the weights after loading for a Linear layer, it is used for vllm 0.11 + + Compared to the original process_weights_after_loading in vllm, we just avoid creation of + new torch.nn.Parameter objects, because that removes the weight_loader attribute which we need for refit. + """ + from torch.nn import Parameter + from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + maybe_post_process_fp8_weight_block, + process_fp8_weight_block_strategy, + ) + from vllm.model_executor.parameter import ( + BlockQuantScaleParameter, + ModelWeightParameter, + ) + + assert self.block_quant and self.quant_config.is_checkpoint_fp8_serialized + assert self.quant_config.activation_scheme == "dynamic" + + def _create_param_from_subclass_attributes(custom_param): + param = Parameter(custom_param.data, requires_grad=False) + base_param_dir = dir(torch.nn.Parameter) + custom_param_dir = dir(custom_param) + # Find the attributes that are unique to the custom parameter + custom_attributes = [ + attr for attr in custom_param_dir if attr not in base_param_dir and not attr.startswith("__") + ] + # Set the custom attributes into the base parameter object + for attr in custom_attributes: + setattr(param, attr, getattr(custom_param, attr)) + + param.subclass_type = type(custom_param) + return param + + weight_scale = layer.weight_scale_inv if hasattr(layer, "weight_scale_inv") else layer.weight_scale + weight, weight_scale = process_fp8_weight_block_strategy(layer.weight, weight_scale) + + layer.weight = _create_param_from_subclass_attributes( + ModelWeightParameter( + data=weight.data, + output_dim=0, + input_dim=1, + weight_loader=layer.weight.weight_loader, + ) + ) + layer.weight_scale = _create_param_from_subclass_attributes( + BlockQuantScaleParameter( + data=weight_scale.data, + output_dim=0, + input_dim=1, + weight_loader=layer.weight_scale_inv.weight_loader, + ) + ) + + del layer.weight_scale_inv + + if version.parse(vllm.__version__) == version.parse("0.11.0"): + maybe_post_process_fp8_weight_block(layer, self.cutlass_block_fp8_supported) + else: + maybe_post_process_fp8_weight_block(layer) + + +def process_weights_after_loading_for_vllm14(self, layer) -> None: + """process_weights_after_loading for vLLM >= 0.14. + + Starting from v0.14, vLLM keeps the scale parameter as `weight_scale_inv` + (instead of renaming it to `weight_scale` like v0.11-v0.12), and `apply()` + accesses `layer.weight_scale_inv`. We preserve `weight_loader` and + `subclass_type` attributes so that refit (repeated weight sync) works. + """ + from torch.nn import Parameter + from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + maybe_post_process_fp8_weight_block, + process_fp8_weight_block_strategy, + ) + from vllm.model_executor.parameter import ( + BlockQuantScaleParameter, + ModelWeightParameter, + ) + + assert self.block_quant and self.quant_config.is_checkpoint_fp8_serialized + assert self.quant_config.activation_scheme == "dynamic" + + def _create_param_from_subclass_attributes(custom_param): + param = Parameter(custom_param.data, requires_grad=False) + base_param_dir = dir(torch.nn.Parameter) + custom_param_dir = dir(custom_param) + custom_attributes = [ + attr for attr in custom_param_dir if attr not in base_param_dir and not attr.startswith("__") + ] + for attr in custom_attributes: + setattr(param, attr, getattr(custom_param, attr)) + + param.subclass_type = type(custom_param) + return param + + weight, weight_scale_inv = process_fp8_weight_block_strategy(layer.weight, layer.weight_scale_inv) + + layer.weight = _create_param_from_subclass_attributes( + ModelWeightParameter( + data=weight.data, + output_dim=0, + input_dim=1, + weight_loader=layer.weight.weight_loader, + ) + ) + layer.weight_scale_inv = _create_param_from_subclass_attributes( + BlockQuantScaleParameter( + data=weight_scale_inv.data, + output_dim=0, + input_dim=1, + weight_loader=layer.weight_scale_inv.weight_loader, + ) + ) + + # vLLM v0.17 removed the `else: register_parameter("input_scale", None)` from + # create_weights() for dynamic activation, but apply() still accesses layer.input_scale. + # Since block_quant always uses dynamic activation, ensure the attribute exists. + if not hasattr(layer, "input_scale"): + layer.input_scale = None + + maybe_post_process_fp8_weight_block(layer) + + +def process_weights_after_loading_moe_for_vllm10(self, layer) -> None: + """This function is used to process the weights after loading for a FusedMoE layer, it is used for vllm v0.10""" + from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import is_rocm_aiter_moe_enabled + from vllm.model_executor.layers.quantization.fp8 import _is_col_major, _swap_w13_to_w31 + from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + get_col_major_tma_aligned_tensor, + requant_weight_ue8m0_inplace, + ) + from vllm.utils.deep_gemm import is_blackwell_deep_gemm_used + + self.rocm_aiter_moe_enabled = is_rocm_aiter_moe_enabled() + assert self.quant_config.activation_scheme == "dynamic" + if self.flashinfer_moe_enabled: + w13_weight = _swap_w13_to_w31(layer.w13_weight.data) + w13_weight_scale_inv = _swap_w13_to_w31(layer.w13_weight_scale_inv.data) + w2_weight = layer.w2_weight.data + w2_weight_scale_inv = layer.w2_weight_scale_inv.data + else: + w13_weight = layer.w13_weight.data + w13_weight_scale_inv = layer.w13_weight_scale_inv.data + w2_weight = layer.w2_weight + w2_weight_scale_inv = layer.w2_weight_scale_inv + + from torch.nn import Parameter + + def _create_param_from_subclass_attributes(custom_data, custom_weight): + param = Parameter(custom_data, requires_grad=False) + base_param_dir = dir(torch.nn.Parameter) + custom_weight_dir = dir(custom_weight) + # Find the attributes that are unique to the custom parameter + custom_attributes = [ + attr for attr in custom_weight_dir if attr not in base_param_dir and not attr.startswith("__") + ] + # Set the custom attributes into the base parameter object + for attr in custom_attributes: + setattr(param, attr, getattr(custom_weight, attr)) + + return param + + layer.w13_weight = _create_param_from_subclass_attributes(w13_weight, layer.w13_weight) + layer.w13_weight_scale_inv = _create_param_from_subclass_attributes( + w13_weight_scale_inv, layer.w13_weight_scale_inv + ) + layer.w2_weight = _create_param_from_subclass_attributes(w2_weight, layer.w2_weight) + layer.w2_weight_scale_inv = _create_param_from_subclass_attributes(w2_weight_scale_inv, layer.w2_weight_scale_inv) + + # DeepGemm scales need to be transposed and aligned. We try to do + # it ahead of time for performance reasons. + if self.allow_deep_gemm and not is_blackwell_deep_gemm_used(): + # Lazy import to avoid CUDA initialization problems. + if _is_col_major(layer.w13_weight_scale_inv): + layer.w13_weight_scale_inv = get_col_major_tma_aligned_tensor(layer.w13_weight_scale_inv).contiguous() + if _is_col_major(layer.w2_weight_scale_inv): + layer.w2_weight_scale_inv = get_col_major_tma_aligned_tensor(layer.w2_weight_scale_inv).contiguous() + + if is_blackwell_deep_gemm_used(): + assert layer.weight_block_size is not None + # Re-quantise the expert weights so their scales are UE8M0. + block_sz = tuple(layer.weight_block_size) + requant_weight_ue8m0_inplace( + layer.w13_weight.data, + layer.w13_weight_scale_inv.data, + block_sz, + ) + requant_weight_ue8m0_inplace( + layer.w2_weight.data, + layer.w2_weight_scale_inv.data, + block_sz, + ) + + if _is_col_major(layer.w13_weight_scale_inv): + layer.w13_weight_scale_inv = get_col_major_tma_aligned_tensor(layer.w13_weight_scale_inv).contiguous() + if _is_col_major(layer.w2_weight_scale_inv): + layer.w2_weight_scale_inv = get_col_major_tma_aligned_tensor(layer.w2_weight_scale_inv).contiguous() + + +def process_weights_after_loading_moe_for_vllm11(self, layer) -> None: + """This function is used to process the weights after loading for a FusedMoE layer, it is used for vllm 0.11""" + from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + swap_w13_to_w31, + ) + from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + expert_weight_is_col_major, + requant_weight_ue8m0_inplace, + ) + from vllm.utils.deep_gemm import ( + get_col_major_tma_aligned_tensor, + is_deep_gemm_e8m0_used, + ) + + try: + from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import is_rocm_aiter_moe_enabled + + self.rocm_aiter_moe_enabled = is_rocm_aiter_moe_enabled() + except ImportError: + from vllm._aiter_ops import rocm_aiter_ops + + self.rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled() + + assert self.block_quant and self.quant_config.is_checkpoint_fp8_serialized + assert self.quant_config.activation_scheme == "dynamic" + + if self.flashinfer_moe_backend is not None: + layer.w13_weight.data = swap_w13_to_w31(layer.w13_weight.data) + layer.w13_weight_scale_inv.data = swap_w13_to_w31(layer.w13_weight_scale_inv.data) + + if self.allow_deep_gemm and not is_deep_gemm_e8m0_used(): + if expert_weight_is_col_major(layer.w13_weight_scale_inv): + layer.w13_weight_scale_inv = get_col_major_tma_aligned_tensor(layer.w13_weight_scale_inv) + if expert_weight_is_col_major(layer.w2_weight_scale_inv): + layer.w2_weight_scale_inv = get_col_major_tma_aligned_tensor(layer.w2_weight_scale_inv) + + if is_deep_gemm_e8m0_used(): + assert layer.weight_block_size is not None + # Re-quantise the expert weights so their scales are UE8M0. + block_sz = tuple(layer.weight_block_size) + requant_weight_ue8m0_inplace( + layer.w13_weight.data, + layer.w13_weight_scale_inv.data, + block_sz, + ) + requant_weight_ue8m0_inplace( + layer.w2_weight.data, + layer.w2_weight_scale_inv.data, + block_sz, + ) + + # Ensure column-major TMA alignment expected by DeepGEMM. + if expert_weight_is_col_major(layer.w13_weight_scale_inv): + layer.w13_weight_scale_inv = get_col_major_tma_aligned_tensor(layer.w13_weight_scale_inv) + if expert_weight_is_col_major(layer.w2_weight_scale_inv): + layer.w2_weight_scale_inv = get_col_major_tma_aligned_tensor(layer.w2_weight_scale_inv) + + +def process_weights_after_loading_moe_for_vllm14(self, layer) -> None: + # removed the reentrancy guard here for refit + from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( + convert_to_fp8_moe_kernel_format, + make_fp8_moe_kernel, + ) + + # Allow for accessing weights and scales in standard way. + w13 = layer.w13_weight + w2 = layer.w2_weight + w13_scale = getattr(layer, f"w13_{self.weight_scale_name}") + w2_scale = getattr(layer, f"w2_{self.weight_scale_name}") + w13_input_scale = layer.w13_input_scale + w2_input_scale = layer.w2_input_scale + + # Shuffle weights to runtime format and setup kernel. + w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format( + fp8_backend=self.fp8_backend, + layer=layer, + w13=w13, + w2=w2, + w13_scale=w13_scale, + w2_scale=w2_scale, + w13_input_scale=w13_input_scale, + w2_input_scale=w2_input_scale, + ) + from torch.nn import Parameter + + def _create_param_from_subclass_attributes(custom_data, custom_weight): + param = Parameter(custom_data, requires_grad=False) + base_param_dir = dir(torch.nn.Parameter) + custom_weight_dir = dir(custom_weight) + # Find the attributes that are unique to the custom parameter + custom_attributes = [ + attr for attr in custom_weight_dir if attr not in base_param_dir and not attr.startswith("__") + ] + # Set the custom attributes into the base parameter object + for attr in custom_attributes: + setattr(param, attr, getattr(custom_weight, attr)) + + return param + + # Replace parameters with updated versions. Note that this helper + # function ensures the replacement is compatible with RL weight reloads. + layer.w13_weight = _create_param_from_subclass_attributes(w13, layer.w13_weight) + layer.w2_weight = _create_param_from_subclass_attributes(w2, layer.w2_weight) + layer.w13_weight_scale_inv = _create_param_from_subclass_attributes(w13_scale, layer.w13_weight_scale_inv) + layer.w2_weight_scale_inv = _create_param_from_subclass_attributes(w2_scale, layer.w2_weight_scale_inv) + + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + if self.moe_quant_config: + assert self.experts_cls is not None + + # Check for the new API by inspecting the function signature, which is more + # robust than version string comparison, especially for dev/pre-release versions. + sig = inspect.signature(make_fp8_moe_kernel) + if "routing_tables" in sig.parameters: + # vLLM >= 0.16+: routing_tables/shared_experts added, returns kernel directly + self.moe_kernel = make_fp8_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + fp8_backend=self.fp8_backend, + experts_cls=self.experts_cls, + routing_tables=layer._maybe_init_expert_routing_tables(), + shared_experts=layer.shared_experts, + ) + else: + # vLLM 0.14/0.15: routing_tables/shared_experts not supported, returns (kernel, use_inplace) + self.kernel, self.use_inplace = make_fp8_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + fp8_backend=self.fp8_backend, + experts_cls=self.experts_cls, + ) + + +def apply_vllm_fp8_patches(): + logger.info("Applying vllm fp8 patches for blockwise quantization") + vllm_ver = version.parse(vllm.__version__) + + # Linear patch: v0.14+ keeps weight_scale_inv, v0.11-v0.12 renames to weight_scale + func1_path = "vllm.model_executor.layers.quantization.fp8.Fp8LinearMethod.process_weights_after_loading" + if vllm_ver >= version.parse("0.14.0"): + linear_patch_fn = process_weights_after_loading_for_vllm14 + elif vllm_ver >= version.parse("0.11.0"): + linear_patch_fn = process_weights_after_loading_for_vllm11 + else: + linear_patch_fn = process_weights_after_loading_for_vllm10 + patcher1 = patch(func1_path, linear_patch_fn) + patcher1.start() + + # MoE patch + func2_path = "vllm.model_executor.layers.quantization.fp8.Fp8MoEMethod.process_weights_after_loading" + if vllm_ver >= version.parse("0.14.0"): + moe_patch_fn = process_weights_after_loading_moe_for_vllm14 + elif vllm_ver >= version.parse("0.11.0"): + moe_patch_fn = process_weights_after_loading_moe_for_vllm11 + else: + moe_patch_fn = process_weights_after_loading_moe_for_vllm10 + patcher2 = patch(func2_path, moe_patch_fn) + patcher2.start() diff --git a/verl/verl/version/version b/verl/verl/version/version new file mode 100644 index 0000000000000000000000000000000000000000..7188dbafb438572b3bd7e02ee7ab16529b1be225 --- /dev/null +++ b/verl/verl/version/version @@ -0,0 +1 @@ +0.8.0.dev diff --git a/verl/verl/workers/__init__.py b/verl/verl/workers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1ce90c5eb352d85c59105c0dc85b5f1dd576f095 --- /dev/null +++ b/verl/verl/workers/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/verl/verl/workers/__pycache__/__init__.cpython-312.pyc b/verl/verl/workers/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..680a5a8853324d1cd525238ed724061b750b0fdc Binary files /dev/null and b/verl/verl/workers/__pycache__/__init__.cpython-312.pyc differ diff --git a/verl/verl/workers/__pycache__/engine_workers.cpython-312.pyc b/verl/verl/workers/__pycache__/engine_workers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d22e6c66cd35ae16bb093bb8b9b8b6b47ca1cdb7 Binary files /dev/null and b/verl/verl/workers/__pycache__/engine_workers.cpython-312.pyc differ diff --git a/verl/verl/workers/config/__init__.py b/verl/verl/workers/config/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..29f621b789db314a9b9f125cfbef795254514b0a --- /dev/null +++ b/verl/verl/workers/config/__init__.py @@ -0,0 +1,36 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from . import actor, critic, disaggregation, engine, model, optimizer, reward, rollout +from .actor import * # noqa: F401 +from .critic import * # noqa: F401 +from .disaggregation import * # noqa: F401 +from .distillation import * # noqa: F401 +from .engine import * # noqa: F401 +from .model import * # noqa: F401 +from .optimizer import * # noqa: F401 +from .reward import * # noqa: F401 +from .rollout import * # noqa: F401 + +__all__ = ( + actor.__all__ + + critic.__all__ + + reward.__all__ + + engine.__all__ + + optimizer.__all__ + + rollout.__all__ + + model.__all__ + + distillation.__all__ + + disaggregation.__all__ +) diff --git a/verl/verl/workers/config/__pycache__/__init__.cpython-312.pyc b/verl/verl/workers/config/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6d040c828e1b3e4f9197bc82b3e007b4f8f28d21 Binary files /dev/null and b/verl/verl/workers/config/__pycache__/__init__.cpython-312.pyc differ diff --git a/verl/verl/workers/config/__pycache__/actor.cpython-312.pyc b/verl/verl/workers/config/__pycache__/actor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f9c8ae3c57dd9a784b110f19a66f1b20f1223298 Binary files /dev/null and b/verl/verl/workers/config/__pycache__/actor.cpython-312.pyc differ diff --git a/verl/verl/workers/config/__pycache__/critic.cpython-312.pyc b/verl/verl/workers/config/__pycache__/critic.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..57803fff5802778ee0c542c030719a1b20fa0d6c Binary files /dev/null and b/verl/verl/workers/config/__pycache__/critic.cpython-312.pyc differ diff --git a/verl/verl/workers/config/__pycache__/disaggregation.cpython-312.pyc b/verl/verl/workers/config/__pycache__/disaggregation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d50e1cdca5ccce9eb7985b76d1be0a0fe13c18b2 Binary files /dev/null and b/verl/verl/workers/config/__pycache__/disaggregation.cpython-312.pyc differ diff --git a/verl/verl/workers/config/__pycache__/distillation.cpython-312.pyc b/verl/verl/workers/config/__pycache__/distillation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cff59b1e94871431cabbec792767ec14374947bf Binary files /dev/null and b/verl/verl/workers/config/__pycache__/distillation.cpython-312.pyc differ diff --git a/verl/verl/workers/config/__pycache__/engine.cpython-312.pyc b/verl/verl/workers/config/__pycache__/engine.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a3cc2ac464951a04e2e85a7de24aca317014905e Binary files /dev/null and b/verl/verl/workers/config/__pycache__/engine.cpython-312.pyc differ diff --git a/verl/verl/workers/config/__pycache__/model.cpython-312.pyc b/verl/verl/workers/config/__pycache__/model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..af58d71b518a9be4e63e76efbb10bc94709010a9 Binary files /dev/null and b/verl/verl/workers/config/__pycache__/model.cpython-312.pyc differ diff --git a/verl/verl/workers/config/__pycache__/optimizer.cpython-312.pyc b/verl/verl/workers/config/__pycache__/optimizer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..869569d8a23e1d99bef993869e04cc491f68b251 Binary files /dev/null and b/verl/verl/workers/config/__pycache__/optimizer.cpython-312.pyc differ diff --git a/verl/verl/workers/config/__pycache__/reward.cpython-312.pyc b/verl/verl/workers/config/__pycache__/reward.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..750fb5d6e418ded0988cba8eb1ada1bb9f4cf042 Binary files /dev/null and b/verl/verl/workers/config/__pycache__/reward.cpython-312.pyc differ diff --git a/verl/verl/workers/config/__pycache__/rollout.cpython-312.pyc b/verl/verl/workers/config/__pycache__/rollout.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..17549aba747dd995ac1b19ecc4c332fa111ff975 Binary files /dev/null and b/verl/verl/workers/config/__pycache__/rollout.cpython-312.pyc differ diff --git a/verl/verl/workers/config/actor.py b/verl/verl/workers/config/actor.py new file mode 100644 index 0000000000000000000000000000000000000000..32149f02375394f1e2b0d1caabe5cf088fd42556 --- /dev/null +++ b/verl/verl/workers/config/actor.py @@ -0,0 +1,404 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass, field +from typing import Any, Optional + +from omegaconf import MISSING + +from verl.base_config import BaseConfig +from verl.trainer.config import CheckpointConfig, RolloutCorrectionConfig +from verl.utils.profiler.config import ProfilerConfig +from verl.utils.qat import QATConfig + +from .engine import ( + FSDPEngineConfig, + McoreEngineConfig, + MindSpeedEngineConfig, + TorchtitanEngineConfig, + VeOmniEngineConfig, +) +from .model import HFModelConfig +from .optimizer import OptimizerConfig + +__all__ = [ + "PolicyLossConfig", + "RouterReplayConfig", + "ActorConfig", + "FSDPActorConfig", + "McoreActorConfig", + "VeOmniActorConfig", + "QATConfig", + "TorchTitanActorConfig", + "MindSpeedActorConfig", +] + + +@dataclass +class RouterReplayConfig(BaseConfig): + """Configuration for router replay in MoE models. + + This configuration controls the routing behavior for Mixture of Experts (MoE) models, + allowing for deterministic training through route recording and replay. + + Args: + mode (str): Router replay mode. Options: 'disabled', 'R2', 'R3'. + - 'disabled': No router replay functionality + - 'R2': Use Router Replay routing strategy + - 'R3': Use Rollout Router Replay routing strategy + record_file (Optional[str]): File path to save recorded routing decisions. + Required when mode is 'record', 'R2', or 'R3'. + replay_file (Optional[str]): File path to load recorded routing decisions for replay. + Required when mode is 'replay'. + """ + + mode: str = "disabled" + record_file: Optional[str] = None + replay_file: Optional[str] = None + + def __post_init__(self): + """Validate router replay configuration.""" + valid_modes = ["disabled", "R2", "R3"] + if self.mode not in valid_modes: + raise ValueError(f"Invalid router_replay mode: {self.mode}. Must be one of {valid_modes}") + + +@dataclass +class PolicyLossConfig(BaseConfig): + """Configuration for policy loss computation. + + The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config. + + Args: + loss_mode (str): Loss function mode. Options: 'vanilla', 'clip-cov', 'kl-cov', 'gpg'. + clip_cov_ratio (float): Ratio of tokens to be clipped for clip-cov loss. + clip_cov_lb (float): Lower bound for clip-cov loss. + clip_cov_ub (float): Upper bound for clip-cov loss. + kl_cov_ratio (float): Ratio of tokens to be applied KL penalty for kl-cov loss. + ppo_kl_coef (float): KL divergence penalty coefficient. + rollout_correction (RolloutCorrectionConfig): Configuration for rollout correction. + """ + + loss_mode: str = "vanilla" + clip_cov_ratio: float = 0.0002 + clip_cov_lb: float = 1.0 + clip_cov_ub: float = 5.0 + kl_cov_ratio: float = 0.0002 + ppo_kl_coef: float = 0.1 + rollout_correction: RolloutCorrectionConfig = field(default_factory=RolloutCorrectionConfig) + + +@dataclass +class ActorConfig(BaseConfig): + """Configuration for actor model training. + + The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config. + + Args: + strategy (str): Training strategy. Must be specified. + ppo_mini_batch_size (int): Mini-batch size for PPO training. + ppo_micro_batch_size (Optional[int]): Micro-batch size for PPO training. + If None, uses ppo_micro_batch_size_per_gpu. + ppo_micro_batch_size_per_gpu (Optional[int]): Micro-batch size per GPU for PPO training. + use_dynamic_bsz (bool): Whether to use dynamic batch sizing. + ppo_max_token_len_per_gpu (int): Maximum token length per GPU for PPO training. + clip_ratio (float): PPO clipping ratio for policy loss. + clip_ratio_low (float): Lower bound for PPO clipping ratio. + clip_ratio_high (float): Upper bound for PPO clipping ratio. + policy_loss (PolicyLossConfig): Configuration for policy loss computation. + clip_ratio_c (float): Clipping ratio for critic loss. + loss_agg_mode (str): Loss aggregation mode. Options: 'token-mean', 'sample-mean'. + loss_scale_factor (Optional[int]): Scale factor for 'seq-mean-token-sum-norm' loss aggregation mode. + If None, uses response_length. Set to a constant to ensure consistent normalization. + entropy_coeff (float): Entropy coefficient for regularization. + tau_pos (float): Positive tau for SAPO smoothing (>= 1.0 keeps rewards stable). + tau_neg (float): Negative tau for SAPO smoothing (> tau_pos for asymmetry). + use_kl_loss (bool): Whether to use KL divergence loss. + use_torch_compile (bool): Whether to use torch.compile for optimization. + kl_loss_coef (float): KL divergence loss coefficient. + kl_loss_type (str): Type of KL loss to use. + ppo_epochs (int): Number of PPO epochs per training step. + shuffle (bool): Whether to shuffle data during training. + checkpoint (CheckpointConfig): Configuration for checkpointing. + optim (OptimizerConfig): Configuration for optimizer. + use_fused_kernels (bool): Whether to use custom fused kernels (e.g., FlashAttention, fused MLP). + data_loader_seed (int): Seed for data loader. If None, uses global seed. + router_replay (RouterReplayConfig): Configuration for router replay in MoE models. + """ + + _mutable_fields = BaseConfig._mutable_fields | { + "ppo_mini_batch_size", + "ppo_micro_batch_size", + "ppo_micro_batch_size_per_gpu", + "ppo_infer_micro_batch_size_per_gpu", + "engine", + "model_config", + } + + strategy: str = MISSING + ppo_mini_batch_size: int = 256 + ppo_micro_batch_size: Optional[int] = None # deprecate + ppo_micro_batch_size_per_gpu: Optional[int] = None + ppo_infer_micro_batch_size_per_gpu: Optional[int] = None + use_dynamic_bsz: bool = False + ppo_max_token_len_per_gpu: int = 16384 + ppo_infer_max_token_len_per_gpu: int = 16384 + clip_ratio: float = 0.2 + clip_ratio_low: float = 0.2 + clip_ratio_high: float = 0.2 + freeze_vision_tower: bool = False + policy_loss: PolicyLossConfig = field(default_factory=PolicyLossConfig) + clip_ratio_c: float = 3.0 + loss_agg_mode: str = "token-mean" + loss_scale_factor: Optional[int] = None + entropy_coeff: float = 0 + tau_pos: float = 1.0 + tau_neg: float = 1.05 + calculate_entropy: bool = False + calculate_sum_pi_squared: bool = False + use_kl_loss: bool = False + # Whether to enable PrefixGrouper-based shared-prefix forward + use_prefix_grouper: bool = False + use_torch_compile: bool = True + kl_loss_coef: float = 0.001 + kl_loss_type: str = "low_var_kl" + ppo_epochs: int = 1 + shuffle: bool = False + data_loader_seed: int = 1 + checkpoint: CheckpointConfig = field(default_factory=CheckpointConfig) + optim: OptimizerConfig = field(default_factory=OptimizerConfig) + use_fused_kernels: bool = False + profiler: ProfilerConfig = field(default_factory=ProfilerConfig) + engine: BaseConfig = field(default_factory=BaseConfig) + rollout_n: int = MISSING # must be override by sampling config + model_config: HFModelConfig = field(default_factory=BaseConfig) + router_replay: RouterReplayConfig = field(default_factory=RouterReplayConfig) + + # Store global batch info for loss aggregation: + # dp_size: data parallel size + # batch_num_tokens: number of valid tokens in global batch + # global_batch_size: global batch size + global_batch_info: dict = field(default_factory=dict) + qat: QATConfig = field(default_factory=QATConfig) + + def __post_init__(self): + """Validate actor configuration parameters.""" + assert self.strategy != MISSING + assert self.rollout_n != MISSING + if not self.use_dynamic_bsz: + if self.ppo_micro_batch_size is not None and self.ppo_micro_batch_size_per_gpu is not None: + raise ValueError( + "[actor] You have set both 'actor.ppo_micro_batch_size' AND 'actor.ppo_micro_batch_size_per_gpu'. " + "Please remove 'actor.ppo_micro_batch_size' because only '*_ppo_micro_batch_size_per_gpu' is " + "supported (the former is deprecated)." + ) + else: + assert not (self.ppo_micro_batch_size is None and self.ppo_micro_batch_size_per_gpu is None), ( + "[actor] Please set at least one of 'actor.ppo_micro_batch_size' or " + "'actor.ppo_micro_batch_size_per_gpu' if use_dynamic_bsz is not enabled." + ) + + valid_loss_agg_modes = [ + "token-mean", + "seq-mean-token-sum", + "seq-mean-token-mean", + "seq-mean-token-sum-norm", + ] + if self.loss_agg_mode not in valid_loss_agg_modes: + raise ValueError(f"Invalid loss_agg_mode: {self.loss_agg_mode}") + + def validate(self, n_gpus: int, train_batch_size: int, model_config: dict = None): + """Validate actor configuration with runtime parameters.""" + if not self.use_dynamic_bsz: + if train_batch_size < self.ppo_mini_batch_size: + raise ValueError( + f"train_batch_size ({train_batch_size}) must be >= " + f"actor.ppo_mini_batch_size ({self.ppo_mini_batch_size})" + ) + + sp_size = getattr(self, "ulysses_sequence_parallel_size", 1) + if self.ppo_micro_batch_size is not None: + if self.ppo_mini_batch_size % self.ppo_micro_batch_size != 0: + raise ValueError( + f"ppo_mini_batch_size ({self.ppo_mini_batch_size}) must be divisible by " + f"ppo_micro_batch_size ({self.ppo_micro_batch_size})" + ) + if self.ppo_micro_batch_size * sp_size < n_gpus: + raise ValueError( + f"ppo_micro_batch_size ({self.ppo_micro_batch_size}) * " + f"ulysses_sequence_parallel_size ({sp_size}) must be >= n_gpus ({n_gpus})" + ) + + @staticmethod + def _check_mutually_exclusive(mbs, mbs_per_gpu, name: str): + """Validate mutually exclusive micro batch size configuration options.""" + param = "ppo_micro_batch_size" + param_per_gpu = f"{param}_per_gpu" + + if mbs is None and mbs_per_gpu is None: + raise ValueError(f"[{name}] Please set at least one of '{name}.{param}' or '{name}.{param_per_gpu}'.") + + if mbs is not None and mbs_per_gpu is not None: + raise ValueError( + f"[{name}] You have set both '{name}.{param}' AND '{name}.{param_per_gpu}'. Please remove " + f"'{name}.{param}' because only '*_{param_per_gpu}' is supported (the former is deprecated)." + ) + + +@dataclass +class McoreActorConfig(ActorConfig): + """Configuration for Megatron actor models. + + The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config. + + Args: + strategy (str): Training strategy set to 'megatron' for Megatron parallelism. + load_weight (bool): Whether to load model weights from checkpoint. + megatron (dict[str, Any]): Configuration for Megatron parallelism settings. + profile (dict[str, Any]): Configuration for profiling settings. + """ + + strategy: str = "megatron" + load_weight: bool = True + megatron: McoreEngineConfig = field(default_factory=McoreEngineConfig) + profile: dict[str, Any] = field(default_factory=dict) + use_rollout_log_probs: bool = False + + def __post_init__(self): + """Validate FSDP actor configuration parameters.""" + super().__post_init__() + self.engine = self.megatron + + +@dataclass +class FSDPActorConfig(ActorConfig): + """Configuration for FSDP actor models. + + The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config. + + Args: + strategy (str): Training strategy set to 'fsdp' for Fully Sharded Data Parallel. + grad_clip (float): Gradient clipping threshold. + ulysses_sequence_parallel_size (int): [DEPRECATED] Ulysses sequence parallel size for long sequences. + entropy_from_logits_with_chunking (bool): Whether to compute entropy from logits + with chunking for memory efficiency. + entropy_checkpointing (bool): Whether to use gradient checkpointing for entropy computation. + fsdp_config (dict[str, Any]): Configuration for FSDP settings. + use_remove_padding (bool): Whether to remove padding tokens in inputs during training + """ + + strategy: str = "fsdp" + grad_clip: float = 1.0 + ulysses_sequence_parallel_size: int = 1 + entropy_from_logits_with_chunking: bool = False + entropy_checkpointing: bool = False + fsdp_config: FSDPEngineConfig = field(default_factory=FSDPEngineConfig) + use_remove_padding: bool = False + use_rollout_log_probs: bool = False + + def __post_init__(self): + """Validate FSDP actor configuration parameters.""" + super().__post_init__() + self.engine = self.fsdp_config + # Sync strategy to engine config so engine_workers can pick the right FSDP version. + # EngineConfig.strategy defaults to None, so without this, engine_workers.py always + # falls back to FSDP1 even when actor.strategy="fsdp2". + object.__setattr__(self.engine, "strategy", self.strategy) + + # backward compatibility + if self.ulysses_sequence_parallel_size > 1: + self.fsdp_config.ulysses_sequence_parallel_size = self.ulysses_sequence_parallel_size + + def validate(self, n_gpus: int, train_batch_size: int, model_config: dict = None): + """Validate FSDP actor configuration with runtime parameters.""" + super().validate(n_gpus, train_batch_size, model_config) + + if self.strategy in {"fsdp", "fsdp2"} and self.ulysses_sequence_parallel_size > 1: + if model_config and not model_config.get("use_remove_padding", False): + raise ValueError( + "When using sequence parallelism for actor/ref policy, you must enable `use_remove_padding`." + ) + + +@dataclass +class VeOmniActorConfig(ActorConfig): + """Configuration for VeOmni actor models. + + The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config. + + Args: + strategy (str): Training strategy set to 'veomni' for VeOmni parallelism. + veomni (dict[str, Any]): Configuration for VeOmni settings. + use_remove_padding (bool): Whether to remove padding tokens in inputs during training + """ + + strategy: str = "veomni" + veomni: VeOmniEngineConfig = field(default_factory=VeOmniEngineConfig) + use_remove_padding: bool = False + use_rollout_log_probs: bool = False + + def __post_init__(self): + """Validate VeOmni actor configuration parameters.""" + super().__post_init__() + self.engine = self.veomni + + +@dataclass +class TorchTitanActorConfig(ActorConfig): + """Configuration for TorchTitan actor models. + + The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config. + + Args: + strategy (str): Training strategy set to 'torchtitan' for TorchTitan parallelism. + torchtitan (TorchtitanEngineConfig): Configuration for TorchTitan engine settings. + use_remove_padding (bool): Whether to remove padding tokens in inputs during training + use_rollout_log_probs (bool): Whether to use log probabilities from rollout engine + """ + + strategy: str = "torchtitan" + torchtitan: TorchtitanEngineConfig = field(default_factory=TorchtitanEngineConfig) + use_remove_padding: bool = False + use_rollout_log_probs: bool = False + + def __post_init__(self): + """Validate TorchTitan actor configuration parameters.""" + super().__post_init__() + self.engine = self.torchtitan + + +@dataclass +class MindSpeedActorConfig(ActorConfig): + """Configuration for mindspeed actor models. + + The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config. + + Args: + strategy (str): Training strategy set to 'mindspeed' for mindspeed parallelism. + load_weight (bool): Whether to load model weights from checkpoint. + mindspeed (dict[str, Any]): Configuration for mindspeed parallelism settings. + profile (dict[str, Any]): Configuration for profiling settings. + use_rollout_log_probs (bool): Whether to use log probabilities from rollout engine. + """ + + strategy: str = "mindspeed" + load_weight: bool = True + mindspeed: MindSpeedEngineConfig = field(default_factory=MindSpeedEngineConfig) + profile: dict[str, Any] = field(default_factory=dict) + use_rollout_log_probs: bool = False + + def __post_init__(self): + """Validate MindSpeed actor configuration parameters.""" + super().__post_init__() + self.engine = self.mindspeed diff --git a/verl/verl/workers/config/critic.py b/verl/verl/workers/config/critic.py new file mode 100644 index 0000000000000000000000000000000000000000..3531ed63c92707e1d0196e288be454e8016f6e56 --- /dev/null +++ b/verl/verl/workers/config/critic.py @@ -0,0 +1,300 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass, field +from typing import Optional + +from omegaconf import MISSING + +from verl.base_config import BaseConfig +from verl.trainer.config import BaseModelConfig, CheckpointConfig +from verl.utils.profiler import ProfilerConfig + +from .engine import FSDPEngineConfig, McoreEngineConfig, MindSpeedEngineConfig, TorchtitanEngineConfig +from .model import HFModelConfig +from .optimizer import OptimizerConfig + +__all__ = [ + "CriticConfig", + "FSDPCriticConfig", + "McoreCriticConfig", + "TorchTitanCriticConfig", + "FSDPCriticModelCfg", + "MindSpeedCriticConfig", +] + + +@dataclass +class CriticConfig(BaseConfig): + """Configuration for critic model training. + + The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config. + + Args: + strategy (str): Strategy used for critic model training (fsdp, fsdp2, megatron). + ppo_micro_batch_size_per_gpu (int): Local per-GPU micro batch size. + rollout_n (int): Number of rollouts per update (mirrors actor rollout_n). + optim (Dict[str, Any]): Optimizer configuration including lr, weight_decay, etc. + model (Dict[str, Any]): Model configuration including path, tokenizer_path, etc. + ppo_mini_batch_size (int): PPO mini-batch size per update. + ppo_micro_batch_size (Optional[int]): Global micro batch size (deprecated). + use_dynamic_bsz (bool): Whether to automatically adjust batch size at runtime. + ppo_max_token_len_per_gpu (int): Max tokens per GPU in one PPO batch. + forward_max_token_len_per_gpu (int): Max token length per GPU in forward pass. + ppo_epochs (int): Number of PPO epochs per batch. + shuffle (bool): Shuffle training data across PPO epochs. + cliprange_value (float): PPO value function clipping range. + loss_agg_mode (str): Loss aggregation mode. + checkpoint (Dict[str, Any]): Checkpoint configuration. + profiler (Dict[str, Any]): Profiler configuration. + enable (Optional[bool]): Whether to enable the critic. + """ + + _mutable_fields = BaseConfig._mutable_fields | { + "ppo_micro_batch_size_per_gpu", + "ppo_mini_batch_size", + "ppo_micro_batch_size", + "engine", + "model_config", + } + + strategy: str = MISSING + ppo_micro_batch_size_per_gpu: Optional[int] = None + enable: Optional[bool] = None + rollout_n: int = 1 + ppo_mini_batch_size: int = 1 + use_dynamic_bsz: bool = False + ppo_max_token_len_per_gpu: int = 32768 + # deprecate this + forward_max_token_len_per_gpu: int = 32768 + ppo_infer_micro_batch_size_per_gpu: Optional[int] = None + ppo_infer_max_token_len_per_gpu: int = 32768 + ppo_epochs: int = 1 + data_loader_seed: int = 1 + shuffle: bool = True + cliprange_value: float = 0.5 + loss_agg_mode: str = "token-mean" + ppo_micro_batch_size: Optional[int] = None + engine: BaseConfig = field(default_factory=BaseConfig) + optim: OptimizerConfig = field(default_factory=OptimizerConfig) + model: HFModelConfig = None + checkpoint: CheckpointConfig = field(default_factory=CheckpointConfig) + profiler: ProfilerConfig = field(default_factory=ProfilerConfig) + + def __post_init__(self): + """Validate critic configuration parameters.""" + assert self.strategy != MISSING + + if not self.use_dynamic_bsz: + self._check_mutually_exclusive(self.ppo_micro_batch_size, self.ppo_micro_batch_size_per_gpu, "critic") + + if self.ppo_micro_batch_size is not None: + if self.ppo_mini_batch_size % self.ppo_micro_batch_size != 0: + raise ValueError( + f"[critic] ppo_mini_batch_size ({self.ppo_mini_batch_size}) must be divisible by " + f"ppo_micro_batch_size ({self.ppo_micro_batch_size})" + ) + + def validate(self, n_gpus: int, train_batch_size: int): + """Validate critic configuration with runtime parameters. + + Args: + n_gpus: Total number of GPUs available + train_batch_size: Training batch size from data config + """ + if not self.use_dynamic_bsz: + if train_batch_size < self.ppo_mini_batch_size: + raise ValueError( + f"train_batch_size ({train_batch_size}) must be >= " + f"critic.ppo_mini_batch_size ({self.ppo_mini_batch_size})" + ) + + @staticmethod + def _check_mutually_exclusive(mbs, mbs_per_gpu, name: str): + """Validate mutually exclusive micro batch size configuration options. + + Ensures that users don't set both deprecated micro_batch_size and + the new micro_batch_size_per_gpu parameters simultaneously. + + Args: + mbs: Deprecated micro batch size parameter value. + mbs_per_gpu: New micro batch size per GPU parameter value. + name (str): Configuration section name for error messages. + + Raises: + ValueError: If both parameters are set or neither is set. + """ + param = "micro_batch_size" + param_per_gpu = f"{param}_per_gpu" + + if mbs is None and mbs_per_gpu is None: + raise ValueError(f"[{name}] Please set at least one of '{name}.{param}' or '{name}.{param_per_gpu}'.") + + if mbs is not None and mbs_per_gpu is not None: + raise ValueError( + f"[{name}] You have set both '{name}.{param}' AND '{name}.{param_per_gpu}'. Please remove " + f"'{name}.{param}' because only '*_{param_per_gpu}' is supported (the former is deprecated)." + ) + + +@dataclass +class McoreCriticConfig(CriticConfig): + """Configuration for Megatron-based critic model training. + + The inheritance from CriticConfig provides all base critic configuration plus Megatron-specific settings. + + Args: + nccl_timeout (int): NCCL timeout in seconds for distributed operations. + megatron (Dict[str, Any]): Megatron-specific parallelism settings. + load_weight (bool): Whether to load initial weights. + """ + + strategy: str = "megatron" + nccl_timeout: int = 600 + megatron: McoreEngineConfig = field(default_factory=McoreEngineConfig) + load_weight: bool = True + + def validate(self, n_gpus: int, train_batch_size: int): + """Validate Megatron critic configuration with runtime parameters.""" + super().validate(n_gpus, train_batch_size) + + def __post_init__(self): + """Validate Megatron critic configuration parameters.""" + super().__post_init__() + self.engine = self.megatron + + +@dataclass +class FSDPCriticConfig(CriticConfig): + """Configuration for FSDP-based critic model training. + + The inheritance from CriticConfig provides all base critic configuration plus FSDP-specific settings. + + Args: + forward_micro_batch_size (int): Forward-only batch size during inference (global). + forward_micro_batch_size_per_gpu (int): Forward-only batch size during inference (per GPU). + ulysses_sequence_parallel_size (int): [DEPRECATED] Ulysses sequence parallel size for long sequences. + grad_clip (float): Gradient clipping for critic updates. + """ + + _mutable_fields = CriticConfig._mutable_fields | { + "forward_micro_batch_size", + "forward_micro_batch_size_per_gpu", + } + + strategy: str = "fsdp" + fsdp: FSDPEngineConfig = field(default_factory=FSDPEngineConfig) + forward_micro_batch_size: int = 1 + forward_micro_batch_size_per_gpu: int = 1 + ulysses_sequence_parallel_size: int = 1 + grad_clip: float = 1.0 + + def __post_init__(self): + """Validate FSDP critic configuration parameters.""" + super().__post_init__() + self.engine = self.fsdp + # Sync strategy to engine config so engine_workers can pick the right FSDP version. + # EngineConfig.strategy defaults to None, so without this, engine_workers.py always + # falls back to FSDP1 even when critic.strategy="fsdp2". + object.__setattr__(self.engine, "strategy", self.strategy) + + if self.strategy in {"fsdp", "fsdp2"}: + if self.ulysses_sequence_parallel_size > 1: + if not self.model.get("use_remove_padding", False): + raise ValueError( + "When using sequence parallelism for critic, you must enable `use_remove_padding`." + ) + + def validate(self, n_gpus: int, train_batch_size: int): + """Validate FSDP critic configuration with runtime parameters.""" + super().validate(n_gpus, train_batch_size) + + if not self.use_dynamic_bsz: + sp_size = self.ulysses_sequence_parallel_size + if self.ppo_micro_batch_size is not None: + if self.ppo_micro_batch_size * sp_size < n_gpus: + raise ValueError( + f"critic.ppo_micro_batch_size ({self.ppo_micro_batch_size}) * " + f"ulysses_sequence_parallel_size ({sp_size}) must be >= n_gpus ({n_gpus})" + ) + + +@dataclass +class TorchTitanCriticConfig(CriticConfig): + """Configuration for TorchTitan-based critic model training. + + The inheritance from CriticConfig provides all base critic configuration plus TorchTitan-specific settings. + + Args: + strategy (str): Training strategy set to 'torchtitan' for TorchTitan parallelism. + torchtitan (TorchtitanEngineConfig): Configuration for TorchTitan engine settings. + """ + + strategy: str = "torchtitan" + torchtitan: TorchtitanEngineConfig = field(default_factory=TorchtitanEngineConfig) + + def __post_init__(self): + """Validate TorchTitan critic configuration parameters.""" + super().__post_init__() + self.engine = self.torchtitan + + +@dataclass +class FSDPCriticModelCfg(BaseModelConfig): + """FSDP-enabled critic model configuration. + Inherits base critic settings and adds distributed-memory and LoRA options. + + Args: + use_shm (bool): Whether to use shared memory for loading the model. + enable_activation_offload (bool): Offload activations to CPU to reduce GPU memory usage. + use_remove_padding (bool): Use remove-padding optimization (saves compute). + enable_gradient_checkpointing (bool): Enable gradient checkpointing for memory efficiency. + fsdp_config (FSDPEngineConfig): FSDP-specific configuration block. + lora_rank (int): Set to positive value to enable LoRA (e.g., 32). + lora_alpha (int): LoRA scaling factor. + target_modules (Union[str, List[str]]): LoRA target modules: "all-linear" or list of layer names. + """ + + use_shm: bool = False + enable_activation_offload: bool = False + use_remove_padding: bool = False + enable_gradient_checkpointing: bool = True + fsdp_config: FSDPEngineConfig = field(default_factory=FSDPEngineConfig) + lora_rank: int = 0 + lora_alpha: int = 16 + target_modules: str | list[str] = "all-linear" + # TiledMLP configuration for memory-efficient MLP computation + tiled_mlp: dict = field(default_factory=lambda: {"enabled": False, "num_shards": 4}) + + +@dataclass +class MindSpeedCriticConfig(CriticConfig): + """Configuration for mindspeed-based critic model training. + + The inheritance from CriticConfig provides all base critic configuration plus mindspeed-specific settings. + + Args: + nccl_timeout (int): NCCL timeout in seconds for distributed operations. + mindspeed (Dict[str, Any]): mindspeed-specific parallelism settings. + load_weight (bool): Whether to load initial weights. + """ + + strategy: str = "mindspeed" + nccl_timeout: int = 600 + mindspeed: MindSpeedEngineConfig = field(default_factory=MindSpeedEngineConfig) + load_weight: bool = True + + def validate(self, n_gpus: int, train_batch_size: int): + """Validate mindspeed critic configuration with runtime parameters.""" + super().validate(n_gpus, train_batch_size) diff --git a/verl/verl/workers/config/disaggregation.py b/verl/verl/workers/config/disaggregation.py new file mode 100644 index 0000000000000000000000000000000000000000..731e13637732de399e87e84406f3e34317d0cfcd --- /dev/null +++ b/verl/verl/workers/config/disaggregation.py @@ -0,0 +1,54 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from dataclasses import dataclass +from typing import Optional + +from verl.base_config import BaseConfig + +__all__ = ["DisaggregationConfig"] + +_ALLOWED_BACKENDS = ("nixl", "mooncake", "ascend", "mori", "fake") + + +@dataclass +class DisaggregationConfig(BaseConfig): + """Prefill-Decode disaggregation knobs (SGLang only).""" + + enabled: bool = False + prefill_replicas: int = 1 + decode_replicas: int = 1 + decode_tensor_model_parallel_size: Optional[int] = None + transfer_backend: str = "nixl" + bootstrap_port: Optional[int] = None + ib_device: Optional[str] = None + + def __post_init__(self) -> None: + if not self.enabled: + return + if self.transfer_backend not in _ALLOWED_BACKENDS: + raise ValueError(f"disaggregation.transfer_backend={self.transfer_backend!r} not in {_ALLOWED_BACKENDS}") + if self.prefill_replicas < 1 or self.decode_replicas < 1: + raise ValueError( + f"disaggregation requires >=1 prefill and >=1 decode replica " + f"(got prefill_replicas={self.prefill_replicas}, decode_replicas={self.decode_replicas})" + ) + if self.bootstrap_port is not None and not (0 < self.bootstrap_port < 65536): + raise ValueError(f"bootstrap_port out of range: {self.bootstrap_port}") + + def effective_decode_tp(self, prefill_tp: int) -> int: + """Resolve decode TP (defaults to ``prefill_tp``). Test-only helper; runtime paths + must inline this because OmegaConf/Ray serialization drops dataclass methods.""" + if self.decode_tensor_model_parallel_size is not None: + return self.decode_tensor_model_parallel_size + return prefill_tp diff --git a/verl/verl/workers/config/distillation.py b/verl/verl/workers/config/distillation.py new file mode 100644 index 0000000000000000000000000000000000000000..cc075d41f75e65f3f605d71e73434f8610c9538b --- /dev/null +++ b/verl/verl/workers/config/distillation.py @@ -0,0 +1,308 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +import os +from dataclasses import dataclass, field +from typing import Optional + +from verl.base_config import BaseConfig +from verl.utils.config import omega_conf_to_dataclass + +from .rollout import RolloutConfig + +__all__ = ["DistillationLossConfig", "DistillationTeacherModelConfig", "DistillationConfig"] + +logger = logging.getLogger(__name__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + + +@dataclass +class DistillationLossConfig(BaseConfig): + """Configuration for distillation loss settings. + + loss_mode (str): + Distillation loss function to use. + topk (int, optional): + Number of top tokens to consider for top-k distillation losses. + use_task_rewards (bool): + Whether to include task rewards alongside distillation loss. + distillation_loss_coef (float): + Coefficient for distillation loss when combined with task rewards. + loss_max_clamp (float, optional): + Maximum value to clamp distillation loss. If None, no clamping is applied. + log_prob_min_clamp (float, optional): + Minimum value to clamp log probabilities for stability, e.g., log q - log p where p or q are + very close to zero. If None, no clamping is applied. + use_policy_gradient (bool): + Whether to incorporate distillation loss as a reward, as done + by https://thinkingmachines.ai/blog/on-policy-distillation/. Recommended to use loss_mode=k1. + Otherwise, distillation loss is directly backpropagated as a supervised loss, + as in https://arxiv.org/abs/2306.13649. Recommended to use loss_mode=k3 or forward_kl_topk. + policy_loss_mode (str): + Name of the policy loss to use when use_policy_gradient is true. + clip_ratio (float): + PPO clipping ratio for policy loss. + clip_ratio_low (float): + Lower bound for PPO clipping ratio. + clip_ratio_high (float): + Upper bound for PPO clipping ratio. + loss_settings (DistillationLossSettings, optional): + Runtime-populated settings based on loss_mode. Not set by user. + """ + + loss_mode: str = "k3" + topk: Optional[int] = 128 + use_task_rewards: bool = True + distillation_loss_coef: float = 1.0 + loss_max_clamp: Optional[float] = 10.0 + log_prob_min_clamp: Optional[float] = -10.0 + + use_policy_gradient: bool = True + policy_loss_mode: str = "vanilla" + clip_ratio: float = 0.2 + clip_ratio_low: float = 0.2 + clip_ratio_high: float = 0.2 + + # Store global batch info for loss aggregation: + # dp_size: data parallel size + # batch_num_tokens: number of valid tokens in global batch + # global_batch_size: global batch size + global_batch_info: dict = field(default_factory=dict) + + # Store distillation loss settings for computing the specified loss_mode + # Not set by user, populated at runtime + loss_settings: Optional[dict] = None + + def __post_init__(self): + self._mutable_fields.add("loss_settings") + from verl.trainer.distillation.losses import DistillationLossSettings, get_distillation_loss_settings + + self.loss_settings: DistillationLossSettings = get_distillation_loss_settings(self.loss_mode) + + if self.policy_loss_mode != "vanilla": + raise NotImplementedError( + f"Only vanilla policy loss is currently supported when use_policy_gradient is True, " + f"but got {self.policy_loss_mode}." + ) + + if self.use_policy_gradient and self.loss_mode == "forward_kl_topk": + print( + "WARNING: forward_kl_topk is most effective as a supervised distillation loss " + "(use_policy_gradient=False). With policy gradient, the update uses only the sampled" + " token's logprob ∇logπ(a), so the top-k distributional signal (how non-sampled logits " + "should move) is largely unused." + ) + + if not self.use_policy_gradient and self.loss_mode == "k1": + raise ValueError( + "Directly backpropagating k1 loss is incorrect since gradient of k1 loss" + " wrt model weights does not depend on teacher log probabilities." + ) + + +@dataclass +class DistillationTeacherModelConfig(BaseConfig): + """Configuration for on-policy distillation teacher. + + key (str, optional): + Identifier to route examples to the teacher model in multi-teacher setting. + model_path (str, optional): + Model path for the teacher model. Can be a local path or a Hugging Face model + inference (RolloutConfig): + Rollout configuration for the teacher model inference during distillation. + num_replicas (int): + Number of inference replicas of this teacher to launch. Each replica occupies + `per_replica_world_size` GPUs (= inference.data_parallel_size * + inference.tensor_model_parallel_size * inference.pipeline_model_parallel_size), + so the teacher's total GPU footprint is + `num_replicas * per_replica_world_size`. + """ + + _mutable_fields = BaseConfig._mutable_fields | {"num_replicas", "key"} + + key: Optional[str] = None + model_path: Optional[str] = None + inference: RolloutConfig = field(default_factory=RolloutConfig) + num_replicas: Optional[int] = 0 + + @property + def per_replica_world_size(self) -> int: + return ( + self.inference.tensor_model_parallel_size + * self.inference.data_parallel_size + * self.inference.pipeline_model_parallel_size + ) + + @property + def world_size(self) -> int: + return self.num_replicas * self.per_replica_world_size + + def check_configured(self): + if self.model_path is None: + raise ValueError("model_path must be specified for distillation teacher model config.") + if self.key is None: + raise ValueError("key must be specified for distillation teacher model config.") + if self.num_replicas is None: + raise ValueError("num_replicas must be specified for distillation teacher model config.") + + def validate_and_prepare_for_distillation(self, use_topk: bool, topk: Optional[int]) -> None: + # Prompt + Response from student are fed into teacher as context + max_model_len = self.inference.max_model_len + student_prompt_length = self.inference.prompt_length + student_response_length = self.inference.response_length + required_context_len = student_prompt_length + student_response_length + 1 + if max_model_len is not None and required_context_len > max_model_len: + raise ValueError( + "Distillation teacher inference requires room for the student prompt, the full student " + f"response, and one generated token, but got {student_prompt_length=}, " + f"{student_response_length=}, {required_context_len=}, {max_model_len=}." + ) + self.inference.prompt_length = self.inference.prompt_length + self.inference.response_length + self.inference.response_length = 1 + self._validate_topk_logprobs(use_topk=use_topk, topk=topk) + + def _validate_topk_logprobs(self, use_topk: bool, topk: Optional[int]) -> None: + if not use_topk: + return + if topk is None: + raise ValueError("topk must be specified when use_topk is True.") + + engine_name = self.inference.name + engine_kwargs = self.inference.engine_kwargs + match engine_name: + case "vllm": + vllm_engine_kwargs = dict(engine_kwargs.get("vllm", {})) + max_logprobs = vllm_engine_kwargs.get("max_logprobs") + if max_logprobs is None: + vllm_engine_kwargs["max_logprobs"] = topk + max_logprobs = topk + if max_logprobs < topk: + raise ValueError( + f"VLLM max_logprobs ({max_logprobs}) must be >= distillation_loss topk " + f"({topk}) to enable distillation loss computation." + ) + engine_kwargs["vllm"] = vllm_engine_kwargs + case "sglang": + # SGLang's top_logprobs_num is a per-request parameter, so there is no + # engine-boot cap to align (unlike vLLM's max_logprobs). The async + # server translates sampling_params["prompt_logprobs"] into + # return_logprob + logprob_start_len=0 + top_logprobs_num at call time. + pass + case _: + raise NotImplementedError( + f"DistillationTeacherModelConfig does not support inference engine {engine_name}" + ) + + +@dataclass +class DistillationConfig(BaseConfig): + """Configuration for on-policy distillation. + + enabled (bool): + Whether on-policy distillation is enabled. + n_gpus_per_node (int): + Number of GPUs per node in the teacher resource pool. + nnodes (int): + Number of nodes in the teacher resource pool. + teacher_models (dict[str, TeacherModelConfig]): + Configurations for teacher models used for multi-teacher distillation. + teacher_key (str): + Key to route examples to the appropriate teacher model in multi-teacher setups. Should correspond to a field in + the data proto, e.g., data_source. + distillation_loss (DistillationLossConfig): + Configuration for distillation loss settings. + + NOTE: The `teacher_model` entry is in the `teacher_models` dict by default. + Since it is popped when other teacher entries are added, using `teacher_model` as + one of several keys silently drops it. For example, the following CLI overrides result + in ONLY `teacher_model2` being used: + + ```bash + distillation.teacher_models.teacher_model.key=openai/gsm8k + distillation.teacher_models.teacher_model.model_path=Qwen/Qwen3-4B + +distillation.teacher_models.teacher_model2.key=hiyouga/geometry3k + +distillation.teacher_models.teacher_model2.model_path=Qwen/Qwen3-VL-4B-Instruct + ``` + Instead, give the first teacher a different name: + + ```bash + +distillation.teacher_models.teacher_model1.key=openai/gsm8k + +distillation.teacher_models.teacher_model1.model_path=Qwen/Qwen3-4B + +distillation.teacher_models.teacher_model2.key=hiyouga/geometry3k + +distillation.teacher_models.teacher_model2.model_path=Qwen/Qwen3-VL-4B-Instruct + ``` + """ + + _mutable_fields = BaseConfig._mutable_fields | {"teacher_models", "n_gpus_per_node", "nnodes"} + + enabled: bool = False + n_gpus_per_node: int = 0 + nnodes: int = 0 + teacher_models: dict[str, DistillationTeacherModelConfig] = field(default_factory=dict) + teacher_key: str = "data_source" + distillation_loss: DistillationLossConfig = field(default_factory=DistillationLossConfig) + + def __post_init__(self): + if not self.enabled: + return + + self.teacher_models = self._resolve_teacher_models() + teacher_world_size_sum = 0 + for teacher_model in self.teacher_models.values(): + teacher_model.validate_and_prepare_for_distillation( + use_topk=self.distillation_loss.loss_settings.use_topk, + topk=self.distillation_loss.topk, + ) + teacher_world_size_sum += teacher_model.world_size + total_pool_size = self.n_gpus_per_node * self.nnodes + if teacher_world_size_sum != total_pool_size: + raise ValueError( + f"Sum of teacher (num_replicas * per_replica_world_size) ({teacher_world_size_sum}) must match " + f"the distillation resource pool size " + f"({self.n_gpus_per_node=} * {self.nnodes=} = {total_pool_size})." + ) + + def _resolve_teacher_models(self) -> dict[str, DistillationTeacherModelConfig]: + assert "teacher_model" in self.teacher_models + if len(self.teacher_models) == 1: + # Single teacher occupies the entire teacher resource pool. + teacher_model = self.teacher_models["teacher_model"] + inference = teacher_model.inference + per_replica = ( + inference.tensor_model_parallel_size + * inference.data_parallel_size + * inference.pipeline_model_parallel_size + ) + pool_size = self.n_gpus_per_node * self.nnodes + if pool_size % per_replica != 0: + raise ValueError( + f"Single teacher's per_replica_world_size ({per_replica}) must divide the distillation " + f"resource pool size ({self.n_gpus_per_node=} * {self.nnodes=} = {pool_size})." + ) + teacher_model.num_replicas = pool_size // per_replica + teacher_model.key = "default" + else: + # Multiple teachers: remove default single teacher config + self.teacher_models.pop("teacher_model") + + # Teacher models dict is keyed by teacher_key instead of YAML entry name + teacher_models = {} + for teacher_config in self.teacher_models.values(): + teacher_config = omega_conf_to_dataclass(teacher_config, dataclass_type=DistillationTeacherModelConfig) + teacher_config.check_configured() + if teacher_config.key in teacher_models: + raise ValueError(f"Duplicate teacher key {teacher_config.key} found in teacher models.") + teacher_models[teacher_config.key] = teacher_config + return teacher_models diff --git a/verl/verl/workers/config/engine.py b/verl/verl/workers/config/engine.py new file mode 100644 index 0000000000000000000000000000000000000000..27f8922cdfad9672564d2438674afbaedf4bf2a5 --- /dev/null +++ b/verl/verl/workers/config/engine.py @@ -0,0 +1,605 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +import os +import warnings +from dataclasses import dataclass, field +from typing import Any, Callable, Literal, Optional + +from verl.base_config import BaseConfig +from verl.trainer.config import CheckpointConfig + +from ...utils.profiler import ProfilerConfig +from .model import HFModelConfig +from .optimizer import OptimizerConfig + +__all__ = [ + "FSDPEngineConfig", + "McoreEngineConfig", + "TrainingWorkerConfig", + "TorchtitanEngineConfig", + "VeOmniEngineConfig", + "AutomodelEngineConfig", + "EngineConfig", + "EngineRouterReplayConfig", + "QATEngineConfig", + "MindSpeedEngineConfig", +] + + +logger = logging.getLogger(__name__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "INFO")) + + +# TODO: rename to RouterReplayConfig after removing the legacy implementation +@dataclass +class EngineRouterReplayConfig(BaseConfig): + """Configuration for router replay in MoE models. + + This configuration controls the routing behavior for Mixture of Experts (MoE) models, + allowing for deterministic training through route recording and replay. + + Args: + mode (str): Router replay mode. Options: 'disabled', 'R2', 'R3'. + - 'disabled': No router replay functionality + - 'R2': Use Router Replay routing strategy + - 'R3': Use Rollout Router Replay routing strategy + record_file (Optional[str]): File path to save recorded routing decisions. + Required when mode is 'record', 'R2', or 'R3'. + replay_file (Optional[str]): File path to load recorded routing decisions for replay. + Required when mode is 'replay'. + """ + + mode: str = "disabled" + record_file: Optional[str] = None + replay_file: Optional[str] = None + + def __post_init__(self): + """Validate router replay configuration.""" + valid_modes = ["disabled", "R2", "R3"] + if self.mode not in valid_modes: + raise ValueError(f"Invalid router_replay mode: {self.mode}. Must be one of {valid_modes}") + + +@dataclass +class EngineConfig(BaseConfig): + _mutable_fields = BaseConfig._mutable_fields | { + "use_dynamic_bsz", + "max_token_len_per_gpu", + "micro_batch_size_per_gpu", + "infer_max_token_len_per_gpu", + "infer_micro_batch_size_per_gpu", + "use_fused_kernels", + "use_remove_padding", + "forward_only", + "param_offload", + } + # whether to offload param + param_offload: bool = False + # whether to offload optimizer + optimizer_offload: bool = False + # whether to offload grad + grad_offload: bool = False + # whether the engine is forward only (e.g., ref policy) + forward_only: bool = False + # the strategy (backend) + strategy: str = None + # model dtype + dtype: str = "bfloat16" # ["bfloat16", "float16"] + # whether to use dynamic bsz + use_dynamic_bsz: bool = True + # for training + max_token_len_per_gpu: int = None + micro_batch_size_per_gpu: int = None + # for inference + infer_max_token_len_per_gpu: int = None + infer_micro_batch_size_per_gpu: int = None + # whether use fuse lm head kernel + use_fused_kernels: bool = False + # TODO (this may conflict with the one in model config) + use_remove_padding: bool = True + + seed: int = 42 + + full_determinism: bool = False + router_replay: EngineRouterReplayConfig = field(default_factory=EngineRouterReplayConfig) + + def __post_init__(self): + pass + # TODO: turn on this check after we reorg config + # if self.use_dynamic_bsz: + # assert self.max_token_len_per_gpu is not None + # else: + # assert self.micro_batch_size_per_gpu is not None + + +@dataclass +class QATEngineConfig(BaseConfig): + """Configuration for QAT (Quantization-Aware Training) within an engine. + + Args: + enable (bool): Whether to enable QAT, default False + mode (str): Quantization mode, "w4a16" or "w4a4", default "w4a16" + group_size (int): Group size for blockwise quantization, default 16 + ignore_patterns (list[str]): Module name patterns to exclude from quantization + activation_observer (str): Observer strategy for activation global_scale (W4A4 only) + quantization_config_path (Optional[str]): Path to quantization config JSON for vLLM + """ + + enable: bool = False + mode: str = "w4a16" + group_size: int = 16 + ignore_patterns: list[str] = field(default_factory=lambda: ["lm_head", "embed_tokens", "re:.*mlp.gate$"]) + activation_observer: str = "static_minmax" + quantization_config_path: Optional[str] = None + + +@dataclass +class McoreEngineConfig(EngineConfig): + """Configuration for Megatron parallelism. + + The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config. + + Args: + param_offload (bool): Whether to offload parameters to CPU. + grad_offload (bool): Whether to offload gradients to CPU. + optimizer_offload (bool): Whether to offload optimizer states to CPU. + tensor_model_parallel_size (int): Tensor model parallel size. + expert_model_parallel_size (int): Expert model parallel size for MoE models. + expert_tensor_parallel_size (Optional[int]): Expert tensor parallel size for MoE models. + pipeline_model_parallel_size (int): Pipeline model parallel size. + virtual_pipeline_model_parallel_size (Optional[int]): Virtual pipeline model parallel size + for interleaved scheduling. + context_parallel_size (int): Context parallel size for long sequences. + dynamic_context_parallel (bool): Whether to enable hybrid context parallelism. + max_seqlen_per_dp_cp_rank (Optional[int]): Maximum sequence length per DPxCP rank. + sequence_parallel (bool): Whether to enable sequence parallelism. + use_distributed_optimizer (bool): Whether to use distributed optimizer. + use_dist_checkpointing (bool): Whether to use distributed checkpointing. + dist_checkpointing_path (Optional[str]): Path for distributed checkpointing. + dist_ckpt_optim_fully_reshardable (bool): Use fully reshardable optimizer checkpoints. + distrib_optim_fully_reshardable_mem_efficient (bool): Use memory-efficient fully reshardable format. + seed (int): Random seed for reproducibility. + override_ddp_config (dict[str, Any]): Override configuration for DDP. + override_transformer_config (dict[str, Any]): Override configuration for transformer. + use_mbridge (bool): Whether to use MBridge for communication. + use_megatron_fsdp (bool): Whether to use Megatron-FSDP (Zero-3 sharding). + dtype (str): Mixed precision training param dtype, default "bfloat16" + """ + + # sequence_parallel is not listed as a frozen field for auto-correction purpose + _mutable_fields = EngineConfig._mutable_fields | {"sequence_parallel"} + # mcore parallelism + tensor_model_parallel_size: int = 1 + expert_model_parallel_size: int = 1 + expert_tensor_parallel_size: Optional[int] = None + pipeline_model_parallel_size: int = 1 + virtual_pipeline_model_parallel_size: Optional[int] = None + context_parallel_size: int = 1 + dynamic_context_parallel: bool = False + max_seqlen_per_dp_cp_rank: Optional[int] = None + sequence_parallel: bool = True + use_distributed_optimizer: bool = True + use_dist_checkpointing: bool = False + dist_checkpointing_path: Optional[str] = None + dist_checkpointing_prefix: str = "" + dist_ckpt_optim_fully_reshardable: bool = False + distrib_optim_fully_reshardable_mem_efficient: bool = False + override_ddp_config: dict[str, Any] = field(default_factory=dict) + override_transformer_config: dict[str, Any] = field(default_factory=dict) + override_mcore_model_config: dict[str, Any] = field(default_factory=dict) + use_mbridge: bool = True + vanilla_mbridge: bool = True + use_megatron_fsdp: bool = False + strategy: str = "megatron" + qat: QATEngineConfig = field(default_factory=QATEngineConfig) + + def __post_init__(self) -> None: + super().__post_init__() + """config validation logics go here""" + assert self.strategy == "megatron" + assert self.dtype in ["bfloat16", "float16"], f"dtype {self.dtype} not supported" + if self.tensor_model_parallel_size == 1: + warnings.warn("set sequence parallel to false as TP size is 1", stacklevel=2) + self.sequence_parallel = False + + +@dataclass +class FSDPEngineConfig(EngineConfig): + """Configuration for FSDP (Fully Sharded Data Parallel). + + The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config. + + Args: + wrap_policy (Dict[str, Any]): Configuration for FSDP wrap policy. + param_offload (bool): Whether to offload parameters to CPU, default False + optimizer_offload (bool): Whether to offload optimizer states to CPU, default False + offload_policy (bool): Whether to offload policy model parameters, default False + reshard_after_forward (bool): Whether to reshard parameters after forward pass, default True + fsdp_size (int): FSDP group size. -1 means use all available GPUs. + forward_prefetch (bool): Whether to prefetch parameters for next forward pass, default False + model_dtype (str): Model data type used to initialize the transformers model. default "fp32" + use_orig_params (bool): Whether to use original parameters when initialize FSDP1, default False + seed (int): Random seed for reproducibility. + full_determinism (bool): If true, enable_full_determinism is called to ensure reproducible results + in distributed training. Important: this will negatively impact performance, so only use it for + debugging. + mixed_precision (Optional[dict[str, Any]]): Mixed precision configuration for FSDP, default None + dtype (str): Mixed precision training param dtype, default "bfloat16" + qat (QATEngineConfig): QAT configuration, default disabled + """ + + # ulysses_sequence_parallel_size is mutable for backward compatibility + _mutable_fields = EngineConfig._mutable_fields | {"ulysses_sequence_parallel_size"} + + # fsdp specific flags + wrap_policy: dict[str, Any] = field(default_factory=dict) + offload_policy: bool = False + reshard_after_forward: bool = True + fsdp_size: int = -1 + forward_prefetch: bool = False + model_dtype: str = "fp32" + use_orig_params: bool = False + mixed_precision: Optional[dict[str, Any]] = None + ulysses_sequence_parallel_size: int = 1 + entropy_from_logits_with_chunking: bool = False + use_torch_compile: bool = True + entropy_checkpointing: bool = False + strategy: str = "fsdp" + qat: QATEngineConfig = field(default_factory=QATEngineConfig) + + def __post_init__(self): + super().__post_init__() + assert self.strategy in ["fsdp", "fsdp2"], f"strategy {self.strategy} not supported" + + +@dataclass +class VeOmniEngineConfig(EngineConfig): + """Configuration for VeOmni. + + The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config. + + Args: + wrap_policy (Dict[str, Any]): Configuration for FSDP wrap policy. + param_offload (bool): Whether to offload parameters to CPU, default False + optimizer_offload (bool): Whether to offload optimizer states to CPU, default False + offload_policy (bool): Whether to offload policy model parameters, default False + reshard_after_forward (bool): Whether to reshard parameters after forward pass, default True + fsdp_size (int): FSDP group size. -1 means use all available GPUs, default -1 + ulysses_parallel_size (int): Ulysses sequence parallel size, default 1 + expert_parallel_size (int): Expert parallel size, default 1 + init_device (str): Device to initialize model weights. + 1. `cpu`: Init parameters on CPU in rank0 only. + 2. `cuda`: Init parameters on GPU. + 3. `meta`: Init parameters on meta. + 4. `npu`: Init parameters on Ascend NPU. + default "meta" + enable_full_shard (bool): Enable fully shard for FSDP training (ZeRO-3), default False + enable_fsdp_offload (bool): Enable CPU offload for FSDP1, default False + enable_reentrant (bool): Use reentrant gradient checkpointing, default False + attn_implementation (str): Attention implementation to use. + 1. `eager` + 2. `sdpa` + 3. `flash_attention_2` + 4. `flash_attention_3` + 5. `veomni_flash_attention_2_with_sp` + 6. `veomni_flash_attention_3_with_sp` + 7. `native-sparse` + default "flash_attention_2" + Note: In case VeOmni add more attn_implementation, please check https://github.com/ByteDance-Seed/VeOmni/ + moe_implementation (str): MoE implementation to use. + 1. `eager` + 2. `fused` + default "fused" + Note: In case VeOmni add more moe_implementation, please check https://github.com/ByteDance-Seed/VeOmni/ + cross_entropy_loss_implementation (str): Cross-entropy kernel selected via VeOmni's + ``OpsImplementationConfig``. Common values: ``"eager"`` (default), ``"liger_kernel"``, + ``"npu"``. See VeOmni docs for the full registry. + rms_norm_implementation (str): RMSNorm kernel. ``"eager"`` (HF default), + ``"triton"`` (batch-invariant Triton kernel — required to keep vexact's rollout + and the FSDP actor bitwise-aligned on DeepSeek-V3 / Moonlight), ``"liger_kernel"``, + ``"npu"``. + swiglu_mlp_implementation (str): SwiGLU MLP kernel. ``"eager"`` (default) or + ``"liger_kernel"``. + rotary_pos_emb_implementation (str): RoPE kernel. ``"eager"`` (default), ``"triton"`` + (deterministic Triton bmm — required for bitwise-aligned RoPE on DeepSeek-V3 / + Moonlight), ``"liger_kernel"``, ``"npu"``. + load_balancing_loss_implementation (str): MoE load-balancing loss kernel. + ``"eager"`` (default) or ``"triton"``. + force_use_huggingface (bool): Force loading model from huggingface, default False + activation_gpu_limit (float): When enabling activation offload, `activation_gpu_limit` GB + activations are allowed to reserve on GPU, default 0.0 + basic_modules (list[str]): List of basic modules to use, default None + forward_prefetch (bool): Whether to prefetch parameters for next forward pass, default False + model_dtype (str): Model data type used to initialize the transformers model. default "fp32" + use_orig_params (bool): Whether to use original parameters when initialize FSDP1, default False + seed (int): Random seed for reproducibility. + full_determinism (bool): If true, enable_full_determinism is called to ensure reproducible results + in distributed training. Important: this will negatively impact performance, so only use it for + debugging. + mixed_precision (Optional[dict[str, Any]]): Mixed precision configuration for FSDP, default None + + """ + + _mutable_fields = EngineConfig._mutable_fields | {"attn_implementation"} + + wrap_policy: dict[str, Any] = field(default_factory=dict) + offload_policy: bool = False + reshard_after_forward: bool = True + forward_prefetch: bool = False + use_orig_params: bool = False + entropy_from_logits_with_chunking: bool = False + use_torch_compile: bool = True + entropy_checkpointing: bool = False + strategy: str = "veomni" + fsdp_size: int = -1 + ulysses_parallel_size: int = 1 + expert_parallel_size: int = 1 + seed: int = 42 + full_determinism: bool = False + mixed_precision: bool = False + init_device: str = "meta" + enable_full_shard: bool = False + ckpt_manager: Literal["dcp"] = "dcp" + load_checkpoint_path: Optional[str] = None + enable_fsdp_offload: bool = False + enable_reentrant: bool = False + attn_implementation: str = "flash_attention_2" + moe_implementation: str = "fused" + # Kernel-backend selectors for VeOmni's per-model patches; passed into + # OpsImplementationConfig and consumed by apply_per_model_patches in each + # model's device_patch.py. Defaults match VeOmni's OpsImplementationConfig + # defaults so existing configs see no change. + cross_entropy_loss_implementation: str = "eager" + rms_norm_implementation: str = "eager" + swiglu_mlp_implementation: str = "eager" + rotary_pos_emb_implementation: str = "eager" + load_balancing_loss_implementation: str = "eager" + force_use_huggingface: bool = False + activation_gpu_limit: float = 0.0 + basic_modules: Optional[list[str]] = field(default_factory=list) + + def __post_init__(self): + super().__post_init__() + assert self.strategy in ["veomni"], f"strategy {self.strategy} not supported" + + replacements = { + "flash_attention_2": "veomni_flash_attention_2_with_sp", + "flash_attention_3": "veomni_flash_attention_3_with_sp", + "flash_attention_4": "veomni_flash_attention_4_with_sp", + } + if self.attn_implementation in replacements: + new_impl = replacements[self.attn_implementation] + logger.info(f"Replacing attn_implementation from '{self.attn_implementation}' to '{new_impl}'") + self.attn_implementation = new_impl + + +@dataclass +class TorchtitanEngineConfig(EngineConfig): + """Configuration for Torchtitan. + + The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config. + + Args: + wrap_policy (Dict[str, Any]): Configuration for FSDP wrap policy. + reshard_after_forward (Literal["default", "always", "never"]): The policy for applying + `reshard_after_forward` within an FSDP setup, default "default" + forward_prefetch (bool): Whether to prefetch parameters for next forward pass, default False + use_orig_params (bool): Whether to use original parameters when initialize FSDP, default False + mixed_precision (bool): Mixed precision configuration for FSDP, default False + offload_policy (bool): Whether to offload policy model parameters, default False + data_parallel_size (int): Data parallel group size, default 1 + data_parallel_replicate_size (int): Data parallel replicate size, default 1 + data_parallel_shard_size (int): Data parallel shard degree, default 1 + tensor_parallel_size (int): Tensor parallel size, default 1 + expert_parallel_size (int): Expert parallel size, default 1 + expert_tensor_parallel_size (int): Expert tensor parallel size, default 1 + pipeline_parallel_size (int): Pipeline parallel size, default 1 + context_parallel_size (int): Context parallel size, default 1 + attn_type (str): Attention type for torchtitan's model (e.g., "sdpa", "flex", "varlen"), + default "flex" + strategy (str): Strategy to use for distributed training, default "torchtitan" + seed (int): Random seed for reproducibility. + full_determinism (bool): If true, enable_full_determinism is called to ensure reproducible results + in distributed training. Important: this will negatively impact performance, so only use it for + debugging. + + """ + + wrap_policy: dict[str, Any] = field(default_factory=dict) + reshard_after_forward: Literal["default", "always", "never"] = "default" + forward_prefetch: bool = False + use_orig_params: bool = False + mixed_precision: bool = False + offload_policy: bool = False + use_torch_compile: bool = True + entropy_from_logits_with_chunking: bool = False + entropy_checkpointing: bool = False + data_parallel_size: int = 1 + data_parallel_replicate_size: int = 1 + data_parallel_shard_size: int = 1 + tensor_parallel_size: int = 1 + expert_parallel_size: int = 1 + expert_tensor_parallel_size: int = 1 + pipeline_parallel_size: int = 1 + context_parallel_size: int = 1 + attn_type: str = "flex" + max_seq_len: Optional[int] = None + strategy: str = "torchtitan" + seed: int = 42 + full_determinism: bool = False + + def __post_init__(self): + super().__post_init__() + assert self.strategy in ["torchtitan"], f"strategy {self.strategy} not supported" + + +@dataclass +class AutomodelEngineConfig(EngineConfig): + """Configuration for Automodel (nemo_automodel) backend. + + The Automodel backend uses NeMoAutoModelForCausalLM for model loading and + supports FSDP2, MegatronFSDP, and DDP distributed strategies with optional + TP, CP, and EP parallelism. + + Args: + strategy (str): Backend strategy identifier, must be "automodel". + distributed_strategy (str): Distributed training strategy: "fsdp2", "megatron_fsdp", or "ddp". + tp_size (int): Tensor parallel size. + pp_size (int): Pipeline parallel size (only pp_size=1 supported initially). + cp_size (int): Context parallel size. + ep_size (int): Expert parallel size for MoE models. + dp_replicate_size (int): Data-parallel replicate size for HSDP. 1 = pure sharding. + sequence_parallel (bool): Enable sequence parallelism in the TP plan. + defer_fsdp_grad_sync (bool): Defer FSDP gradient sync to the final micro-batch. + activation_checkpointing (bool): Whether to enable activation checkpointing. + enable_fp8 (bool): Whether to enable FP8 training. + enable_compile (bool): Whether to enable torch.compile for the model. + model_dtype (str): Model data type for loading weights. "fp32" loads in float32 + (matching FSDP golden), "auto" uses the dtype from the model config. + attn_implementation (str): Attention implementation to use ("sdpa", "flash_attention_2", "eager", "te"). + + Backend settings (nemo_automodel BackendConfig): + backend_config (dict): Dict of kwargs passed directly to + nemo_automodel.components.models.common.BackendConfig(**backend_config). + Controls how model layers are implemented (TE vs PyTorch) and MoE dispatch. + See automodel.yaml for all predefined keys with defaults. + Key fields: + attn (str): Attention backend. "te" = TransformerEngine fused attention, + "sdpa" = PyTorch scaled dot-product attention. Default: "sdpa". + linear (str): Linear layer backend. "te" = TE fused linear (with FP8 support), + "torch" = standard PyTorch linear. Default: "te". + rms_norm (str): RMSNorm backend. "te" = TE fused RMSNorm, "torch" = PyTorch, + "torch_fp32" = PyTorch in FP32 (better numerical stability for MoE). + Default: "torch_fp32". + rope_fusion (bool): Enable fused RoPE kernel (requires CP=1). Default: true. + experts (str): MoE expert computation backend. + "gmm" = grouped_gemm (requires pip install grouped_gemm), + "torch_mm" = torch._grouped_mm (no external dependency), + "te" = TE GroupedLinear. Default: "gmm". + dispatcher (str): MoE token dispatch strategy. + "torch" = standard all-gather + local compute, + "deepep" = DeepEP optimized all-to-all (higher throughput). + Default: "torch". + Note: "deepep" with experts="gmm" matches the legacy enable_deepep=True behavior. + enable_fsdp_optimizations (bool): Enable FSDP-specific optimizations in Automodel. + Default: false. + enable_hf_state_dict_adapter (bool): Enable HuggingFace state dict adapter for + checkpoint compatibility. Default: true. + fake_balanced_gate (bool): Use fake balanced gating for debugging. Default: false. + fake_gate_noise (float): Noise added to fake balanced gate. Default: 0.0. + gate_precision: Gate computation precision. Default: null (auto). + Full reference: nemo_automodel/components/models/common/backend_config.py + + MoE / Expert Parallelism settings: + moe_config (dict): Dict of kwargs passed directly to + nemo_automodel.components.moe.parallelizer.MoEParallelizerConfig(**moe_config). + Controls MoE parallelization behavior within FSDP2. + See automodel.yaml for all predefined keys with defaults. + Key fields: + ignore_router_for_ac (bool): Exclude router from activation checkpointing. + Default: false. + reshard_after_forward (bool): Reshard expert params after forward pass + (trades compute for memory). Default: false. + lm_head_precision: Precision for the LM head. Default: null (auto). + wrap_outer_model (bool): Whether to FSDP-wrap the outermost model module. + Default: true. + Full reference: nemo_automodel/components/moe/parallelizer.py + + Mixed precision policy (FSDP2): + mp_param_dtype (str): Parameter dtype for FSDP2 mixed precision policy. + mp_reduce_dtype (str): Reduce dtype for FSDP2 mixed precision policy. + mp_output_dtype (str): Output dtype for FSDP2 mixed precision policy. + + Entropy computation: + entropy_from_logits_with_chunking (bool): Whether to use chunked entropy computation. + use_torch_compile (bool): Whether to use torch.compile for entropy computation. + entropy_checkpointing (bool): Whether to use checkpointing for entropy computation. + """ + + strategy: str = "automodel" + distributed_strategy: str = "fsdp2" + # Parallelism sizes + tp_size: int = 1 + pp_size: int = 1 + cp_size: int = 1 + ep_size: int = 1 + dp_replicate_size: int = 1 + sequence_parallel: bool = False + defer_fsdp_grad_sync: bool = True + # Model settings + activation_checkpointing: bool = False + enable_fp8: bool = False + enable_compile: bool = False + model_dtype: str = "fp32" + attn_implementation: str = "flash_attention_2" + # Backend settings + backend_config: dict = field(default_factory=dict) + # MoE settings + moe_config: dict = field(default_factory=dict) + # Mixed precision policy + mp_param_dtype: str = "bf16" + mp_reduce_dtype: str = "fp32" + mp_output_dtype: str = "bf16" + # Entropy computation + entropy_from_logits_with_chunking: bool = False + use_torch_compile: bool = True + entropy_checkpointing: bool = False + + def __post_init__(self): + super().__post_init__() + assert self.strategy == "automodel", f"strategy must be 'automodel', got {self.strategy}" + assert self.distributed_strategy in ["fsdp2", "megatron_fsdp", "ddp"], ( + f"distributed_strategy {self.distributed_strategy} not supported" + ) + assert self.pp_size == 1, "Pipeline parallelism (pp_size > 1) is not yet supported for automodel backend" + + +@dataclass +class MindSpeedEngineConfig(McoreEngineConfig): + """Configuration for mindspeed parallelism. + + The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config. + + Args: + llm_kwargs (str): mindspeed_llm engine kwargs. + mm_kwargs (str): mindspeed_mm engine kwargs. + """ + + strategy: str = "mindspeed_llm" + llm_kwargs: dict[str, Any] = field(default_factory=dict) + mm_kwargs: dict[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + """config validation logics go here""" + assert self.strategy in ["mindspeed_llm", "mindspeed_mm"], f"strategy {self.strategy} not supported" + assert self.dtype in ["bfloat16", "float16"], f"dtype {self.dtype} not supported" + if self.tensor_model_parallel_size == 1: + warnings.warn("set sequence parallel to false as TP size is 1", stacklevel=2) + self.sequence_parallel = False + + +@dataclass +class TrainingWorkerConfig(BaseConfig): + model_type: str = None # model type (language_model/value_model) + model_config: HFModelConfig = None + engine_config: EngineConfig = None + optimizer_config: OptimizerConfig = None + checkpoint_config: CheckpointConfig = None + profiler_config: ProfilerConfig = None + # automatically select engine and optimizer function. + # This function takes model config and the device name as parameter. + # Users can pass in a higher-order function to take more parameters + auto_select_engine_optim_fn: Callable[["HFModelConfig", str], tuple["EngineConfig", "OptimizerConfig"]] = None diff --git a/verl/verl/workers/config/megatron_peft.py b/verl/verl/workers/config/megatron_peft.py new file mode 100644 index 0000000000000000000000000000000000000000..15d7ce46eeb66eaa073fa862011c38b181d77928 --- /dev/null +++ b/verl/verl/workers/config/megatron_peft.py @@ -0,0 +1,121 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PEFT configuration of Megatron for VERL.""" + + +def get_peft_cls(model_config, bridge, provider, dtype=None): + """Get PEFT class from model config. + + Args: + model_config: Model configuration object. + bridge: Megatron-Bridge AutoBridge instance. + provider: Provider instance. + + Returns: + PEFT configuration object (LoRAConfig, CanonicalLoRAConfig, DoRAConfig) or None. + """ + + peft_cls = None + if not hasattr(model_config, "lora"): + return peft_cls + + lora_cfg = model_config.lora + # Only enable if rank > 0 + if lora_cfg.get("rank", 0) <= 0: + return peft_cls + + assert bridge is not None and provider is not None, "LoRA/PEFT only supported via Megatron-Bridge" + + from verl.models.mcore.bridge import CanonicalLoRA, DoRA, LoRA, VLMLoRA + + lora_dtype = lora_cfg.get("dtype", dtype) + if lora_dtype is not None: + from verl.utils.torch_dtypes import PrecisionType + + lora_dtype = PrecisionType.to_dtype(lora_dtype) + + lora_type = lora_cfg.get("type", "lora") + if lora_type == "lora": + peft_cls = LoRA( + target_modules=lora_cfg.get("target_modules", ["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2"]), + dim=lora_cfg.get("rank"), + alpha=lora_cfg.get("alpha", 32), + dropout=lora_cfg.get("dropout", 0.0), + dropout_position=lora_cfg.get("dropout_position", "pre"), + lora_A_init_method=lora_cfg.get("lora_A_init_method", "xavier"), + lora_B_init_method=lora_cfg.get("lora_B_init_method", "zero"), + a2a_experimental=lora_cfg.get("a2a_experimental", False), + lora_dtype=lora_dtype, + exclude_modules=lora_cfg.get("exclude_modules", []), + ) + if lora_type == "vlm_lora": + peft_cls = VLMLoRA( + target_modules=lora_cfg.get("target_modules", ["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2"]), + dim=lora_cfg.get("rank"), + alpha=lora_cfg.get("alpha", 32), + dropout=lora_cfg.get("dropout", 0.0), + dropout_position=lora_cfg.get("dropout_position", "pre"), + lora_A_init_method=lora_cfg.get("lora_A_init_method", "xavier"), + lora_B_init_method=lora_cfg.get("lora_B_init_method", "zero"), + a2a_experimental=lora_cfg.get("a2a_experimental", False), + lora_dtype=lora_dtype, + freeze_vision_model=lora_cfg.get("freeze_vision_model", True), + freeze_vision_projection=lora_cfg.get("freeze_vision_projection", True), + freeze_language_model=lora_cfg.get("freeze_language_model", True), + exclude_modules=lora_cfg.get("exclude_modules", []), + ) + elif lora_type == "canonical_lora": + peft_cls = CanonicalLoRA( + target_modules=lora_cfg.get( + "target_modules", + [ + "linear_q", + "linear_k", + "linear_v", + "linear_proj", + "linear_fc1_up", + "linear_fc1_gate", + "linear_fc2", + ], + ), + dim=lora_cfg.get("rank"), + alpha=lora_cfg.get("alpha", 32), + dropout=lora_cfg.get("dropout", 0.0), + dropout_position=lora_cfg.get("dropout_position", "pre"), + lora_A_init_method=lora_cfg.get("lora_A_init_method", "xavier"), + lora_B_init_method=lora_cfg.get("lora_B_init_method", "zero"), + exclude_modules=lora_cfg.get("exclude_modules", []), + ) + elif lora_type == "dora": + peft_cls = DoRA( + target_modules=lora_cfg.get("target_modules", ["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2"]), + dim=lora_cfg.get("rank"), + alpha=lora_cfg.get("alpha", 32), + dropout=lora_cfg.get("dropout", 0.0), + dropout_position=lora_cfg.get("dropout_position", "pre"), + lora_A_init_method=lora_cfg.get("lora_A_init_method", "xavier"), + lora_B_init_method=lora_cfg.get("lora_B_init_method", "zero"), + exclude_modules=lora_cfg.get("exclude_modules", []), + ) + + print( + f"Enabling {lora_type.upper()} with rank={lora_cfg.get('rank')}, " + f"alpha={lora_cfg.get('alpha')}, dropout={lora_cfg.get('dropout')}" + ) + return peft_cls + + +__all__ = [ + "get_peft_cls", +] diff --git a/verl/verl/workers/config/model.py b/verl/verl/workers/config/model.py new file mode 100644 index 0000000000000000000000000000000000000000..95814663dcaa91f7e2984b1d4c39ca042712c485 --- /dev/null +++ b/verl/verl/workers/config/model.py @@ -0,0 +1,247 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from dataclasses import dataclass, field +from typing import Any, Optional + +from omegaconf import MISSING +from transformers import AutoConfig + +from verl.base_config import BaseConfig +from verl.utils import hf_processor, hf_tokenizer +from verl.utils.fs import copy_to_local +from verl.utils.import_utils import import_external_libs +from verl.utils.model import get_generation_config, update_model_config + +__all__ = ["HFModelConfig", "MtpConfig"] + + +@dataclass +class MtpConfig(BaseConfig): + """ + Configuration for MTP model. + + enable: Enable loading and saving of MTP parameters, but do not use them + + enable_train: Whether to enable using MTP parameters during training + enable_rollout: Whether to enable using MTP parameters during rollout + + Training parameters: + detach_encoder: Whether to detach encoder parameters during MTP training + mtp_loss_scaling_factor: Loss scaling factor during MTP training + + vLLM rollout parameters: + method: "mtp" + num-speculative-tokens: 1 + + SGLang rollout parameters: + speculative-algorithm: EAGLE + speculative-num-steps: 3 + speculative-eagle-topk: 1 + speculative-num-draft-tokens: 4 + """ + + enable: bool = False + enable_train: bool = False + enable_rollout: bool = False + + detach_encoder: bool = False + mtp_loss_scaling_factor: float = 0.1 + + speculative_algorithm: str = "EAGLE" + speculative_num_steps: int = 3 + speculative_eagle_topk: int = 1 + speculative_num_draft_tokens: int = 4 + + method: str = "mtp" + num_speculative_tokens: int = 1 + + +@dataclass +class HFModelConfig(BaseConfig): + # note that we separate model_path, model_config_path and tokenizer_path in case they are different + _mutable_fields = { + "model_type", + "hf_config_path", + "tokenizer_path", + "hf_config", + "generation_config", + "tokenizer", + "processor", + "local_path", + "architectures", + "local_hf_config_path", + "local_tokenizer_path", + "mtp", + } + + path: str = MISSING + local_path: Optional[str] = None + hf_config_path: Optional[str] = None + local_hf_config_path: Optional[str] = None + tokenizer_path: Optional[str] = None + local_tokenizer_path: Optional[str] = None + + # model type, e.g., "language_model", "value_model" + model_type: str = "language_model" + + # whether to load tokenizer. This is useful when we only want to load model config + load_tokenizer: bool = True + + hf_config: Any = None + generation_config: Any = None + tokenizer: Any = None + processor: Any = None + + # whether to use shared memory + use_shm: bool = False + trust_remote_code: bool = False + + # custom chat template for the model + custom_chat_template: Optional[str] = None + + external_lib: Optional[str] = None + + override_config: dict = field(default_factory=dict) + + enable_gradient_checkpointing: bool = True + enable_activation_offload: bool = False + + use_remove_padding: bool = True + + # TODO: unify fsdp and megatron lora config + # fsdp lora related. We may setup a separate config later + lora_rank: int = 0 + lora_alpha: int = 16 + target_modules: Optional[Any] = "all-linear" # allow both "all-linear" and ["q_proj","k_proj"] + target_parameters: Optional[list[str]] = None # for lora adapter on nn.Parameter + + exclude_modules: Optional[str] = None + + # megatron lora config + lora: dict[str, Any] = field(default_factory=dict) + + # path to pre-trained LoRA adapter to load for continued training + lora_adapter_path: Optional[str] = None + use_liger: bool = False + + use_fused_kernels: bool = False + fused_kernel_options: dict = field(default_factory=dict) + + # TiledMLP configuration for memory-efficient MLP computation + tiled_mlp: dict = field(default_factory=lambda: {"enabled": False, "num_shards": 4}) + + architectures: Optional[list[str]] = None + + mtp: MtpConfig = field(default_factory=MtpConfig) + + def __post_init__(self): + import_external_libs(self.external_lib) + + if self.hf_config_path is None: + self.hf_config_path = self.path + if self.tokenizer_path is None: + self.tokenizer_path = self.path + + self.local_path = copy_to_local(self.path, use_shm=self.use_shm) + + # construct tokenizer + if self.load_tokenizer: + self.local_tokenizer_path = copy_to_local(self.tokenizer_path, use_shm=self.use_shm) + self.tokenizer = hf_tokenizer(self.local_tokenizer_path, trust_remote_code=self.trust_remote_code) + self.processor = hf_processor(self.local_tokenizer_path, trust_remote_code=self.trust_remote_code) + + # For base models (e.g. Qwen3.5-2b-Base), the processor may not have a chat_template + # while the tokenizer does. Sync it so that processor.apply_chat_template() works. + if ( + self.processor is not None + and not getattr(self.processor, "chat_template", None) + and getattr(self.tokenizer, "chat_template", None) + ): + self.processor.chat_template = self.tokenizer.chat_template + + if self.custom_chat_template is not None: + if self.processor is not None: + self.processor.chat_template = self.custom_chat_template + else: + self.tokenizer.chat_template = self.custom_chat_template + + self.local_hf_config_path = copy_to_local(self.hf_config_path, use_shm=self.use_shm) + self.generation_config = get_generation_config( + self.local_hf_config_path, trust_remote_code=self.trust_remote_code + ) + + # construct hf_config + attn_implementation = self.override_config.get("attn_implementation", "flash_attention_2") + self.hf_config = AutoConfig.from_pretrained( + self.local_hf_config_path, trust_remote_code=self.trust_remote_code, attn_implementation=attn_implementation + ) + + override_config_kwargs = {} + + if self.tokenizer is not None: + override_config_kwargs.update( + { + "bos_token_id": self.tokenizer.bos_token_id, + "eos_token_id": self.tokenizer.eos_token_id, + "pad_token_id": self.tokenizer.pad_token_id, + } + ) + + # TODO: (vermouth1992). self.config.model in megatron differs from that of fsdp in the override_config. + override_config = ( + self.override_config["model_config"] if "model_config" in self.override_config else self.override_config + ) + override_config_kwargs.update(override_config) + update_model_config(self.hf_config, override_config_kwargs=override_config_kwargs) + + self.share_embeddings_and_output_weights = getattr(self.hf_config, "tie_word_embeddings", False) + + # get model architectures + self.architectures = getattr(self.hf_config, "architectures", None) + assert self.architectures is not None and len(self.architectures) == 1, ( + "Expect only one architecture, got {}".format(self.architectures) + ) + + # per model patch + if getattr(self.hf_config, "model_type", None) == "kimi_vl": + self.hf_config.text_config.topk_method = "greedy" + + # When MTP is disabled, zero out MTP layer counts from hf_config so that + # downstream engine/worker code does not need to handle each MTP field format + # individually. Supports both DeepSeek-style (num_nextn_predict_layers) and + # Qwen3.5-style (mtp_num_hidden_layers, possibly nested under text_config). + if not self.mtp.enable: + if hasattr(self.hf_config, "num_nextn_predict_layers"): + self.hf_config.num_nextn_predict_layers = 0 + if hasattr(self.hf_config, "mtp_num_hidden_layers"): + self.hf_config.mtp_num_hidden_layers = 0 + if hasattr(self.hf_config, "text_config") and hasattr(self.hf_config.text_config, "mtp_num_hidden_layers"): + self.hf_config.text_config.mtp_num_hidden_layers = 0 + + # Ensure target_modules is a str or list[str] (only if not None) + if self.target_modules is not None: + if not isinstance(self.target_modules, (str | list)): + raise TypeError( + "target_modules must be a string or a list of strings, " + f"but got {type(self.target_modules).__name__}" + ) + if isinstance(self.target_modules, list): + for x in self.target_modules: + if not isinstance(x, str): + raise TypeError( + f"All elements in target_modules list must be strings, but found {type(x).__name__}" + ) + + def get_processor(self): + return self.processor if self.processor is not None else self.tokenizer diff --git a/verl/verl/workers/config/optimizer.py b/verl/verl/workers/config/optimizer.py new file mode 100644 index 0000000000000000000000000000000000000000..47afdd3bf2e14ca1be39d14720775284d92d19a4 --- /dev/null +++ b/verl/verl/workers/config/optimizer.py @@ -0,0 +1,271 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import warnings +from dataclasses import dataclass +from typing import Optional + +from omegaconf import MISSING + +from verl.base_config import BaseConfig + +__all__ = [ + "OptimizerConfig", + "FSDPOptimizerConfig", + "McoreOptimizerConfig", + "build_optimizer", + "VeOmniOptimizerConfig", + "TorchtitanOptimizerConfig", + "AutomodelOptimizerConfig", +] + + +@dataclass +class OptimizerConfig(BaseConfig): + """Base optimizer configuration. + + Args: + lr (float): learning rate. Must be specified. + lr_warmup_steps_ratio (float): Warmup steps ratio; total steps will be injected at runtime. + total_training_steps (int): Total training steps (must be overridden at runtime). + weight_decay (float): Weight decay factor. + lr_warmup_steps (Optional[int]): Number of warmup steps; None delegates to lr_warmup_steps_ratio. + """ + + _mutable_fields = {"clip_grad", "total_training_steps", "lr_warmup_steps"} + + lr: float = 1e-3 + lr_warmup_steps_ratio: float = 0.0 + total_training_steps: int = -1 + weight_decay: float = 0.01 + lr_warmup_steps: Optional[int] = -1 + betas: tuple[float, float] = (0.9, 0.999) + clip_grad: float = 1.0 + # deprecate grad_clip + grad_clip: Optional[float] = None + + def __post_init__(self): + assert self.lr != MISSING + if self.grad_clip is not None: + warnings.warn("`grad_clip` is deprecated, use `clip_grad` instead.", DeprecationWarning, stacklevel=2) + self.clip_grad = self.grad_clip + + +@dataclass +class VeOmniOptimizerConfig(OptimizerConfig): + """VeOmni optimizer configuration extending base OptimizerConfig. + + Args: + optimizer (str): Optimizer name; default is "adamw". + lr (float): Learning rate. + lr_min (float): Minimum learning rate. + lr_start (float): Starting learning rate for warmup. + lr_decay_ratio (float): LR decay ratio. + lr_scheduler_type (str): LR scheduler type: "constant" or "cosine". + """ + + _mutable_fields = OptimizerConfig._mutable_fields.copy() + + optimizer: str = "adamw" + lr_min: float = 0.0 + lr_start: float = 0.0 + lr_decay_ratio: float = 1.0 + lr_scheduler_type: str = "constant" + override_optimizer_config: Optional[dict] = None + + +@dataclass +class FSDPOptimizerConfig(OptimizerConfig): + """FSDP optimizer configuration extending base OptimizerConfig. + + Args: + optimizer (str): Optimizer class name (e.g., "AdamW", "AdamW8bit", "_AdamW"). + optimizer_impl (str): Module path to import optimizer from (e.g., "torch.optim", "torchao.optim", + "bitsandbytes.optim"). + lr (float): Learning rate. + min_lr_ratio (Optional[float]): Minimum LR ratio for cosine schedule. + lr_scheduler_type (str): LR scheduler type: "constant" or "cosine". + num_cycles (float): Number of cosine cycles in LR schedule. + zero_indexed_step (bool): Whether the LR schedule uses 0-indexed steps. If True (default), + step counting starts at 0. If False, step counting starts at 1. + """ + + _mutable_fields = OptimizerConfig._mutable_fields.copy() + _mutable_fields.add("lr_scheduler_type") + + optimizer: str = "AdamW" + optimizer_impl: str = "torch.optim" + min_lr_ratio: Optional[float] = None + # deprecate warmup_style + warmup_style: Optional[str] = None + lr_scheduler_type: str = "constant" + num_cycles: float = 0.5 + override_optimizer_config: Optional[dict] = None + zero_indexed_step: bool = True + + def __post_init__(self): + if self.warmup_style is not None: + assert self.warmup_style in ["constant", "cosine"] + warnings.warn( + "`warmup_style` is deprecated, use `lr_scheduler_type` instead.", DeprecationWarning, stacklevel=2 + ) + self.lr_scheduler_type = self.warmup_style + assert self.lr_scheduler_type in ["constant", "cosine"] + return super().__post_init__() + + +@dataclass +class McoreOptimizerConfig(OptimizerConfig): + """Mcore optimizer configuration extending base OptimizerConfig. + + Args: + optimizer (str): Optimizer name; default is "adam". + lr (float): Learning rate. + clip_grad (float): Gradient clipping norm. + lr_warmup_init (float): Initial learning rate for warmup; defaults to 0.0. + lr_decay_steps (Optional[int]): Number of decay steps. + lr_decay_style (str): LR decay style: "constant", "linear", "cosine", or "inverse_square_root". + min_lr (float): Minimum learning rate. + weight_decay_incr_style (str): Weight decay increment style: "constant" or "cosine". + lr_wsd_decay_style (str): Weight-standard-deviation decay style: "constant", "exponential", or "cosine". + lr_wsd_decay_steps (Optional[int]): Number of steps for weight-standard-deviation decay. + use_checkpoint_opt_param_scheduler (bool): Whether to use checkpoint optimizer parameter scheduler. + """ + + optimizer: str = "adam" + lr_warmup_init: float = 0.0 + lr_decay_steps: Optional[int] = None + lr_decay_style: str = "linear" + min_lr: float = 0.0 + weight_decay_incr_style: str = "constant" + lr_wsd_decay_style: str = "exponential" + lr_wsd_decay_steps: Optional[int] = None + use_checkpoint_opt_param_scheduler: bool = False + override_optimizer_config: Optional[dict] = None + + +@dataclass +class TorchtitanOptimizerConfig(OptimizerConfig): + """Torchtitan optimizer configuration extending base OptimizerConfig. + + Args: + name (str): Optimizer name; default is "AdamW". + eps (float): Epsilon value for AdamW optimizer, default 1e-8. + decay_type (str): Weight decay type: "linear", "sqrt", or "cosine". + min_lr_factor (float): Minimum learning rate factor. + """ + + name: str = "AdamW" + eps: float = 1e-8 + decay_type: str = "linear" + min_lr_factor: float = 0.0 + + +@dataclass +class AutomodelOptimizerConfig(OptimizerConfig): + """Automodel optimizer configuration extending base OptimizerConfig. + + Uses the same optimizer building mechanism as FSDP (dynamic import from optimizer_impl). + LR scheduling is handled by Automodel's OptimizerParamScheduler. + + Args: + optimizer (str): Optimizer class name (e.g., "AdamW"). + optimizer_impl (str): Module path to import optimizer from (e.g., "torch.optim"). + lr (float): Learning rate (maps to max_lr in OptimizerParamScheduler). + init_lr_ratio (Optional[float]): Initial LR ratio for warmup start (init_lr = lr * init_lr_ratio). + min_lr_ratio (Optional[float]): Minimum LR ratio after decay (min_lr = lr * min_lr_ratio). + lr_scheduler_type (str): LR decay style: "constant", "cosine", "linear", or "inverse-square-root". + wd_incr_style (str): Weight decay increment style: "constant", "linear", or "cosine". + num_cycles (float): Kept for backward compatibility (unused by Automodel scheduler). + zero_indexed_step (bool): Kept for backward compatibility (unused by Automodel scheduler). + """ + + _mutable_fields = OptimizerConfig._mutable_fields.copy() + _mutable_fields.add("lr_scheduler_type") + + optimizer: str = "AdamW" + optimizer_impl: str = "torch.optim" + init_lr_ratio: Optional[float] = 0.1 + min_lr_ratio: Optional[float] = 0.01 + lr_scheduler_type: str = "cosine" + wd_incr_style: str = "constant" + num_cycles: float = 0.5 + zero_indexed_step: bool = True + # Common optimizer kwargs + eps: float = 1e-8 + master_weights: bool = False + store_param_remainders: bool = False + exp_avg_dtype: Optional[str] = None # "fp32", "bf16", "fp16", or "torch.float32" etc. + exp_avg_sq_dtype: Optional[str] = None # "fp32", "bf16", "fp16", or "torch.float32" etc. + master_weight_dtype: Optional[str] = None # "fp32", "bf16", "fp16", or "torch.float32" etc. + override_optimizer_config: Optional[dict] = None + + def __post_init__(self): + assert self.lr_scheduler_type in ["constant", "cosine", "linear", "inverse-square-root"] + return super().__post_init__() + + +def build_optimizer(parameters, config: FSDPOptimizerConfig): + """Build an optimizer based on the configuration. + + Dynamically imports and instantiates an optimizer class from the specified module. + + Args: + parameters: Model parameters to optimize + config: FSDPOptimizerConfig with optimizer settings + + Returns: + Optimizer instance + + Examples: + # PyTorch AdamW + config.optimizer_impl = "torch.optim" + config.optimizer = "AdamW" + + # TorchAO AdamW with bf16 stochastic rounding + config.optimizer_impl = "torchao.optim" + config.optimizer = "_AdamW" + config.override_optimizer_config = {"bf16_stochastic_round": True} + + # BitsAndBytes AdamW 8bit + config.optimizer_impl = "bitsandbytes.optim" + config.optimizer = "AdamW8bit" + """ + import importlib + + optimizer_args = { + "lr": config.lr, + "weight_decay": config.weight_decay, + } + + optimizer_name_lower = config.optimizer.lower() + if "adam" in optimizer_name_lower or "ademamix" in optimizer_name_lower: + optimizer_args["betas"] = config.betas + + if config.override_optimizer_config is not None: + optimizer_args.update(config.override_optimizer_config) + + try: + module = importlib.import_module(config.optimizer_impl) + optimizer_cls = getattr(module, config.optimizer) + except ImportError as e: + raise ImportError( + f"Failed to import module '{config.optimizer_impl}'. Make sure the package is installed. Error: {e}" + ) from e + except AttributeError as e: + raise AttributeError( + f"Optimizer '{config.optimizer}' not found in module '{config.optimizer_impl}'. " + f"Available optimizers: {dir(module)}" + ) from e + + return optimizer_cls(parameters, **optimizer_args) diff --git a/verl/verl/workers/config/reward.py b/verl/verl/workers/config/reward.py new file mode 100644 index 0000000000000000000000000000000000000000..f6b84385fa70f3a9d2a99add90ec610a9da56e0d --- /dev/null +++ b/verl/verl/workers/config/reward.py @@ -0,0 +1,105 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +import os +from dataclasses import dataclass, field +from typing import Optional + +from verl.base_config import BaseConfig +from verl.trainer.config.config import ModuleConfig + +from .rollout import RolloutConfig + +__all__ = ["SandboxFusionConfig", "RewardConfig", "RewardModelConfig"] + +logger = logging.getLogger(__name__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + + +@dataclass +class RewardManagerConfig(BaseConfig): + """Configuration for reward manager. + + A reward manager defines the mechanism of computing rule-based reward and handling different reward sources. + + Args: + source (str): Source of the reward manager. Options: ``"register"``, ``"importlib"``. Default: ``"register"``. + name (str, optional): + - When ``source`` is ``"register"``, the name is used in `get_reward_manager_cls(name)``. + See ``verl/experimental/reward/reward_manager.py`` for options. Default: ``"naive"``. + - When ``source`` is ``"importlib"``, the name is used in ``getattr(module, name)``, + e.g., ``"DAPORewardManager"``. + module (ModuleConfig, optional): Optional configuration for the external module defining the reward manager, + """ + + source: str = "register" + name: str = "naive" + module: Optional[ModuleConfig] = field(default_factory=ModuleConfig) + + def __post_init__(self): + super().__post_init__() + if self.source == "register": + from verl.experimental.reward_loop.reward_manager.registry import REWARD_MANAGER + + assert self.name in REWARD_MANAGER, ( + f"Reward manager is not registered: {self.name=} ,{REWARD_MANAGER.keys()=}" + ) + elif self.source == "importlib": + # NOTE: The existence is not checked since it depends on which machine the config is initialized on. + assert self.module is not None and self.module.path is not None, ( + "When source is importlib, module.path should be set." + ) + + +@dataclass +class SandboxFusionConfig(BaseConfig): + """Configuration for cloud/local sandbox fusion. + + Args: + url (Optional[str]): Cloud/local function URL for sandbox execution. + max_concurrent (int): Max concurrent requests allowed to sandbox. + memory_limit_mb (int): Max memory limit for each sandbox process in MB. + """ + + url: Optional[str] = None + max_concurrent: int = 64 + memory_limit_mb: int = 1024 + + +@dataclass +class RewardModelConfig(BaseConfig): + _mutable_fields = BaseConfig._mutable_fields + + enable: bool = False + enable_resource_pool: bool = False + n_gpus_per_node: int = 0 + nnodes: int = 0 + model_path: Optional[str] = None + rollout: RolloutConfig = field(default_factory=RolloutConfig) + + +@dataclass +class RewardConfig(BaseConfig): + _mutable_fields = BaseConfig._mutable_fields + + # reward manager args + num_workers: int = 8 + reward_manager: RewardManagerConfig = field(default_factory=RewardManagerConfig) + + # reward model args + reward_model: RewardModelConfig = field(default_factory=RewardModelConfig) + + # sandbox fusion args + sandbox_fusion: SandboxFusionConfig = field(default_factory=SandboxFusionConfig) diff --git a/verl/verl/workers/config/rollout.py b/verl/verl/workers/config/rollout.py new file mode 100644 index 0000000000000000000000000000000000000000..306e6d60086ed7008780fd0e1ab3b3fc348dc584 --- /dev/null +++ b/verl/verl/workers/config/rollout.py @@ -0,0 +1,353 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import warnings +from dataclasses import dataclass, field +from typing import Optional + +from omegaconf import MISSING + +from verl.base_config import BaseConfig +from verl.utils.profiler import ProfilerConfig +from verl.workers.config.disaggregation import DisaggregationConfig +from verl.workers.config.model import MtpConfig + +__all__ = [ + "SamplingConfig", + "MultiTurnConfig", + "CustomAsyncServerConfig", + "AgentLoopConfig", + "TraceConfig", + "ServerConfig", + "PrometheusConfig", + "RolloutConfig", + "CheckpointEngineConfig", + "SkipConfig", +] + + +@dataclass +class SkipConfig(BaseConfig): + """ + Configuration for rollout skip: load/dump previously generated rollout data + instead of computing new rollouts (e.g. for debugging or reuse). + """ + + enable: bool = False + dump_dir: str = "~/.verl/rollout_dump" + max_dump_step: int = 1 + action: str = "cache" # cache | repeat | repeat_last + + def get(self, key: str, default=None): + """Dict-like get for compatibility with code that uses skip.get('enable', False).""" + return getattr(self, key, default) + + +@dataclass +class SamplingConfig(BaseConfig): + temperature: float = 1.0 + top_k: int = -1 + top_p: float = 1.0 + do_sample: bool = True + n: int = 1 + + +@dataclass +class MultiTurnConfig(BaseConfig): + _mutable_fields = {"max_assistant_turns", "max_user_turns"} + + enable: bool = False + max_assistant_turns: Optional[int] = None + tool_config_path: Optional[str] = None + function_tool_path: Optional[str] = None + max_user_turns: Optional[int] = None + max_parallel_calls: int = 1 + max_tool_response_length: int = 256 + tool_response_truncate_side: str = "middle" + use_inference_chat_template: bool = False + tokenization_sanity_check_mode: str = "strict" + format: str = "hermes" + num_repeat_rollouts: Optional[int] = None + + +@dataclass +class CustomAsyncServerConfig(BaseConfig): + path: Optional[str] = None + name: Optional[str] = None + + +@dataclass +class AgentLoopConfig(BaseConfig): + num_workers: int = 8 + default_agent_loop: str = "single_turn_agent" + agent_loop_config_path: Optional[str] = None + custom_async_server: CustomAsyncServerConfig = field(default_factory=CustomAsyncServerConfig) + # Fully qualified class name for custom AgentLoopManager (e.g., "mypackage.module.MyManager"). + # Security: This class will be dynamically imported via importlib. Only use trusted class paths. + agent_loop_manager_class: Optional[str] = None + + +@dataclass +class TraceConfig(BaseConfig): + project_name: Optional[str] = None + experiment_name: Optional[str] = None + backend: Optional[str] = None + token2text: bool = False + max_samples_per_step_per_worker: Optional[int] = None + + def __post_init__(self): + if self.max_samples_per_step_per_worker is not None and self.max_samples_per_step_per_worker < 0: + raise ValueError("`max_samples_per_step_per_worker` must be a non-negative integer or null.") + + +@dataclass +class ServerConfig(BaseConfig): + """ + Configuration for SGLang server when running in server mode + """ + + timeout: float = 60.0 + max_attempts: int = 3 + retry_delay: float = 2.0 + max_connections: int = 1000 + max_start_wait_time: float = 300.0 + + +@dataclass +class PrometheusConfig(BaseConfig): + """ + Configuration for Prometheus server + """ + + # whether enable prometheus on server mode rollout + enable: bool = False + # Port number that Prometheus listens on, default is 9090 + port: int = 9090 + # Path to Prometheus configuration file + file: str = "/tmp/ray/session_latest/metrics/prometheus/prometheus.yml" + # Specify served_model_name to avoid displaying overly long model paths in Grafana + served_model_name: Optional[str] = None + + +@dataclass +class CheckpointEngineConfig(BaseConfig): + """ + Configuration for checkpoint engine to update weights from trainer to rollout + """ + + # Backend for checkpoint engine: naive, nccl, nixl, hccl + backend: Optional[str] = "naive" + # Bucket size in MB to transfer multiple weights at one time + update_weights_bucket_megabytes: int = 2048 + # Additional keyword arguments for checkpoint engine + engine_kwargs: dict = field(default_factory=dict) + # If set, this Python module is imported on every worker process before the + # backend is instantiated, allowing custom backends to register themselves + # in CheckpointEngineRegistry. + custom_backend_module: Optional[str] = None + + +@dataclass +class RolloutConfig(BaseConfig): + _mutable_fields = { + "max_model_len", + "load_format", + "engine_kwargs", + "prompt_length", + "response_length", + "expert_parallel_size", + "moe_tensor_parallel_size", + } + + name: Optional[str] = MISSING + mode: str = "async" + nnodes: int = 0 + n_gpus_per_node: int = 8 + + temperature: float = 1.0 + top_k: int = -1 + top_p: float = 1.0 + do_sample: bool = True + n: int = 1 + repetition_penalty: float = 1.0 + + # Early termination threshold for multi-turn rollout in sglang. + # Abort remaining requests when (1 - over_sample_rate) * total_requests are completed. + over_sample_rate: float = 0.0 + + prompt_length: int = 512 + response_length: int = 512 + + dtype: str = "bfloat16" + gpu_memory_utilization: float = 0.5 + ignore_eos: bool = False + enforce_eager: bool = True + cudagraph_capture_sizes: Optional[list] = None + free_cache_engine: bool = True + data_parallel_size: int = 1 + expert_parallel_size: int = 1 + tensor_model_parallel_size: int = 2 + pipeline_model_parallel_size: int = 1 + moe_tensor_parallel_size: int = 1 + max_num_batched_tokens: int = 8192 + logprobs_mode: Optional[str] = "processed_logprobs" + scheduling_policy: Optional[str] = "fcfs" + + # TODO: enable train_kwargs + # train_sampling_config: SamplingConfig = field(default_factory=SamplingConfig) + + val_kwargs: SamplingConfig = field(default_factory=SamplingConfig) + + max_model_len: Optional[int] = None + max_num_seqs: int = 1024 + + # note that the logprob computation should belong to the actor + log_prob_micro_batch_size: Optional[int] = None + log_prob_micro_batch_size_per_gpu: Optional[int] = None + log_prob_use_dynamic_bsz: bool = False + log_prob_max_token_len_per_gpu: int = 16384 + + disable_log_stats: bool = True + + multi_stage_wake_up: bool = False + engine_kwargs: dict = field(default_factory=dict) + + calculate_log_probs: bool = False + + agent: AgentLoopConfig = field(default_factory=AgentLoopConfig) + + trace: TraceConfig = field(default_factory=TraceConfig) + + multi_turn: MultiTurnConfig = field(default_factory=MultiTurnConfig) + + # Server configuration for sglang server mode + server: ServerConfig = field(default_factory=ServerConfig) + + # Use Prometheus to collect and monitor rollout statistics + prometheus: PrometheusConfig = field(default_factory=PrometheusConfig) + + # Extension point for custom configurations + custom: Optional[dict] = None + + # Fully qualified class name for a custom CheckpointEngineManager. When set, the trainer + # loads this class instead of the built-in CheckpointEngineManager. + checkpoint_manager_class: Optional[str] = None + + # Checkpoint Engine config for update weights from trainer to rollout + checkpoint_engine: CheckpointEngineConfig = field(default_factory=CheckpointEngineConfig) + + # Rollout skip config (load/dump rollout data) + skip: SkipConfig = field(default_factory=SkipConfig) + + profiler: Optional[ProfilerConfig] = None + + enable_chunked_prefill: bool = True + + enable_prefix_caching: bool = True + + load_format: str = "dummy" + + layered_summon: bool = False + + layer_name_map: dict = field(default_factory=dict) + + sglang_engine_mode: str = "local" + + limit_images: Optional[int] = None + + skip_tokenizer_init: bool = False + + quantization: Optional[str] = None + + quantization_config_file: Optional[str] = None + + enable_rollout_routing_replay: bool = False + + enable_sleep_mode: bool = True + + mtp: MtpConfig = field(default_factory=MtpConfig) + + qat: Optional[dict] = None + + disaggregation: DisaggregationConfig = field(default_factory=DisaggregationConfig) + + def __post_init__(self): + """Validate the rollout config""" + # Deprecation warning for mode field - only async mode is supported + if self.mode == "sync": + raise ValueError( + "Rollout mode 'sync' has been removed. Please set " + "`actor_rollout_ref.rollout.mode=async` or remove the mode setting entirely." + ) + if self.mode != "async": + warnings.warn( + f"Unknown rollout mode '{self.mode}'. Only 'async' mode is supported. " + "The 'mode' field is deprecated and will be removed in a future version.", + DeprecationWarning, + stacklevel=2, + ) + + if self.name != "trtllm" and self.expert_parallel_size > 1: + assert self.expert_parallel_size == (self.tensor_model_parallel_size * self.data_parallel_size), ( + "expert_parallel_size must be equal to tensor_model_parallel_size * data_parallel_size" + ) + + if self.moe_tensor_parallel_size is not None and self.moe_tensor_parallel_size > 1: + assert self.name == "trtllm", "moe_tensor_parallel_size is only supported for trtllm" + + if self.name == "trtllm": + # If either expert_parallel_size or moe_tensor_parallel_size is at default 1, + # convert to None so TensorRT-LLM treats it as unspecified. + # When both unspecified: moe_ep_size=1, moe_tp_size=moe_world_size (no EP, all TP). + # When only one set: the other is auto-derived from tensor_model_parallel_size. + if self.expert_parallel_size is not None and self.expert_parallel_size == 1: + self.expert_parallel_size = None + if self.moe_tensor_parallel_size is not None and self.moe_tensor_parallel_size == 1: + self.moe_tensor_parallel_size = None + if self.expert_parallel_size is not None and self.moe_tensor_parallel_size is not None: + assert self.moe_tensor_parallel_size * self.expert_parallel_size == self.tensor_model_parallel_size, ( + "moe_tensor_parallel_size * expert_parallel_size must equal tensor_model_parallel_size " + f"(got {self.moe_tensor_parallel_size} * {self.expert_parallel_size} = " + f"{self.moe_tensor_parallel_size * self.expert_parallel_size}, " + f"tensor_model_parallel_size={self.tensor_model_parallel_size})" + ) + + if self.pipeline_model_parallel_size > 1: + if self.name == "vllm" or self.name == "sglang" or self.name == "trtllm": + raise NotImplementedError( + f"Current rollout {self.name=} not implemented pipeline_model_parallel_size > 1 yet." + ) + + # Hydra passes this as dict/DictConfig; coerce to dataclass so + # downstream .enabled etc. work. BaseConfig is frozen, hence object.__setattr__. + if isinstance(self.disaggregation, dict): + object.__setattr__(self, "disaggregation", DisaggregationConfig(**self.disaggregation)) + elif not isinstance(self.disaggregation, DisaggregationConfig): + from omegaconf import DictConfig, OmegaConf + + if not isinstance(self.disaggregation, DictConfig): + raise TypeError( + f"rollout.disaggregation must be dict, DictConfig, or DisaggregationConfig; " + f"got {type(self.disaggregation).__name__}." + ) + object.__setattr__( + self, + "disaggregation", + DisaggregationConfig(**OmegaConf.to_container(self.disaggregation, resolve=True)), + ) + + if self.disaggregation.enabled and self.name != "sglang": + raise ValueError( + f"rollout.disaggregation.enabled=True is currently only supported with " + f"rollout.name='sglang'; got {self.name!r}. (vLLM PD is a tracked follow-up.)" + ) diff --git a/verl/verl/workers/engine/__init__.py b/verl/verl/workers/engine/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..eceba017f301f2c8720c7c9675248a8969c5a45e --- /dev/null +++ b/verl/verl/workers/engine/__init__.py @@ -0,0 +1,64 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from .base import BaseEngine, EngineRegistry +from .fsdp import FSDPEngine, FSDPEngineWithLMHead + +__all__ = [ + "BaseEngine", + "EngineRegistry", + "FSDPEngine", + "FSDPEngineWithLMHead", +] + +try: + from .torchtitan import TorchTitanEngine, TorchTitanEngineWithLMHead + + __all__ += ["TorchTitanEngine", "TorchTitanEngineWithLMHead"] +except ImportError: + TorchTitanEngine = None + TorchTitanEngineWithLMHead = None + +try: + from .veomni import VeOmniEngine, VeOmniEngineWithLMHead + + __all__ += ["VeOmniEngine", "VeOmniEngineWithLMHead"] +except ImportError: + VeOmniEngine = None + VeOmniEngineWithLMHead = None + +try: + from .automodel import AutomodelEngine, AutomodelEngineWithLMHead + + __all__ += ["AutomodelEngine", "AutomodelEngineWithLMHead"] +except ImportError: + AutomodelEngine = None + AutomodelEngineWithLMHead = None + +# Mindspeed must be imported before Megatron to ensure the related monkey patches take effect as expected +try: + from .mindspeed import MindspeedEngineWithLMHead, MindspeedEngineWithValueHead, MindSpeedLLMEngineWithLMHead + + __all__ += ["MindspeedEngineWithLMHead", "MindspeedEngineWithValueHead", "MindSpeedLLMEngineWithLMHead"] +except ImportError: + MindspeedEngineWithLMHead = None + MindspeedEngineWithValueHead = None + MindSpeedLLMEngineWithLMHead = None + +try: + from .megatron import MegatronEngine, MegatronEngineWithLMHead, MegatronEngineWithValueHead + + __all__ += ["MegatronEngine", "MegatronEngineWithLMHead", "MegatronEngineWithValueHead"] +except ImportError: + MegatronEngine = None + MegatronEngineWithLMHead = None diff --git a/verl/verl/workers/engine/__pycache__/__init__.cpython-312.pyc b/verl/verl/workers/engine/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a7e39d4f2a28757767185b9a787d39c4db23c4c0 Binary files /dev/null and b/verl/verl/workers/engine/__pycache__/__init__.cpython-312.pyc differ diff --git a/verl/verl/workers/engine/__pycache__/base.cpython-312.pyc b/verl/verl/workers/engine/__pycache__/base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8db40e6020227498ca865b102a8e112539f90180 Binary files /dev/null and b/verl/verl/workers/engine/__pycache__/base.cpython-312.pyc differ diff --git a/verl/verl/workers/engine/__pycache__/utils.cpython-312.pyc b/verl/verl/workers/engine/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6c2d37f351f51c0a3871d4769b52066011bfb27f Binary files /dev/null and b/verl/verl/workers/engine/__pycache__/utils.cpython-312.pyc differ diff --git a/verl/verl/workers/engine/automodel/__init__.py b/verl/verl/workers/engine/automodel/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a839342706ba5be7fe7f748d2db291ca467ae4c3 --- /dev/null +++ b/verl/verl/workers/engine/automodel/__init__.py @@ -0,0 +1,20 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .transformer_impl import AutomodelEngine, AutomodelEngineWithLMHead + +__all__ = [ + "AutomodelEngine", + "AutomodelEngineWithLMHead", +] diff --git a/verl/verl/workers/engine/automodel/__pycache__/__init__.cpython-312.pyc b/verl/verl/workers/engine/automodel/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..60c5abde85c964e81c5a04a12f9daeb424fbbd89 Binary files /dev/null and b/verl/verl/workers/engine/automodel/__pycache__/__init__.cpython-312.pyc differ diff --git a/verl/verl/workers/engine/automodel/__pycache__/transformer_impl.cpython-312.pyc b/verl/verl/workers/engine/automodel/__pycache__/transformer_impl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7fce232c23e87777dff0c28b7da4fedb303dbc60 Binary files /dev/null and b/verl/verl/workers/engine/automodel/__pycache__/transformer_impl.cpython-312.pyc differ diff --git a/verl/verl/workers/engine/automodel/transformer_impl.py b/verl/verl/workers/engine/automodel/transformer_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..fc71384a323566e9d6c51644638c7babae9c3c61 --- /dev/null +++ b/verl/verl/workers/engine/automodel/transformer_impl.py @@ -0,0 +1,713 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Automodel (nemo_automodel) engine for verl SFT training. + +This engine delegates model building, parallelization, optimizer sharding, +LR scheduling, gradient clipping, and checkpointing to Automodel's +infrastructure while using verl's training loop, data pipeline, and loss function. +""" + +import gc +import logging +import os +from contextlib import nullcontext +from typing import Any, Callable, Optional + +import torch +import torch.distributed +from huggingface_hub.constants import HF_HUB_CACHE +from nemo_automodel.components.checkpoint.checkpointing import Checkpointer, CheckpointingConfig +from nemo_automodel.components.optim.scheduler import OptimizerParamScheduler +from nemo_automodel.components.training.utils import ( + prepare_for_final_backward, + prepare_for_grad_accumulation, + scale_grads_and_clip_grad_norm, +) +from tensordict import TensorDict +from torch.distributed.tensor import DTensor + +import verl.utils.torch_functional as verl_F +from verl.trainer.config import CheckpointConfig +from verl.utils import tensordict_utils as tu +from verl.utils.dataset.dataset_utils import DatasetPadMode +from verl.utils.debug import log_gpu_memory_usage +from verl.utils.device import get_device_id, get_device_name +from verl.utils.model import convert_weight_keys, extract_multi_modal_inputs +from verl.utils.torch_functional import logprobs_from_logits +from verl.workers.config import AutomodelEngineConfig, AutomodelOptimizerConfig, HFModelConfig + +from ..base import BaseEngine, BaseEngineCtx, EngineRegistry +from ..utils import enable_full_determinism, postprocess_batch_func, prepare_micro_batches +from .utils import ( + build_automodel_model, + build_distributed_config_from_engine_config, + get_dp_group_size, + get_dp_rank, + get_pp_rank, + get_tp_rank, + load_automodel_model_to_gpu, + load_automodel_optimizer, + maybe_fully_shard_optimizer, + offload_automodel_model_to_cpu, + offload_automodel_optimizer, +) + +logger = logging.getLogger(__file__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + + +class AutomodelEngine(BaseEngine): + """Engine implementation using Automodel for distributed training.""" + + def __init__( + self, + model_config: HFModelConfig, + engine_config: AutomodelEngineConfig, + optimizer_config: AutomodelOptimizerConfig, + checkpoint_config: CheckpointConfig, + **kwargs, + ): + super().__init__() + + self.model_config = model_config + self.engine_config = engine_config + self.optimizer_config = optimizer_config + self.checkpoint_config = checkpoint_config + + self.mode = None + self.rank = torch.distributed.get_rank() + + # Apply compatibility patches early in the process + from nemo_automodel._transformers.utils import apply_cache_compatibility_patches + from nemo_automodel.shared.te_patches import apply_te_patches + + apply_cache_compatibility_patches() + apply_te_patches() + + world_size = torch.distributed.get_world_size() + self.distributed_config, self.device_mesh, self.moe_mesh = build_distributed_config_from_engine_config( + self.engine_config, world_size + ) + + if self.engine_config.full_determinism: + enable_full_determinism(seed=self.engine_config.seed) + + self._is_offload_param = self.engine_config.param_offload + self._is_offload_optimizer = self.engine_config.optimizer_offload + + if self.engine_config.entropy_from_logits_with_chunking: + entropy_from_logits = verl_F.entropy_from_logits_with_chunking + else: + entropy_from_logits = verl_F.entropy_from_logits + + self.compute_entropy_from_logits = ( + torch.compile(entropy_from_logits, dynamic=True) + if self.engine_config.use_torch_compile + else entropy_from_logits + ) + + @property + def is_param_offload_enabled(self) -> bool: + return self._is_offload_param + + @property + def is_optimizer_offload_enabled(self) -> bool: + return self._is_offload_optimizer + + def initialize(self): + """Build the model, optimizer, LR scheduler, and checkpointer using Automodel infrastructure.""" + self.module = build_automodel_model( + self.model_config, self.engine_config, self.distributed_config, self.device_mesh, self.moe_mesh + ) + log_gpu_memory_usage("After Automodel model build", logger=logger) + + if not self.engine_config.forward_only: + self.optimizer = self._build_optimizer(self.module) + # maybe shard optimizer for MegatronFSDP + maybe_fully_shard_optimizer(self.module, self.optimizer, self.distributed_config) + self.lr_scheduler = self._build_lr_scheduler(self.optimizer) + else: + self.optimizer = None + self.lr_scheduler = None + self._build_checkpointer() + + self.to( + device="cpu", + model=self._is_offload_param, + optimizer=self._is_offload_optimizer, + grad=self._is_offload_param, + ) + + log_gpu_memory_usage("After offload model/optimizer/grad during init", logger=logger) + torch.cuda.empty_cache() + + def _build_optimizer(self, module): + """Build optimizer via Automodel's build_optimizer.""" + from nemo_automodel.components.config.loader import ConfigNode + from nemo_automodel.recipes.llm.train_ft import build_optimizer as automodel_build_optimizer + + config = self.optimizer_config + + opt_dict = { + "_target_": f"{config.optimizer_impl}.{config.optimizer}", + "lr": config.lr, + "weight_decay": config.weight_decay, + "eps": config.eps, + "betas": list(config.betas), + } + + if config.master_weights: + opt_dict["master_weights"] = config.master_weights + if config.store_param_remainders: + opt_dict["store_param_remainders"] = config.store_param_remainders + + _short_to_torch = {"bf16": "torch.bfloat16", "fp32": "torch.float32", "fp16": "torch.float16"} + for attr in ("exp_avg_dtype", "exp_avg_sq_dtype", "master_weight_dtype"): + val = getattr(config, attr, None) + if val is not None: + opt_dict[attr] = _short_to_torch.get(val, val) + + if config.override_optimizer_config: + opt_dict.update(config.override_optimizer_config) + + cfg_opt = ConfigNode(opt_dict) + optimizers = automodel_build_optimizer(module, cfg_opt, self.distributed_config, self.device_mesh) + assert len(optimizers) == 1, f"Expected 1 optimizer, got {len(optimizers)}" + return optimizers[0] + + def _build_lr_scheduler(self, optimizer): + cfg = self.optimizer_config + total_steps = cfg.total_training_steps + num_warmup_steps = cfg.lr_warmup_steps + + if num_warmup_steps <= 0: + num_warmup_steps = int(cfg.lr_warmup_steps_ratio * total_steps) + + base_lr = cfg.lr + init_lr_ratio = cfg.init_lr_ratio if cfg.init_lr_ratio is not None else 0.1 + min_lr_ratio = cfg.min_lr_ratio if cfg.min_lr_ratio is not None else 0.01 + + if self.rank == 0: + print( + f"Automodel LR Scheduler: total_steps={total_steps}, warmup={num_warmup_steps}, " + f"decay_style={cfg.lr_scheduler_type}, init_lr={base_lr * init_lr_ratio:.2e}, " + f"max_lr={base_lr:.2e}, min_lr={base_lr * min_lr_ratio:.2e}" + ) + + scheduler = OptimizerParamScheduler( + optimizer=optimizer, + init_lr=base_lr * init_lr_ratio, + max_lr=base_lr, + min_lr=base_lr * min_lr_ratio, + lr_warmup_steps=num_warmup_steps, + lr_decay_steps=total_steps, + lr_decay_style=cfg.lr_scheduler_type, + start_wd=cfg.weight_decay, + end_wd=cfg.weight_decay, + wd_incr_steps=total_steps, + wd_incr_style=getattr(cfg, "wd_incr_style", "constant"), + ) + return scheduler + + def forward_backward_batch(self, data: TensorDict, loss_function: Callable, forward_only=False) -> Any: + batch_num_tokens = data["loss_mask"].sum().to(get_device_id()) + torch.distributed.all_reduce( + batch_num_tokens, op=torch.distributed.ReduceOp.SUM, group=self.get_data_parallel_group() + ) + tu.assign_non_tensor(data, batch_num_tokens=batch_num_tokens.item()) + tu.assign_non_tensor(data, dp_size=self.get_data_parallel_size()) + + micro_batches, indices = prepare_micro_batches( + data=data, dp_group=self.get_data_parallel_group(), same_micro_num_in_dp=True + ) + + output_lst = [] + ctx = torch.no_grad() if forward_only else nullcontext() + + if not forward_only: + prepare_for_grad_accumulation([self.module]) + + # Set MoE aux loss backward scale to counteract FSDP's gradient allreduce. + if self.engine_config.ep_size > 1: + from nemo_automodel.components.moe.megatron.moe_utils import MoEAuxLossAutoScaler + + MoEAuxLossAutoScaler.main_loss_backward_scale = torch.tensor( + float(get_dp_group_size(self.device_mesh, include_cp=True)) + ) + + num_micro_batches = len(micro_batches) + for i, micro_batch in enumerate(micro_batches): + # Signal final backward for MoE + if not forward_only and i == num_micro_batches - 1: + prepare_for_final_backward([self.module]) + + with ctx: + loss, meta_info = self.forward_step(micro_batch, loss_function=loss_function, forward_only=forward_only) + if not forward_only: + loss.backward() + output_lst.append(meta_info) + + return postprocess_batch_func(output_lst=output_lst, indices=indices, data=data) + + def forward_step(self, micro_batch: TensorDict, loss_function, forward_only): + raise NotImplementedError("forward_step must be implemented in subclass") + + def optimizer_zero_grad(self): + self.optimizer.zero_grad() + + def optimizer_step(self): + grad_norm = scale_grads_and_clip_grad_norm( + max_grad_norm=self.optimizer_config.clip_grad, + model_parts=[self.module], + norm_type=2.0, + pp_enabled=False, + device_mesh=self.device_mesh, + moe_mesh=self.moe_mesh, + ep_axis_name="ep" if self.moe_mesh is not None and "ep" in self.moe_mesh.mesh_dim_names else None, + pp_axis_name=None, + foreach=True, + num_label_tokens=None, + dp_group_size=get_dp_group_size(self.device_mesh, include_cp=True), + ) + + if isinstance(grad_norm, torch.Tensor): + grad_norm_val = grad_norm.item() + else: + grad_norm_val = float(grad_norm) + + # If grad_norm is not finite, skip the update + if not torch.isfinite(torch.tensor(grad_norm_val)): + print(f"WARN: grad_norm is not finite: {grad_norm_val}") + self.optimizer.zero_grad() + else: + self.optimizer.step() + if hasattr(self.module, "update_moe_gate_bias"): + self.module.update_moe_gate_bias() + + return grad_norm_val + + def lr_scheduler_step(self): + """Step Automodel's OptimizerParamScheduler and return current LR.""" + self.lr_scheduler.step(increment=1) + lr = self.optimizer.param_groups[0]["lr"] + return lr + + def get_data_parallel_rank(self): + if self.device_mesh is not None: + return self.device_mesh.get_local_rank("dp") + return torch.distributed.get_rank() + + def get_data_parallel_size(self): + if self.device_mesh is not None: + return self.device_mesh["dp"].size() + return torch.distributed.get_world_size() + + def get_data_parallel_group(self): + if self.device_mesh is not None: + return self.device_mesh.get_group(mesh_dim="dp") + return torch.distributed.group.WORLD + + def is_mp_src_rank_with_outputs(self): + if self.device_mesh is not None and "tp" in self.device_mesh.mesh_dim_names: + if self.device_mesh["tp"].size() > 1: + return self.device_mesh.get_local_rank("tp") == 0 + return True + + def train_mode(self, **kwargs): + return AutomodelTrainModeCtx(self, **kwargs) + + def eval_mode(self, **kwargs): + return AutomodelEvalModeCtx(self, **kwargs) + + def to(self, device: str, model: bool = True, optimizer: bool = True, grad: bool = True): + super().to(device=device, model=model, optimizer=optimizer, grad=grad) + + if self.engine_config.forward_only: + return + + device_name = get_device_name() + assert device in (device_name, "cpu") + + if device == device_name: + if model: + load_automodel_model_to_gpu(self.module) + if optimizer and self.optimizer is not None: + load_automodel_optimizer(self.optimizer, get_device_id()) + gc.collect() + elif device == "cpu": + if model: + offload_automodel_model_to_cpu(self.module) + if optimizer and self.optimizer is not None: + offload_automodel_optimizer(self.optimizer) + else: + raise ValueError(f"Invalid device type: {device}") + + def _build_checkpointer(self): + ckpt_config = CheckpointingConfig( + enabled=True, + checkpoint_dir="checkpoints/", + model_save_format="safetensors", + model_cache_dir=HF_HUB_CACHE, + model_repo_id=self.model_config.path, + save_consolidated=True, + is_peft=False, + ) + self.checkpointer = Checkpointer( + config=ckpt_config, + dp_rank=get_dp_rank(self.device_mesh, include_cp=True), + tp_rank=get_tp_rank(self.device_mesh), + pp_rank=get_pp_rank(self.device_mesh), + moe_mesh=self.moe_mesh, + ) + + def save_checkpoint( + self, + local_path: str, + hdfs_path: Optional[str] = None, + global_step: int = 0, + max_ckpt_to_keep: Optional[int] = None, + **kwargs, + ) -> None: + """Save model, optimizer, and LR scheduler using Automodel's Checkpointer.""" + origin_module_device = next(self.module.parameters()).device.type + if self._is_offload_param or origin_module_device == "cpu": + load_automodel_model_to_gpu(self.module) + + # Save model weights + self.checkpointer.save_model(self.module, local_path) + + # Save optimizer and LR scheduler state + if self.optimizer is not None: + scheduler_list = [self.lr_scheduler] if self.lr_scheduler is not None else None + self.checkpointer.save_optimizer(self.optimizer, self.module, local_path, scheduler=scheduler_list) + + torch.distributed.barrier() + if self._is_offload_param: + offload_automodel_model_to_cpu(self.module) + + def load_checkpoint( + self, local_path: str, hdfs_path: Optional[str] = None, del_local_after_load: int = True, **kwargs + ) -> None: + """Load model, optimizer, and LR scheduler using Automodel's Checkpointer.""" + if self._is_offload_param: + load_automodel_model_to_gpu(self.module) + + model_path = os.path.join(local_path, "model") + if not os.path.isdir(model_path): + model_path = local_path + self.checkpointer.load_model(self.module, model_path) + + if self.optimizer is not None: + scheduler_list = [self.lr_scheduler] if self.lr_scheduler is not None else None + self.checkpointer.load_optimizer(self.optimizer, self.module, local_path, scheduler=scheduler_list) + + torch.distributed.barrier() + if self._is_offload_param: + offload_automodel_model_to_cpu(self.module) + + if self._is_offload_optimizer and self.optimizer is not None: + offload_automodel_optimizer(self.optimizer) + + def get_per_tensor_param(self, **kwargs): + load_automodel_model_to_gpu(self.module) + + params = self.module.state_dict() + params = convert_weight_keys(params, getattr(self.module, "_fsdp_wrapped_module", self.module)) + + if self._is_offload_param: + offload_automodel_model_to_cpu(self.module) + + def param_generator(): + for name, param in params.items(): + unsharded_tensor = param.full_tensor() if isinstance(param, DTensor) else param + yield name, unsharded_tensor + + return param_generator(), None + + +class AutomodelEvalModeCtx(BaseEngineCtx): + def __init__(self, engine: AutomodelEngine, **kwargs): + super().__init__(engine=engine, mode="eval", **kwargs) + + def __enter__(self): + assert isinstance(self.engine, AutomodelEngine) + super().__enter__() + self.engine.module.eval() + + def __exit__(self, exc_type, exc_value, traceback): + assert isinstance(self.engine, AutomodelEngine) + # Reshard the root FSDP module + if hasattr(self.engine.module, "reshard"): + self.engine.module.reshard() + super().__exit__(exc_type, exc_value, traceback) + + +class AutomodelTrainModeCtx(BaseEngineCtx): + def __init__(self, engine: AutomodelEngine, **kwargs): + super().__init__(engine=engine, mode="train", **kwargs) + + def __enter__(self): + assert isinstance(self.engine, AutomodelEngine) + super().__enter__() + self.engine.module.train() + + def __exit__(self, exc_type, exc_value, traceback): + assert isinstance(self.engine, AutomodelEngine) + self.engine.optimizer_zero_grad() + super().__exit__(exc_type, exc_value, traceback) + + +@EngineRegistry.register(model_type="language_model", backend=["automodel"], device=["cuda"]) +class AutomodelEngineWithLMHead(AutomodelEngine): + """Automodel engine for language model with LM head training.""" + + def prepare_model_inputs(self, micro_batch: TensorDict): + use_remove_padding = tu.get_non_tensor_data(data=micro_batch, key="use_remove_padding", default=True) + pad_mode = tu.get_non_tensor_data(data=micro_batch, key="pad_mode", default=DatasetPadMode.NO_PADDING) + use_fused_kernels = tu.get_non_tensor_data(data=micro_batch, key="use_fused_kernels", default=False) + temperature = micro_batch["temperature"] + temperature_item = temperature + if use_fused_kernels: + assert not isinstance(temperature, torch.Tensor), ( + "use_fused_kernels does not support per sample temperature yet" + ) + assert pad_mode == DatasetPadMode.NO_PADDING, f"pad_mode {pad_mode} not supported" + + multi_modal_inputs = extract_multi_modal_inputs(micro_batch.get("multi_modal_inputs", [])) + input_ids = micro_batch["input_ids"] + position_ids = micro_batch["position_ids"] + + if not isinstance(temperature, torch.Tensor): + temperature = torch.tensor([temperature] * input_ids.shape[0], device=input_ids.device) + + temperature = temperature.to(torch.float32) + assert temperature.shape[0] == input_ids.shape[0] + + output_args = {} + + if use_remove_padding: + temperature_rmpad = verl_F.expand_as_nested(temperature, input_ids).values() + temperature_rmpad = temperature_rmpad.unsqueeze(0) + + if pad_mode == DatasetPadMode.NO_PADDING: + input_ids_rmpad = input_ids.values().unsqueeze(0) + if position_ids.dim() == 3: + position_ids_rmpad = position_ids.values().unsqueeze(1) + else: + position_ids_rmpad = position_ids.values().unsqueeze(0) + else: + raise NotImplementedError(f"pad_mode {pad_mode} not implemented") + + input_ids_rmpad_rolled = torch.roll(input_ids_rmpad, shifts=-1, dims=1) + + input_ids_rmpad_rolled = input_ids_rmpad_rolled.squeeze(0) + temperature_rmpad = temperature_rmpad.squeeze(0) + output_args["input_ids_rmpad_rolled"] = input_ids_rmpad_rolled + output_args["temperature_rmpad"] = temperature_rmpad + + model_inputs = { + "input_ids": input_ids_rmpad, + "attention_mask": None, + "position_ids": position_ids_rmpad, + } + + # For TE attention backend, pass cu_seqlens + if self.engine_config.attn_implementation == "te": + cu_seqlens = input_ids.offsets().to(torch.int32) + max_seqlen = cu_seqlens.diff().max().item() + model_inputs["qkv_format"] = "thd" + model_inputs["cu_seqlens"] = cu_seqlens.unsqueeze(0) + model_inputs["max_seqlen"] = max_seqlen + + else: + if pad_mode == DatasetPadMode.NO_PADDING: + input_ids = micro_batch["input_ids"] + position_ids = micro_batch["position_ids"] + loss_mask = micro_batch["loss_mask"] + + pad_token_id = tu.get_non_tensor_data(data=micro_batch, key="pad_token_id", default=0) + batch_size = micro_batch.batch_size[0] + seq_len_effective = input_ids.offsets().diff() + max_seq_len = max(seq_len_effective) + + input_ids_rmpad_rolled = torch.roll(input_ids.values(), shifts=-1, dims=0) + output_args["input_ids_rmpad_rolled"] = input_ids_rmpad_rolled + output_args["temperature"] = temperature + + input_ids = torch.nested.to_padded_tensor( + input_ids, padding=pad_token_id, output_size=(batch_size, max_seq_len) + ) + + if position_ids.dim() == 3: + position_ids = torch.nested.to_padded_tensor( + position_ids, padding=0, output_size=(batch_size, 4, max_seq_len) + ).transpose(0, 1) + else: + position_ids = torch.nested.to_padded_tensor( + position_ids, padding=0, output_size=(batch_size, max_seq_len) + ) + + attention_mask_list = [torch.ones_like(t, dtype=torch.int32) for t in loss_mask] + attention_mask = torch.nested.as_nested_tensor(attention_mask_list, layout=torch.jagged) + attention_mask = torch.nested.to_padded_tensor( + attention_mask, padding=0, output_size=(batch_size, max_seq_len) + ) + + model_inputs = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "position_ids": position_ids, + } + + else: + raise NotImplementedError(f"pad_mode {pad_mode} not implemented") + + extra_args = {} + if use_fused_kernels: + extra_args["temperature"] = temperature_item + extra_args["return_dict"] = True + + model_inputs.update(multi_modal_inputs) + model_inputs.update(extra_args) + + return model_inputs, output_args + + def prepare_model_outputs(self, output, output_args, micro_batch: TensorDict): + use_remove_padding = tu.get_non_tensor_data(data=micro_batch, key="use_remove_padding", default=True) + pad_mode = tu.get_non_tensor_data(data=micro_batch, key="pad_mode", default=DatasetPadMode.NO_PADDING) + use_fused_kernels = tu.get_non_tensor_data(data=micro_batch, key="use_fused_kernels", default=False) + calculate_entropy = tu.get_non_tensor_data(data=micro_batch, key="calculate_entropy", default=False) + + if isinstance(output, torch.Tensor): + from types import SimpleNamespace + + output = SimpleNamespace(logits=output) + + model_output = {} + input_ids = micro_batch["input_ids"] + + if use_remove_padding: + input_ids_rmpad_rolled = output_args["input_ids_rmpad_rolled"] + temperature_rmpad = output_args["temperature_rmpad"] + + if use_fused_kernels: + log_probs = output.log_probs.squeeze(0) + entropy_rmpad = output.entropy.squeeze(0) + else: + logits_rmpad = output.logits.squeeze(0) + # With TP, logits are DTensors sharded on vocab dim; gather for log_softmax. + if isinstance(logits_rmpad, DTensor): + logits_rmpad = logits_rmpad.full_tensor() + logits_rmpad = logits_rmpad / temperature_rmpad.clamp(min=1e-8).unsqueeze(-1).to(logits_rmpad.dtype) + + inplace_backward = True + if calculate_entropy: + inplace_backward = False + log_probs = logprobs_from_logits( + logits=logits_rmpad, + labels=input_ids_rmpad_rolled, + inplace_backward=inplace_backward, + ) + + if calculate_entropy: + if not self.engine_config.entropy_checkpointing: + entropy_rmpad = self.compute_entropy_from_logits(logits_rmpad) + else: + entropy_rmpad = torch.utils.checkpoint.checkpoint( + self.compute_entropy_from_logits, logits_rmpad + ) + + if pad_mode == DatasetPadMode.NO_PADDING: + cu_seqlens = input_ids.offsets() + log_probs = torch.nested.nested_tensor_from_jagged(log_probs, cu_seqlens) + if calculate_entropy: + entropy = torch.nested.nested_tensor_from_jagged(entropy_rmpad, cu_seqlens) + else: + raise NotImplementedError(f"pad_mode {pad_mode} not implemented") + + else: + response_length = tu.get_non_tensor_data(data=micro_batch, key="max_response_length", default=1024) + if use_fused_kernels: + log_probs = output.log_probs[:, -response_length - 1 : -1] + entropy = output.entropy[:, -response_length - 1 : -1] + else: + logits = output.logits + # With TP, logits are DTensors sharded on vocab dim; gather for log_softmax. + if isinstance(logits, DTensor): + logits = logits.full_tensor() + temperature = output_args["temperature"] + temperature = temperature.unsqueeze(-1).unsqueeze(-1) + logits = logits / temperature.clamp(min=1e-8).to(logits.dtype) + + if calculate_entropy: + if not self.engine_config.entropy_checkpointing: + entropy = verl_F.entropy_from_logits(logits) + else: + entropy = torch.utils.checkpoint.checkpoint(verl_F.entropy_from_logits, logits) + + if pad_mode == DatasetPadMode.NO_PADDING: + cu_seqlens = input_ids.offsets() + seq_lengths = cu_seqlens.diff() + starts = torch.zeros_like(seq_lengths, dtype=torch.int64) + logits = torch.nested.narrow(logits, 1, starts, seq_lengths, layout=torch.jagged) + logits_rmpad = torch.cat([t for t in logits.unbind()]) + input_ids_rmpad_rolled = output_args["input_ids_rmpad_rolled"] + log_probs = logprobs_from_logits(logits=logits_rmpad, labels=input_ids_rmpad_rolled) + log_probs = torch.nested.nested_tensor_from_jagged(log_probs, cu_seqlens) + if calculate_entropy: + entropy = torch.nested.narrow(entropy, 1, starts, seq_lengths, layout=torch.jagged) + entropy_rmpad = torch.cat([t for t in entropy.unbind()]) + entropy = torch.nested.nested_tensor_from_jagged(entropy_rmpad, cu_seqlens) + else: + raise NotImplementedError(f"pad_mode {pad_mode} not implemented") + + model_output["log_probs"] = log_probs + if calculate_entropy: + model_output["entropy"] = entropy + + return model_output + + def forward_step(self, micro_batch: TensorDict, loss_function, forward_only): + """Run forward pass, compute loss, and return outputs.""" + device_name = get_device_name() + micro_batch = micro_batch.to(get_device_id()) + model_inputs, output_args = self.prepare_model_inputs(micro_batch=micro_batch) + + with torch.autocast(device_type=device_name, dtype=torch.bfloat16): + raw_output = self.module( + **model_inputs, + use_cache=False, + ) + + model_output = self.prepare_model_outputs( + output=raw_output, output_args=output_args, micro_batch=micro_batch + ) + + if loss_function is not None: + loss, metrics = loss_function( + model_output=model_output, data=micro_batch, dp_group=self.get_data_parallel_group() + ) + else: + assert forward_only, "forward_only must be True when loss_function is None" + loss = torch.tensor(1.0, device=device_name) + metrics = {} + + output = { + "model_output": model_output, + "loss": loss.detach().item(), + "metrics": metrics, + } + + return loss, output diff --git a/verl/verl/workers/engine/automodel/utils.py b/verl/verl/workers/engine/automodel/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c10cf9a2db2bcd67473a78e5df0ce3f3570984fb --- /dev/null +++ b/verl/verl/workers/engine/automodel/utils.py @@ -0,0 +1,250 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Utility functions for the Automodel engine integration.""" + +import torch +import torch.distributed + +from verl.utils.device import get_device_id, get_torch_device + + +def get_dp_rank(device_mesh, include_cp=False): + """Get data-parallel rank from device mesh.""" + if device_mesh is None: + return 0 + if include_cp and "cp" in device_mesh.mesh_dim_names and device_mesh["cp"].size() > 1: + return device_mesh.get_local_rank("dp_cp") + return device_mesh.get_local_rank("dp") + + +def get_tp_rank(device_mesh): + """Get tensor-parallel rank from device mesh.""" + if device_mesh is None or "tp" not in device_mesh.mesh_dim_names or device_mesh["tp"].size() == 1: + return 0 + return device_mesh.get_local_rank("tp") + + +def get_pp_rank(device_mesh): + """Get pipeline-parallel rank from device mesh.""" + if device_mesh is None or "pp" not in device_mesh.mesh_dim_names or device_mesh["pp"].size() == 1: + return 0 + return device_mesh.get_local_rank("pp") + + +def get_dp_group_size(device_mesh, include_cp=False): + """Get data-parallel group size from device mesh.""" + if device_mesh is None: + return torch.distributed.get_world_size() + if include_cp and "cp" in device_mesh.mesh_dim_names and device_mesh["cp"].size() > 1: + return device_mesh["dp_cp"].size() + if "dp" in device_mesh.mesh_dim_names: + return device_mesh["dp"].size() + return torch.distributed.get_world_size() + + +def maybe_fully_shard_optimizer(model, optimizer, distributed_config): + """Call fully_shard_optimizer for MegatronFSDP strategy.""" + from nemo_automodel.components.distributed.config import MegatronFSDPConfig + + if isinstance(distributed_config, MegatronFSDPConfig) and torch.distributed.get_world_size() > 1: + from megatron_fsdp.fully_shard import fully_shard_optimizer + + fully_shard_optimizer(model, optimizer) + + +def build_distributed_config_from_engine_config(engine_config, world_size): + """Build v5 distributed config, device_mesh, and moe_mesh from engine config. + + Args: + engine_config: AutomodelEngineConfig instance. + world_size: Total number of processes in the job. + + Returns: + Tuple of (distributed_config, device_mesh, moe_mesh). + """ + from nemo_automodel.components.distributed.config import DDPConfig, FSDP2Config, MegatronFSDPConfig + from nemo_automodel.components.distributed.mesh_utils import create_device_mesh + + strategy = engine_config.distributed_strategy + + if strategy == "fsdp2": + from torch.distributed.fsdp import MixedPrecisionPolicy + + from verl.utils.torch_dtypes import PrecisionType + + mp_policy = MixedPrecisionPolicy( + param_dtype=PrecisionType.to_dtype(engine_config.mp_param_dtype), + reduce_dtype=PrecisionType.to_dtype(engine_config.mp_reduce_dtype), + output_dtype=PrecisionType.to_dtype(engine_config.mp_output_dtype), + cast_forward_inputs=True, + ) + + distributed_config = FSDP2Config( + sequence_parallel=engine_config.sequence_parallel, + mp_policy=mp_policy, + activation_checkpointing=engine_config.activation_checkpointing, + defer_fsdp_grad_sync=engine_config.defer_fsdp_grad_sync, + ) + + elif strategy == "megatron_fsdp": + distributed_config = MegatronFSDPConfig( + activation_checkpointing=engine_config.activation_checkpointing, + ) + + elif strategy == "ddp": + distributed_config = DDPConfig( + activation_checkpointing=engine_config.activation_checkpointing, + ) + + else: + raise ValueError(f"Unsupported distributed_strategy: {strategy}") + + device_mesh, moe_mesh = create_device_mesh( + distributed_config, + tp_size=engine_config.tp_size, + pp_size=engine_config.pp_size, + cp_size=engine_config.cp_size, + ep_size=engine_config.ep_size, + dp_replicate_size=engine_config.dp_replicate_size, + world_size=world_size, + ) + + return distributed_config, device_mesh, moe_mesh + + +def build_automodel_model(model_config, engine_config, distributed_config, device_mesh, moe_mesh): + """Build a model using NeMoAutoModelForCausalLM.from_pretrained(). + + Args: + model_config: HFModelConfig with model path and settings. + engine_config: AutomodelEngineConfig with distributed settings. + distributed_config: FSDP2Config, MegatronFSDPConfig, or DDPConfig instance. + device_mesh: Pre-created device mesh (or None for DDP). + moe_mesh: Pre-created MoE mesh (or None). + + Returns: + A HuggingFace model with Automodel's distributed infrastructure applied. + """ + from nemo_automodel._transformers.auto_model import NeMoAutoModelForCausalLM + + kwargs = {} + + if engine_config.enable_fp8: + from nemo_automodel.components.quantization.fp8 import FP8Config + + kwargs["fp8_config"] = FP8Config() + + if engine_config.enable_compile: + from nemo_automodel.components.utils.compile_utils import CompileConfig + + kwargs["compile_config"] = CompileConfig() + + # Qwen/Llama with ep_size<=1: use HF implementation. + from transformers import AutoConfig + + _cfg = AutoConfig.from_pretrained(model_config.path, trust_remote_code=model_config.trust_remote_code) + _arch = (getattr(_cfg, "architectures", None) or [""])[0].lower() + if engine_config.ep_size <= 1 and ("qwen" in _arch or "llama" in _arch): + kwargs["force_hf"] = True + + if engine_config.backend_config and not kwargs.get("force_hf", False): + from nemo_automodel.components.models.common.utils import BackendConfig + + backend_kwargs = dict(engine_config.backend_config) + kwargs["backend"] = BackendConfig(**backend_kwargs) + + # MoE config for MoEParallelizerConfig + if engine_config.ep_size > 1: + from nemo_automodel.components.moe.config import MoEParallelizerConfig + + moe_kwargs = dict(engine_config.moe_config) if engine_config.moe_config else {} + if hasattr(distributed_config, "mp_policy"): + moe_kwargs.setdefault("mp_policy", distributed_config.mp_policy) + + kwargs["moe_config"] = MoEParallelizerConfig(**moe_kwargs) + + kwargs["attn_implementation"] = engine_config.attn_implementation + + from verl.utils.torch_dtypes import PrecisionType + + kwargs["torch_dtype"] = PrecisionType.to_dtype(engine_config.model_dtype) + + model = NeMoAutoModelForCausalLM.from_pretrained( + pretrained_model_name_or_path=model_config.path, + device_mesh=device_mesh, + moe_mesh=moe_mesh, + distributed_config=distributed_config, + activation_checkpointing=engine_config.activation_checkpointing, + trust_remote_code=model_config.trust_remote_code, + **kwargs, + ) + + return model + + +@torch.no_grad() +def offload_automodel_model_to_cpu(model, empty_cache=True): + """Offload an FSDP2-wrapped model to CPU (reshard, move to CPU, optional cache clear).""" + from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState + from torch.distributed.fsdp._fully_shard._fsdp_state import _get_module_fsdp_state + + for module in model.modules(): + state = _get_module_fsdp_state(module) + if state is None: + continue + fsdp_param_group = state._fsdp_param_group + + if fsdp_param_group is None: + continue + + fsdp_param_group._training_state = TrainingState.IDLE + + model.reshard() + model.cpu() + if empty_cache: + get_torch_device().empty_cache() + + +@torch.no_grad() +def load_automodel_model_to_gpu(model): + """Load model back to GPU.""" + device = get_device_id() + model.to(device, non_blocking=True) + + +@torch.no_grad() +def offload_automodel_optimizer(optimizer): + """Offload optimizer state to CPU.""" + if not optimizer.state: + return + for param_group in optimizer.param_groups: + for param in param_group["params"]: + state = optimizer.state[param] + for key, value in state.items(): + if isinstance(value, torch.Tensor): + state[key] = value.to("cpu", non_blocking=True) + + +@torch.no_grad() +def load_automodel_optimizer(optimizer, device_id): + """Load optimizer state back to GPU.""" + if not optimizer.state: + return + for param_group in optimizer.param_groups: + for param in param_group["params"]: + state = optimizer.state[param] + for key, value in state.items(): + if isinstance(value, torch.Tensor): + state[key] = value.to(device_id, non_blocking=True) diff --git a/verl/verl/workers/engine/base.py b/verl/verl/workers/engine/base.py new file mode 100644 index 0000000000000000000000000000000000000000..ff72de5834b68fb1da536c5cde59f0569e1a9076 --- /dev/null +++ b/verl/verl/workers/engine/base.py @@ -0,0 +1,337 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +The abstract base class defining the interface for model training engines. +""" + +from abc import abstractmethod +from contextlib import nullcontext +from typing import Any, Callable, ContextManager, Generator, Optional + +import torch +from tensordict import TensorDict + +from verl.utils.device import get_device_name +from verl.utils.tensordict_utils import maybe_fix_3d_position_ids + + +class BaseEngine: + """ + Abstract base class defining the interface for model training engines. Interface is subject to + change before release. + + Engine implementations must subclass BaseEngine and provide concrete behavior for all methods. + """ + + def initialize(self): + """ + Instantiate or load the model, optimizer, and learning rate scheduler. + + Should prepare all components necessary for training or evaluation. + """ + raise NotImplementedError + + @property + @abstractmethod + def is_param_offload_enabled(self) -> bool: + """Whether parameter offloading is enabled.""" + raise NotImplementedError + + @property + @abstractmethod + def is_optimizer_offload_enabled(self) -> bool: + """Whether optimizer offloading is enabled.""" + raise NotImplementedError + + def train_mode(self, **kwargs): + """ + Context manager entry for switching the engine and model into training mode. + + Usage: + with engine.train_mode(): + # runs in training mode + """ + raise NotImplementedError + + def eval_mode(self, **kwargs): + """ + Context manager entry for switching the engine and model into evaluation mode. + + Usage: + with engine.eval_mode(): + # runs in evaluation mode + """ + raise NotImplementedError + + def optimizer_zero_grad(self): + """ + Zero the gradients of the optimizer. + """ + raise NotImplementedError + + def optimizer_step(self): + """ + Perform an optimization step using the optimizer. + """ + raise NotImplementedError + + def lr_scheduler_step(self): + """ + Advance the learning rate scheduler by one step. + + Returns: + current_lr (float or list[float]): Updated learning rate(s). + """ + raise NotImplementedError + + def forward_backward_batch(self, data: TensorDict, loss_function: Callable, forward_only=False) -> Any: + """ + Perform a forward pass and optionally a backward pass on a batch of data. + + Args: + data: The input data for the forward pass, typically containing tensors and metadata. + loss_function: The loss function to optimize. See `verl.workers.roles.utils.losses` for examples. + forward_only: If True, perform only the forward pass. If False, perform forward and backward pass. + + Returns: + Any: The output of the forward pass, which can be used for loss computation or other purposes. + """ + raise NotImplementedError + + def train_batch(self, data: TensorDict, loss_function: Callable) -> Any: + """ + Perform a training step on a batch of data. + + Args: + data: The input data for training, typically containing tensors and metadata. + loss_function: A function that computes the loss and metrics given a batch and predictions. + + Returns: + dict[str, torch.Tensor]: A dictionary containing the aggregated training metrics for the batch. + """ + maybe_fix_3d_position_ids(data) + + self.optimizer_zero_grad() + outputs = self.forward_backward_batch(data, loss_function, forward_only=False) + grad_norm = self.optimizer_step() + if self.is_mp_src_rank_with_outputs(): + assert "grad_norm" not in outputs["metrics"] + outputs["metrics"]["grad_norm"] = grad_norm + return outputs + + def infer_batch(self, data: TensorDict, loss_function: Optional[Callable] = None) -> Any: + """ + Perform inference on a batch of data. + + Args: + data: The input data for inference, typically containing tensors and metadata. + + Returns: + Any: The output of the inference, which can be used for predictions or other purposes. + """ + # see comments from train_batch + maybe_fix_3d_position_ids(data) + + with torch.no_grad(): + outputs = self.forward_backward_batch(data, loss_function, forward_only=True) + return outputs + + def get_per_tensor_param(self) -> tuple[Generator[tuple[str, torch.Tensor], None, None], Optional[dict]]: + """ + Get a generator that yields per-tensor parameters and optional peft config. + + Returns: + Generator[tuple[str, torch.Tensor]]: A generator that yields tuples of parameter names and tensors. + Optional[dict]: Optional peft config. + """ + raise NotImplementedError + + def get_data_parallel_size(self): + raise NotImplementedError + + def get_data_parallel_rank(self): + raise NotImplementedError + + def get_data_parallel_group(self): + raise NotImplementedError + + def to(self, device: str, model: bool = True, optimizer: bool = True, grad: bool = True): + """ + Move model parameters, optimizer states, or both to the specified device. + + Args: + device: Target device identifier. + model: If True, move the model. + optimizer: If True, move the optimizer states. + grad: If True, move the gradient buffer. + """ + if grad: + assert model, "Gradient buffers must be moved to device along with model parameters" + + def save_checkpoint( + self, + local_path: str, + hdfs_path: Optional[str] = None, + global_step: int = 0, + max_ckpt_to_keep: Optional[int] = None, + **kwargs, + ) -> None: + """ + Save model, optimizer, and scheduler states to a checkpoint. + + Args: + local_path: Local filesystem path to save checkpoint. + hdfs_path: Optional HDFS path to copy checkpoint. + global_step: Integer training step number for naming. + max_ckpt_to_keep: Maximum number of recent checkpoints to retain. + **kwargs: Arbitrary keyword arguments. + """ + raise NotImplementedError + + def load_checkpoint( + self, local_path: str, hdfs_path: Optional[str] = None, del_local_after_load: bool = True, **kwargs + ) -> None: + """ + Load model, optimizer, and scheduler states from a checkpoint. + + Args: + local_path: Local filesystem path of the checkpoint. + hdfs_path: Optional HDFS path where checkpoint is stored. + del_local_after_load: Whether to delete local copy after loading. + **kwargs: Arbitrary keyword arguments. + """ + raise NotImplementedError + + def is_mp_src_rank_with_outputs(self): + """ + Whether the current rank is the first rank in model parallel group that contains model outputs + """ + raise NotImplementedError + + def disable_adapter(self) -> ContextManager: + """ + Disable all adapters temporarily under the context in the model for LoRA + """ + return nullcontext() + + +class BaseEngineCtx: + def __init__(self, engine: BaseEngine, mode, **kwargs): + """Base Engine context that handles load and offload + + Args: + engine: + **kwargs: + """ + self.engine = engine + self.mode = mode + assert self.mode in ("train", "eval") + self.disable_auto_offload = kwargs.pop("disable_auto_offload", False) + + def _context_switch(self, device): + if self.disable_auto_offload: + return + if device != "cpu": + if not self.engine.is_param_offload_enabled and not self.engine.is_optimizer_offload_enabled: + return + if self.mode == "eval": + self.engine.to(device=device, model=self.engine.is_param_offload_enabled, optimizer=False, grad=False) + elif self.mode == "train": + self.engine.to( + device=device, + model=self.engine.is_param_offload_enabled, + optimizer=self.engine.is_optimizer_offload_enabled, + grad=self.engine.is_param_offload_enabled, + ) + + def __enter__(self): + self._context_switch(get_device_name()) + self.engine.mode = self.mode + + def __exit__(self, exc_type, exc_val, exc_tb): + self._context_switch("cpu") + self.engine.mode = None + + +class EngineRegistry: + """ + A registry for managing and instantiating different types of training engines. + + This class uses a dictionary to store engine classes, mapping a string key to each class. + It provides a decorator `register` to add new engines to the registry and a `new` method + to create an instance of a registered engine. + """ + + _engines = {} + + @classmethod + def register(cls, model_type: str, backend: list[str] | str, device: list[str] | str = "cuda"): + """ + A class method decorator that registers an engine class with a given key. + + This allows for dynamic instantiation of engine classes by their registered key. + + Args: + model_type (str): The type of the model + backend (list[str] | str): The backend to use for the model type + device (list[str] | str): The device type (e.g., "cuda", "npu", "cpu") this engine supports, + default is "cuda" + + Returns: + A decorator function that takes an engine class and registers it. + """ + + def decorator(engine_class): + assert issubclass(engine_class, BaseEngine) + if model_type not in cls._engines: + cls._engines[model_type] = {} + + backends = backend if isinstance(backend, list) else [backend] + devices = device if isinstance(device, list) else [device] + for current_backend in backends: + for current_device in devices: + if current_backend not in cls._engines[model_type]: + cls._engines[model_type][current_backend] = {} + if current_device not in cls._engines[model_type][current_backend]: + cls._engines[model_type][current_backend][current_device] = engine_class + + return engine_class + + return decorator + + @classmethod + def get_engine_cls(cls, model_type: str, backend: str): + assert model_type in cls._engines, f"Unknown model_type: {model_type}" + assert backend in cls._engines[model_type], f"Unknown backend: {backend}" + device = get_device_name() + assert device in cls._engines[model_type][backend], ( + f"Unknown device: {device} for model_type: {model_type} and backend: {backend}" + ) + return cls._engines[model_type][backend][device] + + @classmethod + def new(cls, model_type, backend, *args, **kwargs): + """ + Function to create a new training engine instance based on the provided config. + Args: + key: A configuration object containing the engine key and other settings. + *args: Variable length argument list. + **kwargs: Arbitrary keyword arguments. + Returns: + engine: An instance of the training engine corresponding to the config. + Raises: + NotImplementedError: If the engine key in the config does not match any known engines. + """ + engine_cls = cls.get_engine_cls(model_type, backend) + return engine_cls(*args, **kwargs) diff --git a/verl/verl/workers/engine/fsdp/__init__.py b/verl/verl/workers/engine/fsdp/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a1bdb16b47cec72d684b5c9fbc61ab787e7e81c1 --- /dev/null +++ b/verl/verl/workers/engine/fsdp/__init__.py @@ -0,0 +1,16 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from .transformer_impl import FSDPEngine, FSDPEngineWithLMHead + +__all__ = ["FSDPEngine", "FSDPEngineWithLMHead"] diff --git a/verl/verl/workers/engine/fsdp/__pycache__/__init__.cpython-312.pyc b/verl/verl/workers/engine/fsdp/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a49334dcb3f62a476744acb268636688a5c14e29 Binary files /dev/null and b/verl/verl/workers/engine/fsdp/__pycache__/__init__.cpython-312.pyc differ diff --git a/verl/verl/workers/engine/fsdp/__pycache__/transformer_impl.cpython-312.pyc b/verl/verl/workers/engine/fsdp/__pycache__/transformer_impl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..94c43f44a0ade6e6280429e36411758b5cfb0f71 Binary files /dev/null and b/verl/verl/workers/engine/fsdp/__pycache__/transformer_impl.cpython-312.pyc differ diff --git a/verl/verl/workers/engine/fsdp/__pycache__/utils.cpython-312.pyc b/verl/verl/workers/engine/fsdp/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9deeaf38ede0c5c26954d81ea0579e3c80608da4 Binary files /dev/null and b/verl/verl/workers/engine/fsdp/__pycache__/utils.cpython-312.pyc differ diff --git a/verl/verl/workers/engine/fsdp/transformer_impl.py b/verl/verl/workers/engine/fsdp/transformer_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..7b106ff8aa5cf9e88eb2c82507801d342c1f595c --- /dev/null +++ b/verl/verl/workers/engine/fsdp/transformer_impl.py @@ -0,0 +1,1321 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +The concrete Engine implementation using PyTorch FullyShardedDataParallel (FSDP) +""" + +import gc +import logging +import os +import warnings +from contextlib import nullcontext +from typing import Callable, ContextManager, Optional + +import torch +import torch.distributed +from peft import LoraConfig, TaskType, get_peft_model +from tensordict import TensorDict +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +from torch.distributed.fsdp.api import FullStateDictConfig, ShardedStateDictConfig, StateDictType +from torch.distributed.tensor import DTensor + +import verl.utils.torch_functional as verl_F +from verl.models.transformers.monkey_patch import apply_monkey_patch +from verl.trainer.config import CheckpointConfig +from verl.utils import tensordict_utils as tu +from verl.utils.activation_offload import enable_activation_offloading +from verl.utils.checkpoint.fsdp_checkpoint_manager import FSDPCheckpointManager +from verl.utils.dataset.dataset_utils import DatasetPadMode +from verl.utils.debug import log_gpu_memory_usage +from verl.utils.device import get_device_id, get_device_name +from verl.utils.fsdp_utils import ( + CPUOffloadPolicy, + FSDPModule, + MixedPrecisionPolicy, + apply_fsdp2, + collect_lora_params, + fsdp2_clip_grad_norm_, + fsdp2_load_full_state_dict, + fsdp_version, + get_fsdp_wrap_policy, + get_init_weight_context_manager, + init_fn, + load_fsdp_model_to_gpu, + load_fsdp_optimizer, + merged_lora_context, + normalize_peft_param_name, + offload_fsdp_model_to_cpu, + offload_fsdp_optimizer, + replace_lora_wrapper, +) +from verl.utils.model import convert_weight_keys, extract_multi_modal_inputs +from verl.utils.py_functional import convert_to_regular_types +from verl.utils.torch_functional import logprobs_from_logits +from verl.utils.ulysses import ( + gather_outputs_and_unpad, + get_ulysses_sequence_parallel_group, + set_ulysses_sequence_parallel_group, + ulysses_pad, + ulysses_pad_and_slice_inputs, +) +from verl.workers.config import FSDPEngineConfig, FSDPOptimizerConfig, HFModelConfig +from verl.workers.utils.padding import build_attention_mask_from_nested + +from ..base import BaseEngine, BaseEngineCtx, EngineRegistry +from ..utils import enable_full_determinism, postprocess_batch_func, prepare_micro_batches +from .utils import create_device_mesh, get_sharding_strategy + +logger = logging.getLogger(__file__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + +device_name = get_device_name() + + +class FSDPEngine(BaseEngine): + """ + Concrete Engine implementation using PyTorch FullyShardedDataParallel (FSDP). + + Supports model sharding, activation/optimizer offloading, LoRA, and sequence parallelism. + """ + + def __init__( + self, + model_config: HFModelConfig, + engine_config: FSDPEngineConfig, + optimizer_config: FSDPOptimizerConfig, + checkpoint_config: CheckpointConfig, + ): + """ + Initialize the FSDPEngine. + + Sets up distributed device meshes, LoRA, and offload policies based on config. + + Args: + config: Configuration object with FSDP and model settings. + """ + super().__init__() + + self.model_config = model_config + self.engine_config = engine_config + self.optimizer_config = optimizer_config + self.checkpoint_config = checkpoint_config + + self.mode = None + + self.rank = torch.distributed.get_rank() + + # Apply NPU patches for FSDP backend + from .utils import apply_npu_fsdp_patches + + apply_npu_fsdp_patches() + + # build device mesh for Ulysses Sequence Parallel + + self.use_remove_padding = self.model_config.use_remove_padding + + self._init_device_mesh() + + if self.engine_config.full_determinism: + enable_full_determinism(seed=self.engine_config.seed) + + # set FSDP offload params + self._is_offload_param = self.engine_config.param_offload + self._is_offload_optimizer = self.engine_config.optimizer_offload + self._is_lora = self.model_config.lora_rank > 0 + + # Defaults for mixed-precision state. _build_fsdp_module overrides these when it + # runs; subclasses that bypass _build_fsdp_module (e.g. VeOmniEngine) keep the + # defaults so forward_step / optimizer_step can still read them safely. + self._autocast_dtype = torch.bfloat16 + self.scaler = None + + # QAT (Quantization-Aware Training) + self._qat_config = getattr(self.engine_config, "qat", None) + self._qat_enabled = self._qat_config is not None and getattr(self._qat_config, "enable", False) + if self._qat_enabled: + logger.info(f"QAT enabled: mode={self._qat_config.mode}, group_size={self._qat_config.group_size}") + + if self.engine_config.entropy_from_logits_with_chunking: + entropy_from_logits = verl_F.entropy_from_logits_with_chunking + else: + entropy_from_logits = verl_F.entropy_from_logits + + self.compute_entropy_from_logits = ( + torch.compile(entropy_from_logits, dynamic=True) + if self.engine_config.use_torch_compile # use torch compile by default + else entropy_from_logits + ) + + @property + def is_param_offload_enabled(self) -> bool: + return self._is_offload_param + + @property + def is_optimizer_offload_enabled(self) -> bool: + return self._is_offload_optimizer + + def is_mp_src_rank_with_outputs(self): + if self.ulysses_device_mesh is not None: + is_collect = self.ulysses_device_mesh["sp"].get_local_rank() == 0 + else: + is_collect = True + return is_collect + + def initialize(self): + """ + Build the model, optimizer, and learning rate scheduler under FSDP. + + Applies device, dtype, and precision configurations, including mixed precision. + Sets up checkpoint manager and FLOPs counter. + """ + # This is used to import external_lib into the huggingface systems + self._build_model_optimizer() + + self.checkpoint_manager = FSDPCheckpointManager( + model=self.module, + optimizer=self.optimizer, + lr_scheduler=self.lr_scheduler, + processing_class=self.model_config.get_processor(), + checkpoint_config=self.checkpoint_config, + trust_remote_code=self.model_config.trust_remote_code, + ) + + self.to( + device="cpu", + model=self._is_offload_param, + optimizer=self._is_offload_optimizer, + grad=self._is_offload_param, + ) + + log_gpu_memory_usage("After offload model/optimizer/grad during init", logger=logger) + + def _init_device_mesh(self): + world_size = torch.distributed.get_world_size() + from torch.distributed.device_mesh import init_device_mesh + + fsdp_size = self.engine_config.fsdp_size + + self.device_mesh = create_device_mesh(world_size=world_size, fsdp_size=fsdp_size) + self.ulysses_device_mesh = None + self.ulysses_parallel_group = None + self.ulysses_sequence_parallel_size = self.engine_config.ulysses_sequence_parallel_size + dp_size = self.get_data_parallel_size() + if self.ulysses_sequence_parallel_size > 1: + self.ulysses_device_mesh = init_device_mesh( + device_name, mesh_shape=(dp_size, self.ulysses_sequence_parallel_size), mesh_dim_names=["dp", "sp"] + ) + self.ulysses_parallel_group = self.ulysses_device_mesh["sp"].get_group() + + self.use_ulysses_sp = self.ulysses_sequence_parallel_size > 1 + + def _build_module(self): + from verl.utils.model import get_hf_auto_model_class + from verl.utils.torch_dtypes import PrecisionType + + torch_dtype = self.engine_config.model_dtype + + if torch_dtype is None: + # if it is training, we force torch_dtype to fp32 + torch_dtype = torch.float32 if not self.engine_config.forward_only else torch.bfloat16 + + torch_dtype = PrecisionType.to_dtype(torch_dtype) + + init_context = get_init_weight_context_manager( + use_meta_tensor=not self.model_config.hf_config.tie_word_embeddings, mesh=self.device_mesh + ) + + with init_context(), warnings.catch_warnings(): + warnings.simplefilter("ignore") + + if self.model_config.model_type == "language_model": + auto_class = get_hf_auto_model_class(hf_config=self.model_config.hf_config) + + module = auto_class.from_pretrained( + pretrained_model_name_or_path=self.model_config.local_path, + torch_dtype=torch_dtype, + config=self.model_config.hf_config, + trust_remote_code=self.model_config.trust_remote_code, + ) + else: + from verl.utils.model import load_valuehead_model + + assert self.model_config.model_type == "value_model", ( + f"Unsupported model type: {self.model_config.model_type}" + ) + self.model_config.hf_config.num_labels = 1 + self.model_config.hf_config.classifier_dropout = 0.0 + self.model_config.hf_config.hidden_dropout = "0" + self.model_config.hf_config.summary_dropout_prob = 0.0 + module = load_valuehead_model( + local_path=self.model_config.local_path, + torch_dtype=torch_dtype, + model_config=self.model_config.hf_config, + trust_remote_code=self.model_config.trust_remote_code, + ) + + use_liger = self.model_config.use_liger + # Apply Liger kernel; disable fused_linear_cross_entropy (conflicts with verl's forward patching) + if use_liger: + from liger_kernel.transformers.monkey_patch import _apply_liger_kernel_to_instance + + _apply_liger_kernel_to_instance( + model=module, + fused_linear_cross_entropy=False, + swiglu=True, + ) + + fused_kernel_options = self.model_config.fused_kernel_options + fused_kernels_backend = ( + fused_kernel_options.get("impl_backend", None) if fused_kernel_options is not None else None + ) + + use_fused_kernels = self.model_config.use_fused_kernels + apply_monkey_patch( + model=module, + use_remove_padding=self.use_remove_padding, + ulysses_sp_size=self.ulysses_sequence_parallel_size, + use_fused_kernels=use_fused_kernels, + fused_kernels_backend=fused_kernels_backend, + ) + + # some parameters may not in torch_dtype + module.to(torch_dtype) + + if self.model_config.enable_gradient_checkpointing: + module.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) + return module + + def _build_lora_module(self, module): + module.enable_input_require_grads() + + lora_adapter_path = getattr(self.model_config, "lora_adapter_path", None) + if lora_adapter_path is not None: + from peft import PeftModel + + from verl.utils.fs import copy_to_local + + print(f"Loading pre-trained LoRA adapter to from: {lora_adapter_path}") + # Copy adapter to local if needed + local_adapter_path = copy_to_local(lora_adapter_path, use_shm=self.model_config.use_shm) + + module = PeftModel.from_pretrained(module, local_adapter_path, is_trainable=True) + peft_config = module.peft_config["default"] + # Ensure task_type is TaskType enum, not string + if isinstance(peft_config.task_type, str): + peft_config.task_type = TaskType.CAUSAL_LM + else: + # Convert config to regular Python types before creating PEFT model + lora_config = { + "task_type": TaskType.CAUSAL_LM, + "r": self.model_config.lora_rank, + "lora_alpha": self.model_config.lora_alpha, + "target_modules": convert_to_regular_types(self.model_config.target_modules), + "target_parameters": convert_to_regular_types(self.model_config.target_parameters), + "exclude_modules": convert_to_regular_types(self.model_config.exclude_modules), + "bias": "none", + } + module = get_peft_model(module, LoraConfig(**lora_config)) + + return module + + def _build_fsdp_module(self, module): + # TODO(ziheng): need to improve + from torch.distributed.fsdp import CPUOffload, MixedPrecision + + from verl.utils.torch_dtypes import PrecisionType + + mixed_precision_config = self.engine_config.mixed_precision + if mixed_precision_config is not None: + param_dtype = PrecisionType.to_dtype(mixed_precision_config.get("param_dtype", "bf16")) + reduce_dtype = PrecisionType.to_dtype(mixed_precision_config.get("reduce_dtype", "fp32")) + buffer_dtype = PrecisionType.to_dtype(mixed_precision_config.get("buffer_dtype", "fp32")) + else: + param_dtype = torch.bfloat16 + reduce_dtype = torch.float32 + buffer_dtype = torch.float32 + + mixed_precision = MixedPrecision(param_dtype=param_dtype, reduce_dtype=reduce_dtype, buffer_dtype=buffer_dtype) + + self._autocast_dtype = param_dtype + # fp16 training requires loss scaling to avoid gradient underflow. Mirror the pattern + # landed in #4036 for the legacy dp_actor path. bf16 / fp32 do not need a scaler. + if param_dtype == torch.float16: + from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler + + self.scaler = ShardedGradScaler(growth_interval=400) + else: + self.scaler = None + + auto_wrap_policy = get_fsdp_wrap_policy( + module=module, + config=self.engine_config.wrap_policy, + is_lora=self.model_config.lora_rank > 0, + ) + + fsdp_mesh = self.device_mesh + sharding_strategy = get_sharding_strategy(fsdp_mesh) + + # Note: We force turn off CPUOffload because it causes incorrect results when using grad accumulation + if self.engine_config.strategy == "fsdp": + # cpu_offload: + # - actor: None + # - critic: None + # - ref: CPUOffload(offload_params=True) + + # We force reference policy to use CPUOffload to save memory. + # We force turn off CPUOffload for actor because it causes incorrect results when using grad accumulation + cpu_offload = None + if self.engine_config.forward_only: + cpu_offload = CPUOffload(offload_params=True) + self._is_offload_param = False + self._is_offload_optimizer = False + + module = FSDP( + module, + param_init_fn=init_fn, + auto_wrap_policy=auto_wrap_policy, + device_id=get_device_id(), + sharding_strategy=sharding_strategy, + mixed_precision=mixed_precision, + sync_module_states=True, + device_mesh=self.device_mesh, + forward_prefetch=self.engine_config.forward_prefetch, + use_orig_params=self.engine_config.use_orig_params, + cpu_offload=cpu_offload, + ) + elif self.engine_config.strategy == "fsdp2": + # - actor: offload_policy + # - critic: offload_policy + # - ref: CPUOffloadPolicy(pin_memory=True) + assert CPUOffloadPolicy is not None, "PyTorch version >= 2.4 is required for using fully_shard API (FSDP2)" + mp_policy = MixedPrecisionPolicy( + param_dtype=param_dtype, reduce_dtype=reduce_dtype, cast_forward_inputs=True + ) + offload_policy = None + if self.engine_config.offload_policy or self.engine_config.forward_only: + self._is_offload_param = False + self._is_offload_optimizer = False + offload_policy = CPUOffloadPolicy(pin_memory=True) + + fsdp_kwargs = { + "mesh": fsdp_mesh, + "mp_policy": mp_policy, + "offload_policy": offload_policy, + "reshard_after_forward": self.engine_config.reshard_after_forward, + } + full_state = module.state_dict() + apply_fsdp2(module, fsdp_kwargs, self.engine_config) + fsdp2_load_full_state_dict(module, full_state, fsdp_mesh, offload_policy) + else: + raise NotImplementedError(f"Unknown strategy {self.engine_config.strategy}") + + if self.model_config.enable_activation_offload: + enable_gradient_checkpointing = self.model_config.enable_gradient_checkpointing + enable_activation_offloading(module, self.engine_config.strategy, enable_gradient_checkpointing) + + if torch.distributed.get_world_size() == 1 and fsdp_version(module) == 1: + FSDP.set_state_dict_type( + module, + state_dict_type=StateDictType.FULL_STATE_DICT, + state_dict_config=FullStateDictConfig(), + ) + elif fsdp_version(module) == 1: + FSDP.set_state_dict_type( + module, + state_dict_type=StateDictType.SHARDED_STATE_DICT, + state_dict_config=ShardedStateDictConfig(), + ) + + return module + + def _build_optimizer(self, module): + from verl.workers.config.optimizer import build_optimizer + + optimizer = build_optimizer(module.parameters(), self.optimizer_config) + + return optimizer + + def _build_lr_scheduler(self, optimizer): + from verl.utils.torch_functional import get_constant_schedule_with_warmup, get_cosine_schedule_with_warmup + + optim_config = self.optimizer_config + + total_steps = optim_config.total_training_steps + num_warmup_steps = optim_config.lr_warmup_steps + lr_scheduler_type = optim_config.lr_scheduler_type + min_lr_ratio = optim_config.min_lr_ratio + num_cycles = optim_config.num_cycles + zero_indexed_step = optim_config.zero_indexed_step + if num_warmup_steps <= 0: + num_warmup_steps_ratio = optim_config.lr_warmup_steps_ratio + num_warmup_steps = int(num_warmup_steps_ratio * total_steps) + + if self.rank == 0: + print(f"Total steps: {total_steps}, num_warmup_steps: {num_warmup_steps}") + + if lr_scheduler_type == "constant": + lr_scheduler = get_constant_schedule_with_warmup(optimizer=optimizer, num_warmup_steps=num_warmup_steps) + elif lr_scheduler_type == "cosine": + lr_scheduler = get_cosine_schedule_with_warmup( + optimizer=optimizer, + num_warmup_steps=num_warmup_steps, + num_training_steps=total_steps, + min_lr_ratio=min_lr_ratio, + num_cycles=num_cycles, + zero_indexed_step=zero_indexed_step, + ) + else: + raise NotImplementedError(f"LR scheduler type {lr_scheduler_type} is not supported") + return lr_scheduler + + def _apply_qat(self, module): + """Apply QAT transformations to the model before FSDP wrapping.""" + from verl.utils.qat.core import apply_qat, enable_qat_fuse + + module = apply_qat( + module, + { + "enable": self._qat_config.enable, + "mode": self._qat_config.mode, + "group_size": self._qat_config.group_size, + "ignore_patterns": list(self._qat_config.ignore_patterns), + "activation_observer": self._qat_config.activation_observer, + }, + ) + enable_qat_fuse(module) + + if self._qat_config.mode == "w4a4": + self._restore_w4a4_input_scales(module, self.model_config.local_path) + + return module + + def _restore_w4a4_input_scales(self, model, model_path): + """Restore input_global_scale and input_amax from checkpoint for W4A4 mode.""" + import glob + + from safetensors import safe_open + + safetensor_files = glob.glob(f"{model_path}/model*.safetensors") + loaded_count = 0 + + for sf_path in safetensor_files: + with safe_open(sf_path, framework="pt") as f: + for key in f.keys(): + if "input_global_scale" in key: + module_path = key.replace(".input_global_scale", "") + amax_key = f"{module_path}.input_amax" + + module = model + for part in module_path.split("."): + module = module[int(part)] if part.isdigit() else getattr(module, part) + + scale_val = f.get_tensor(key) + val = scale_val.item() if scale_val.numel() == 1 else scale_val.max().item() + module.input_global_scale.fill_(val) + + amax_val = f.get_tensor(amax_key) + amax = amax_val.item() if amax_val.numel() == 1 else amax_val.max().item() + module.input_amax.fill_(amax) + loaded_count += 1 + + logger.info(f"[QAT W4A4] Restored {loaded_count} input_global_scale/input_amax from {model_path}") + + def _build_model_optimizer(self): + from verl.utils.model import print_model_size + + # Load base model with specified configuration and dtype + module = self._build_module() + # Apply LoRA adapters if low-rank adaptation is enabled + if self._is_lora: + module = self._build_lora_module(module) + + # Apply QAT before FSDP wrapping (training only) + if self._qat_enabled and not self.engine_config.forward_only: + module = self._apply_qat(module) + + # Synchronize all distributed processes before proceeding + torch.distributed.barrier() + if self.rank == 0: + print_model_size(module) + log_gpu_memory_usage("After init model from HF AutoModel", logger=logger) + + # Wrap model with FSDP for distributed training (sharding, mixed precision, etc.) + log_gpu_memory_usage("Before FSDP", logger=None) + module = self._build_fsdp_module(module) + log_gpu_memory_usage("After FSDP", logger=None) + + if not self.engine_config.forward_only: + # Initialize optimizer with model parameters and config settings + optimizer = self._build_optimizer(module) + # Create learning rate scheduler with warmup and decay settings + lr_scheduler = self._build_lr_scheduler(optimizer) + else: + optimizer = None + lr_scheduler = None + + self.module = module + self.optimizer = optimizer + self.lr_scheduler = lr_scheduler + + def train_mode(self, **kwargs): + """ + Return a context manager that switches to training mode with FSDP-specific handling. + + Includes parameter and optimizer offload entry/exit. + """ + return EngineTrainModeCtx(self, **kwargs) + + def eval_mode(self, **kwargs): + """ + Return a context manager that switches to evaluation mode with FSDP-specific handling. + + Includes activation offload entry/exit. + """ + return EngineEvalModeCtx(self, **kwargs) + + def get_data_parallel_rank(self): + if self.ulysses_device_mesh is not None: + return self.ulysses_device_mesh["dp"].get_local_rank() + else: + return torch.distributed.get_rank() + + def get_data_parallel_size(self): + return torch.distributed.get_world_size() // self.ulysses_sequence_parallel_size + + def get_data_parallel_group(self): + if self.ulysses_device_mesh is not None: + return self.ulysses_device_mesh.get_group(mesh_dim="dp") + else: + return torch.distributed.group.WORLD + + def get_model_parallel_group(self): + raise NotImplementedError + + def get_context_parallel_group(self): + raise NotImplementedError + + def forward_backward_batch(self, data: TensorDict, loss_function: Callable, forward_only=False) -> list[TensorDict]: + # note that the global_batch_size should include data on all the dp + tu.assign_non_tensor(data, sp_size=self.ulysses_sequence_parallel_size) + + # compute num_tokens in global batch for loss normalization + batch_num_tokens = data["loss_mask"].sum().to(get_device_id()) + torch.distributed.all_reduce( + batch_num_tokens, op=torch.distributed.ReduceOp.SUM, group=self.get_data_parallel_group() + ) + tu.assign_non_tensor(data, batch_num_tokens=batch_num_tokens.item()) + tu.assign_non_tensor(data, dp_size=self.get_data_parallel_size()) + + micro_batches, indices = prepare_micro_batches( + data=data, dp_group=self.get_data_parallel_group(), same_micro_num_in_dp=True + ) + + output_lst = [] + + ctx = torch.no_grad() if forward_only else nullcontext() + + # getattr fallback: some subclasses (e.g. VeOmniEngine) bypass FSDPEngine.__init__ + # and _build_fsdp_module, so self.scaler may not be set. + scaler = getattr(self, "scaler", None) + + for micro_batch in micro_batches: + with ctx: + loss, meta_info = self.forward_step(micro_batch, loss_function=loss_function, forward_only=forward_only) + + if not forward_only: + if scaler is not None: + scaler.scale(loss).backward() + else: + loss.backward() + + output_lst.append(meta_info) + + # postprocess and return + return postprocess_batch_func(output_lst=output_lst, indices=indices, data=data) + + def forward_step(self, micro_batch: TensorDict, loss_function, forward_only): + raise NotImplementedError("forward_step must be implemented in subclass") + + def optimizer_zero_grad(self): + """ + Zero gradients and enforce FSDP grad-clipping logic. + """ + self.optimizer.zero_grad() + + def optimizer_step(self): + """ + Clip gradients, skip update if non-finite, and step optimizer. + + Returns: + grad_norm (float): Norm of gradients before clipping. + """ + assert self.optimizer_config.clip_grad is not None + + # getattr fallback: some subclasses (e.g. VeOmniEngine) bypass FSDPEngine.__init__. + scaler = getattr(self, "scaler", None) + + # Unscale gradients before clip so the clip threshold is applied to true gradient + # magnitudes, not scaled ones. scaler.step() will skip the update if any grad is inf/nan. + if scaler is not None: + scaler.unscale_(self.optimizer) + + if isinstance(self.module, FSDP): + grad_norm = self.module.clip_grad_norm_(self.optimizer_config.clip_grad) + elif isinstance(self.module, FSDPModule): + grad_norm = fsdp2_clip_grad_norm_(self.module.parameters(), max_norm=self.optimizer_config.clip_grad) + else: + grad_norm = torch.nn.utils.clip_grad_norm_( + self.module.parameters(), max_norm=self.optimizer_config.clip_grad + ) + + if isinstance(grad_norm, DTensor): + grad_norm = grad_norm.full_tensor() + + if scaler is not None: + # scaler handles inf/nan skipping internally via _check_inf_per_device. + scaler.step(self.optimizer) + scaler.update() + else: + # if grad_norm is not finite, skip the update + if not torch.isfinite(grad_norm): + print(f"WARN: grad_norm is not finite: {grad_norm}") + self.optimizer.zero_grad() + else: + self.optimizer.step() + + if self._qat_enabled: + from verl.utils.qat.core import invalidate_all_scales + + invalidate_all_scales(self.module) + + return grad_norm.item() + + def lr_scheduler_step(self): + """ + Advance FSDP scheduler and return updated learning rate. + """ + self.lr_scheduler.step() + lr = self.lr_scheduler.get_last_lr()[0] # only return the first group + return lr + + def to(self, device: str, model: bool = True, optimizer: bool = True, grad: bool = True): + """ + Move FSDP model and/or optimizer to CPU or GPU with offload support. + Note that this function executes irrespective of offload config. It serves as manual control + """ + super().to(device=device, model=model, optimizer=optimizer, grad=grad) + + if self.engine_config.forward_only: + # force cpu_offload + return + + device_name = get_device_name() + + assert device in (device_name, "cpu") + if device == device_name: + if model: + load_fsdp_model_to_gpu(self.module) + if optimizer and self.optimizer is not None: + load_fsdp_optimizer(self.optimizer, device) + gc.collect() + elif device == "cpu": + if model: + offload_fsdp_model_to_cpu(self.module) + if optimizer and self.optimizer is not None: + offload_fsdp_optimizer(self.optimizer) + else: + raise ValueError(f"Invalid device type: {device}") + + def save_checkpoint( + self, + local_path: str, + hdfs_path: Optional[str] = None, + global_step: int = 0, + max_ckpt_to_keep: Optional[int] = None, + **kwargs, + ) -> None: + """ + Save FSDP checkpoint, handling parameter offload as needed. + """ + origin_module_device = next(self.module.parameters()).device.type + if self._is_offload_param or origin_module_device == "cpu": + load_fsdp_model_to_gpu(self.module) + + self.checkpoint_manager.save_checkpoint( + local_path=local_path, hdfs_path=hdfs_path, global_step=global_step, max_ckpt_to_keep=max_ckpt_to_keep + ) + + torch.distributed.barrier() + if self._is_offload_param: + offload_fsdp_model_to_cpu(self.module) + + def load_checkpoint( + self, local_path: str, hdfs_path: Optional[str] = None, del_local_after_load: int = True, **kwargs + ) -> None: + """ + Load FSDP checkpoint, restoring parameters and optimizer state. + """ + import torch + + if self._is_offload_param: + load_fsdp_model_to_gpu(self.module) + + self.checkpoint_manager.load_checkpoint( + local_path=local_path, hdfs_path=hdfs_path, del_local_after_load=del_local_after_load + ) + + torch.distributed.barrier() + if self._is_offload_param: + offload_fsdp_model_to_cpu(self.module) + + if self._is_offload_optimizer: + offload_fsdp_optimizer(self.optimizer) + + def get_per_tensor_param(self, layered_summon=False, base_sync_done=False, **kwargs): + log_gpu_memory_usage("Before load_fsdp_model_to_gpu", logger=logger) + + load_fsdp_model_to_gpu(self.module) + + log_gpu_memory_usage("After load_fsdp_model_to_gpu", logger=logger) + + peft_config = None + merge_lora = self.model_config.lora.get("merge", False) + + peft_model = getattr(self.module, "_fsdp_wrapped_module", self.module) + if hasattr(peft_model, "peft_config"): # LoRA + if not merge_lora: + peft_config = peft_model.peft_config.get("default", None) + params = collect_lora_params( + module=self.module, + layered_summon=layered_summon, + base_sync_done=base_sync_done, + ) + if not base_sync_done: + params = {replace_lora_wrapper(k, peft_config): v for k, v in params.items()} + else: # merge lora + with merged_lora_context(self.module, backup_adapters=True): + params = self.module.state_dict() + params = normalize_peft_param_name(params) + else: + params = self.module.state_dict() + + params = convert_weight_keys(params, getattr(self.module, "_fsdp_wrapped_module", self.module)) + + log_gpu_memory_usage("Before offload_fsdp_model_to_cpu", logger=logger) + if self._is_offload_param: + offload_fsdp_model_to_cpu(self.module) + log_gpu_memory_usage("After offload_fsdp_model_to_cpu", logger=logger) + + if peft_config is not None and base_sync_done: + per_tensor_param = params.items() + else: + device = get_device_id() # used when fsdp2 set cpu_offload_policy + # TODO: cast fp32 to bf16 to reduce weight sync overhead, need more fine-grained control, e.g MoE gate + per_tensor_param = ( + ( + name, + param.to(device, non_blocking=True).full_tensor().to(torch.bfloat16, non_blocking=True) + if isinstance(param, DTensor) + else param, + ) + for name, param in params.items() + ) + + if self._qat_enabled: + from verl.utils.qat.quantizer import QATQuantizer + from verl.utils.torch_dtypes import PrecisionType + + mixed_precision_config = self.engine_config.mixed_precision + if mixed_precision_config is not None: + param_dtype = PrecisionType.to_dtype(mixed_precision_config.get("param_dtype", "bf16")) + else: + param_dtype = torch.bfloat16 + + quantizer = QATQuantizer( + mode=self._qat_config.mode, + group_size=self._qat_config.group_size, + ignore_patterns=list(self._qat_config.ignore_patterns), + device=torch.device(get_device_id()), + param_dtype=param_dtype, + ) + per_tensor_param = quantizer.quantize_with_fusion( + per_tensor_param, + target_device=torch.device("cpu"), + ) + + peft_config_dict = peft_config.to_dict() if peft_config is not None else None + return per_tensor_param, peft_config_dict + + def disable_adapter(self) -> ContextManager: + return self.module.disable_adapter() + + +class EngineEvalModeCtx(BaseEngineCtx): + def __init__(self, engine: FSDPEngine, **kwargs): + super().__init__(engine=engine, mode="eval", **kwargs) + + def __enter__(self): + assert isinstance(self.engine, FSDPEngine) + super().__enter__() + self.prev_sp_group = get_ulysses_sequence_parallel_group() + set_ulysses_sequence_parallel_group(self.engine.ulysses_parallel_group) + self.engine.module.eval() + + def __exit__(self, exc_type, exc_value, traceback): + assert isinstance(self.engine, FSDPEngine) + set_ulysses_sequence_parallel_group(self.prev_sp_group) + + # https://pytorch.org/docs/stable/notes/fsdp.html#fsdp-notes + # unshard the root FSDP module + if self.engine.engine_config.fsdp_size > 1: + if fsdp_version(self.engine.module) == 1: + self.engine.module._handle.reshard(True) + elif fsdp_version(self.engine.module) == 2: + self.engine.module.reshard() + + super().__exit__(exc_type, exc_value, traceback) + + +class EngineTrainModeCtx(BaseEngineCtx): + def __init__(self, engine: FSDPEngine, **kwargs): + super().__init__(engine=engine, mode="train", **kwargs) + + def __enter__(self): + assert isinstance(self.engine, FSDPEngine) + super().__enter__() + self.prev_sp_group = get_ulysses_sequence_parallel_group() + set_ulysses_sequence_parallel_group(self.engine.ulysses_parallel_group) + self.engine.module.train() + + def __exit__(self, exc_type, exc_value, traceback): + assert isinstance(self.engine, FSDPEngine) + set_ulysses_sequence_parallel_group(self.prev_sp_group) + self.engine.optimizer_zero_grad() + super().__exit__(exc_type, exc_value, traceback) + + +@EngineRegistry.register(model_type="language_model", backend=["fsdp", "fsdp2"], device=["cuda", "npu"]) +class FSDPEngineWithLMHead(FSDPEngine): + def prepare_model_inputs(self, micro_batch: TensorDict): + use_remove_padding = tu.get_non_tensor_data(data=micro_batch, key="use_remove_padding", default=True) + pad_mode = tu.get_non_tensor_data(data=micro_batch, key="pad_mode", default=DatasetPadMode.NO_PADDING) + use_fused_kernels = tu.get_non_tensor_data(data=micro_batch, key="use_fused_kernels", default=False) + temperature = micro_batch["temperature"] + temperature_item = temperature + if use_fused_kernels: + assert not isinstance(temperature, torch.Tensor), ( + "use_fused_kernels does not support per sample temperature yet" + ) + assert pad_mode == DatasetPadMode.NO_PADDING, f"pad_mode {pad_mode} not supported" + + multi_modal_inputs = extract_multi_modal_inputs(micro_batch.get("multi_modal_inputs", [])) + input_ids = micro_batch["input_ids"] + position_ids = micro_batch["position_ids"] + + if not isinstance(temperature, torch.Tensor): + temperature = torch.tensor([temperature] * input_ids.shape[0], device=input_ids.device) + + temperature = temperature.to(torch.float32) + assert temperature.shape[0] == input_ids.shape[0] + + # args used to get outputs + output_args = {} + + if use_remove_padding: + # support per sample temperature + # temperature (bsz,) + # input_ids (bsz, j1) + temperature_rmpad = verl_F.expand_as_nested(temperature, input_ids).values() # (total_nnz,) + temperature_rmpad = temperature_rmpad.unsqueeze(0) # (1, total_nnz) + + if pad_mode == DatasetPadMode.NO_PADDING: + input_ids_rmpad = input_ids.values().unsqueeze(0) # (1, total_nnz) + if position_ids.dim() == 3: + position_ids_rmpad = position_ids.values().unsqueeze(1) # (4, 1, total_nnz) + else: + position_ids_rmpad = position_ids.values().unsqueeze(0) # (1, total_nnz) + else: + raise NotImplementedError(f"pad_mode {pad_mode} not implemented") + + # for compute the log_prob + input_ids_rmpad_rolled = torch.roll(input_ids_rmpad, shifts=-1, dims=1) # (1, total_nnz) + + # pad and slice the inputs if sp > 1 + if self.use_ulysses_sp: + is_vlm_model = hasattr(getattr(self.module, "module", self.module).config, "vision_config") + if is_vlm_model: + # vlm model's inputs will be sliced after embedding + input_ids_rmpad, position_ids_rmpad, pad_size = ulysses_pad( + input_ids_rmpad, + position_ids_rmpad=position_ids_rmpad, + sp_size=self.ulysses_sequence_parallel_size, + ) + else: + input_ids_rmpad, position_ids_rmpad, pad_size = ulysses_pad_and_slice_inputs( + input_ids_rmpad, + position_ids_rmpad=position_ids_rmpad, + sp_size=self.ulysses_sequence_parallel_size, + skip_position_ids_rmpad=True if self.__class__.__name__ == "VeOmniEngineWithLMHead" else False, + ) + input_ids_rmpad_rolled, _, _ = ulysses_pad_and_slice_inputs( + input_ids_rmpad_rolled, + position_ids_rmpad=None, + sp_size=self.ulysses_sequence_parallel_size, + ) + + temperature_rmpad, _, _ = ulysses_pad_and_slice_inputs( + temperature_rmpad, position_ids_rmpad=None, sp_size=self.ulysses_sequence_parallel_size, pad_value=1 + ) + + output_args["pad_size"] = pad_size + + input_ids_rmpad_rolled = input_ids_rmpad_rolled.squeeze(0) # ((total_nnz / sp) + pad) + temperature_rmpad = temperature_rmpad.squeeze(0) + output_args["input_ids_rmpad_rolled"] = input_ids_rmpad_rolled + output_args["temperature_rmpad"] = temperature_rmpad + + # only pass input_ids and position_ids to enable flash_attn_varlen + + model_inputs = { + "input_ids": input_ids_rmpad, + "attention_mask": None, + "position_ids": position_ids_rmpad, + } + + else: + if pad_mode == DatasetPadMode.NO_PADDING: + input_ids = micro_batch["input_ids"] + position_ids = micro_batch["position_ids"] + pad_token_id = tu.get_non_tensor_data(data=micro_batch, key="pad_token_id", default=0) + batch_size = micro_batch.batch_size[0] + seq_len_effective = input_ids.offsets().diff() + max_seq_len = int(seq_len_effective.max().item()) + + input_ids_rmpad_rolled = torch.roll(input_ids.values(), shifts=-1, dims=0) + output_args["input_ids_rmpad_rolled"] = input_ids_rmpad_rolled + # we store the per sample temperature + output_args["temperature"] = temperature + + input_ids = torch.nested.to_padded_tensor( + input_ids, padding=pad_token_id, output_size=(batch_size, max_seq_len) + ) + + if position_ids.dim() == 3: + position_ids = torch.nested.to_padded_tensor( + position_ids, padding=0, output_size=(batch_size, 4, max_seq_len) + ).transpose(0, 1) # (4, batch_size, max_seq_len) + else: + position_ids = torch.nested.to_padded_tensor( + position_ids, padding=0, output_size=(batch_size, max_seq_len) + ) + + attention_mask = build_attention_mask_from_nested( + input_ids=micro_batch["input_ids"], max_seq_len=max_seq_len + ) + + model_inputs = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "position_ids": position_ids, + } + + else: + raise NotImplementedError(f"pad_mode {pad_mode} not implemented") + + extra_args = {} + if use_fused_kernels: + extra_args["temperature"] = temperature_item + extra_args["return_dict"] = True + if use_remove_padding: + # We have already computed `input_ids_rmpad_rolled` from the *full* + # global sequence and (when SP>1) SP-sliced it. Pass it into the model + # so the fused forward uses these labels verbatim instead of redoing + # `torch.roll` on the local SP shard, which would wrap around the + # shard boundary rather than the global sequence (issue #6068). This + # mirrors what the veomni engine already does for fused kernels. + extra_args["shift_labels"] = output_args["input_ids_rmpad_rolled"].unsqueeze(0) + + model_inputs.update(multi_modal_inputs) + model_inputs.update(extra_args) + + return model_inputs, output_args + + def prepare_model_outputs(self, output, output_args, micro_batch: TensorDict, logits_processor_func): + use_remove_padding = tu.get_non_tensor_data(data=micro_batch, key="use_remove_padding", default=True) + pad_mode = tu.get_non_tensor_data(data=micro_batch, key="pad_mode", default=DatasetPadMode.NO_PADDING) + use_fused_kernels = tu.get_non_tensor_data(data=micro_batch, key="use_fused_kernels", default=False) + calculate_entropy = tu.get_non_tensor_data(data=micro_batch, key="calculate_entropy", default=False) + calculate_sum_pi_squared = tu.get_non_tensor_data( + data=micro_batch, key="calculate_sum_pi_squared", default=False + ) + distillation_use_topk = tu.get_non_tensor_data(data=micro_batch, key="distillation_use_topk", default=False) + + if calculate_sum_pi_squared and use_fused_kernels: + raise NotImplementedError( + "calculate_sum_pi_squared=True is not supported with use_fused_kernels=True: " + "fused kernels do not materialize the full logits tensor needed for Σπ²." + ) + + model_output = {} + + input_ids = micro_batch["input_ids"] + + if use_remove_padding: + input_ids_rmpad_rolled = output_args["input_ids_rmpad_rolled"] + temperature_rmpad = output_args["temperature_rmpad"] + + if use_fused_kernels: + # temperature is singleton + log_probs = output.log_probs.squeeze(0) # (total_nnz,) + entropy_rmpad = output.entropy.squeeze(0) # (total_nnz,) + else: + logits_rmpad = output.logits.squeeze(0) # (total_nnz, vocab_size) + logits_rmpad.div_(temperature_rmpad.clamp(min=1e-8).unsqueeze(-1).to(logits_rmpad.dtype)) + + # if use_sp: ((total_nnz / sp) + pad) ; if not use_sp: (batch, seqlen) + inplace_backward = True + if calculate_entropy: + inplace_backward = False + log_probs = logprobs_from_logits( + logits=logits_rmpad, + labels=input_ids_rmpad_rolled, + inplace_backward=inplace_backward, + ) + + # compute entropy + if calculate_entropy: + if not self.engine_config.entropy_checkpointing: + entropy_rmpad = self.compute_entropy_from_logits(logits_rmpad) # ((total_nnz / sp) + pad) + else: + entropy_rmpad = torch.utils.checkpoint.checkpoint( + self.compute_entropy_from_logits, logits_rmpad + ) + + # compute sum_pi_squared (Σπ²) for optimal-baseline advantage estimators + if calculate_sum_pi_squared: + sum_pi_squared_rmpad = verl_F.calculate_sum_pi_squared_from_logits(logits_rmpad) + + # logits_processor_func return tensors with shape (1, total_nnz/sp_size) + if distillation_use_topk: + outputs = logits_processor_func(student_logits=logits_rmpad.unsqueeze(0), data=micro_batch) + cu_seqlens = input_ids.offsets() + for k, v in outputs.items(): + v = v.squeeze(0) + assert v.shape == log_probs.shape, f"log_probs shape: {log_probs.shape}, {k} shape: {v.shape}" + if self.use_ulysses_sp: + pad_size = output_args["pad_size"] + v = gather_outputs_and_unpad(v, gather_dim=0, unpad_dim=0, padding_size=pad_size) + model_output[k] = torch.nested.nested_tensor_from_jagged(v, cu_seqlens) + + # gather log_prob if sp > 1 + if self.use_ulysses_sp: + pad_size = output_args["pad_size"] + + # gather and unpad for the ulysses sp + log_probs = gather_outputs_and_unpad( + log_probs, + gather_dim=0, + unpad_dim=0, + padding_size=pad_size, + ) + if calculate_entropy: + entropy_rmpad = gather_outputs_and_unpad( + entropy_rmpad, + gather_dim=0, + unpad_dim=0, + padding_size=pad_size, + ) + if calculate_sum_pi_squared: + sum_pi_squared_rmpad = gather_outputs_and_unpad( + sum_pi_squared_rmpad, + gather_dim=0, + unpad_dim=0, + padding_size=pad_size, + ) + + if pad_mode == DatasetPadMode.NO_PADDING: + cu_seqlens = input_ids.offsets() + # (bsz, j1), for each sample, is the length of each sample: [real_prompt length + real_response length] + log_probs = torch.nested.nested_tensor_from_jagged(log_probs, cu_seqlens) + if calculate_entropy: + entropy = torch.nested.nested_tensor_from_jagged(entropy_rmpad, cu_seqlens) + if calculate_sum_pi_squared: + sum_pi_squared = torch.nested.nested_tensor_from_jagged(sum_pi_squared_rmpad, cu_seqlens) + else: + raise NotImplementedError(f"pad_mode {pad_mode} not implemented") + + else: # not using rmpad and no ulysses sp + response_length = tu.get_non_tensor_data(data=micro_batch, key="max_response_length", default=1024) + if use_fused_kernels: + log_probs = output.log_probs[:, -response_length - 1 : -1] + entropy = output.entropy[:, -response_length - 1 : -1] # (bsz, response_length) + + else: + logits = output.logits # (bsz, response_length, vocab_size) + temperature = output_args["temperature"] # (bsz,) + temperature = temperature.unsqueeze(-1).unsqueeze(-1) + logits.div_(temperature.clamp(min=1e-8).to(logits.dtype)) + + if calculate_entropy: + if not self.engine_config.entropy_checkpointing: + entropy = verl_F.entropy_from_logits(logits) + else: + entropy = torch.utils.checkpoint.checkpoint(verl_F.entropy_from_logits, logits) + + if calculate_sum_pi_squared: + sum_pi_squared = verl_F.calculate_sum_pi_squared_from_logits(logits) + + if pad_mode == DatasetPadMode.NO_PADDING: + cu_seqlens = input_ids.offsets() + seq_lengths = cu_seqlens.diff() + starts = torch.zeros_like(seq_lengths, dtype=torch.int64) + logits = torch.nested.narrow(logits, 1, starts, seq_lengths, layout=torch.jagged) + logits_rmpad = torch.cat([t for t in logits.unbind()]) + input_ids_rmpad_rolled = output_args["input_ids_rmpad_rolled"] + log_probs = logprobs_from_logits(logits=logits_rmpad, labels=input_ids_rmpad_rolled) + + # Mirror the use_remove_padding=True branch (see verl#6293). + # No Ulysses SP gather here: this branch is the no-SP path + # (log_probs is also not gathered) and pad_size is only + # populated in output_args along the use_remove_padding=True + # path of prepare_model_inputs. + if distillation_use_topk: + outputs = logits_processor_func(student_logits=logits_rmpad.unsqueeze(0), data=micro_batch) + for k, v in outputs.items(): + v = v.squeeze(0) + assert v.shape == log_probs.shape, ( + f"log_probs shape: {log_probs.shape}, {k} shape: {v.shape}" + ) + model_output[k] = torch.nested.nested_tensor_from_jagged(v, cu_seqlens) + + # (bsz, j1), for each sample, length of each sample: [real_prompt_length + real_response_length] + log_probs = torch.nested.nested_tensor_from_jagged(log_probs, cu_seqlens) + if calculate_entropy: + entropy = torch.nested.narrow(entropy, 1, starts, seq_lengths, layout=torch.jagged) + entropy_rmpad = torch.cat([t for t in entropy.unbind()]) + entropy = torch.nested.nested_tensor_from_jagged(entropy_rmpad, cu_seqlens) + if calculate_sum_pi_squared: + sum_pi_squared = torch.nested.narrow( + sum_pi_squared, 1, starts, seq_lengths, layout=torch.jagged + ) + sum_pi_squared_rmpad = torch.cat([t for t in sum_pi_squared.unbind()]) + sum_pi_squared = torch.nested.nested_tensor_from_jagged(sum_pi_squared_rmpad, cu_seqlens) + else: + raise NotImplementedError(f"pad_mode {pad_mode} not implemented") + + model_output["log_probs"] = log_probs + if calculate_entropy: + model_output["entropy"] = entropy + if calculate_sum_pi_squared: + model_output["sum_pi_squared"] = sum_pi_squared + + return model_output + + def forward_step(self, micro_batch: TensorDict, loss_function, forward_only): + device_name = get_device_name() + # actually, we should avoid assigning like this... + micro_batch = micro_batch.to(get_device_id()) + model_inputs, output_args = self.prepare_model_inputs(micro_batch=micro_batch) + + # Honor mixed_precision.param_dtype resolved during FSDP setup. When dtype is fp32, + # autocast is a no-op at best and a footgun at worst, so skip it entirely. + # getattr fallback: some subclasses (e.g. VeOmniEngine) bypass FSDPEngine.__init__ + # and _build_fsdp_module, so self._autocast_dtype may not be set. + autocast_dtype = getattr(self, "_autocast_dtype", torch.bfloat16) + autocast_ctx: ContextManager = ( + nullcontext() + if autocast_dtype == torch.float32 + else torch.autocast(device_type=device_name, dtype=autocast_dtype) + ) + with autocast_ctx: + raw_output = self.module( + **model_inputs, + use_cache=False, + ) # prevent model thinks we are generating + + model_output = self.prepare_model_outputs( + output=raw_output, output_args=output_args, micro_batch=micro_batch, logits_processor_func=loss_function + ) + + if loss_function is not None: + loss, metrics = loss_function( + model_output=model_output, data=micro_batch, dp_group=self.get_data_parallel_group() + ) + else: + assert forward_only, "forward_only must be True when loss_function is None" + loss = torch.tensor(1.0, device=device_name) + metrics = {} + + output = { + "model_output": model_output, + "loss": loss.detach().item(), + "metrics": metrics, + } + + return loss, output + + +@EngineRegistry.register(model_type="value_model", backend=["fsdp", "fsdp2"], device=["cuda", "npu"]) +class FSDPEngineWithValueHead(FSDPEngineWithLMHead): + """ + The only difference between critic and actor is how the raw model output is processed + """ + + def prepare_model_outputs(self, output, output_args, micro_batch: TensorDict, logits_processor_func): + use_remove_padding = tu.get_non_tensor_data(data=micro_batch, key="use_remove_padding", default=True) + pad_mode = tu.get_non_tensor_data(data=micro_batch, key="pad_mode", default=DatasetPadMode.NO_PADDING) + + input_ids = micro_batch["input_ids"] + if use_remove_padding: + if hasattr(self.module, "v_head"): + # For trl.AutoModelForCausalLMWithValueHead + values_rmpad = output[2].squeeze(0) + else: + values_rmpad = output.logits + values_rmpad = values_rmpad.squeeze(0) # (total_nnz, 1) + # critic model arch is like Qwen3ForTokenClassfication and num_labels=1 + # so we squeeze the last dimension here to get the value for each token + values_rmpad = values_rmpad.squeeze(-1) + + # gather output if sp > 1 + if self.use_ulysses_sp: + pad_size = output_args["pad_size"] + values_rmpad = gather_outputs_and_unpad(values_rmpad, gather_dim=0, unpad_dim=0, padding_size=pad_size) + + if pad_mode == DatasetPadMode.NO_PADDING: + cu_seqlens = input_ids.offsets() + # (bsz, j1), for each sample, is the length of each sample: [real_prompt length + real_response length] + values = torch.nested.nested_tensor_from_jagged(values_rmpad, cu_seqlens) + else: + raise NotImplementedError(f"pad_mode {pad_mode} not implemented") + + else: + if hasattr(self.module, "v_head"): + # For trl.AutoModelForCausalLMWithValueHead + values = output[2] + else: + values = output.logits.squeeze(-1) + + if pad_mode == DatasetPadMode.NO_PADDING: + cu_seqlens = input_ids.offsets() + seq_lengths = cu_seqlens.diff() + starts = torch.zeros_like(seq_lengths, dtype=torch.int64) + values = torch.nested.narrow(values, 1, starts, seq_lengths, layout=torch.jagged) + values_rmpad = torch.cat([t for t in values.unbind()]) + # (bsz, j1), for each sample, length of each sample: [real_prompt_length + real_response_length] + values = torch.nested.nested_tensor_from_jagged(values_rmpad, cu_seqlens) + else: + raise NotImplementedError(f"pad_mode {pad_mode} not implemented") + + return {"values": values} diff --git a/verl/verl/workers/engine/fsdp/utils.py b/verl/verl/workers/engine/fsdp/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..6336bd8794cc915af0aa03380fd92a5fecd25795 --- /dev/null +++ b/verl/verl/workers/engine/fsdp/utils.py @@ -0,0 +1,80 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import logging +import os + +import torch +from torch.distributed.device_mesh import init_device_mesh + +from verl.utils.device import get_device_name, is_npu_available + +logger = logging.getLogger(__file__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + + +def apply_npu_fsdp_patches(): + """Apply NPU patches for FSDP backend if NPU is available.""" + if is_npu_available: + try: + import verl.models.transformers.npu_patch # noqa + + if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0: + logger.info("Applied NPU patches for FSDP backend") + except Exception as e: + logger.warning(f"Failed to apply NPU patches: {e}") + + +def create_device_mesh(world_size, fsdp_size): + """ + Create a device mesh for distributed training based on the world size and FSDP size. + + Args: + world_size (int): Total number of processes in the distributed training setup. + fsdp_size (int): Size of the Fully Sharded Data Parallel (FSDP) group. + + Returns: + torch.distributed.device_mesh.DeviceMesh: The initialized device mesh. + """ + device_name = get_device_name() + if fsdp_size < 0 or fsdp_size >= world_size: + device_mesh = init_device_mesh(device_name, mesh_shape=(world_size,), mesh_dim_names=["fsdp"]) + else: + device_mesh = init_device_mesh( + device_name, mesh_shape=(world_size // fsdp_size, fsdp_size), mesh_dim_names=["ddp", "fsdp"] + ) + return device_mesh + + +def get_sharding_strategy(device_mesh): + """ + Determine the appropriate sharding strategy based on the number of dimensions of the device mesh. + + Args: + device_mesh (torch.distributed.device_mesh.DeviceMesh): The device mesh used for distributed training. + + Returns: + torch.distributed.fsdp.ShardingStrategy: The sharding strategy to be used with FSDP. + + Raises: + NotImplementedError: If the number of dimensions of the device mesh is neither 1 nor 2. + """ + from torch.distributed.fsdp import ShardingStrategy + + if device_mesh.ndim == 1: + sharding_strategy = ShardingStrategy.FULL_SHARD + elif device_mesh.ndim == 2: + sharding_strategy = ShardingStrategy.HYBRID_SHARD + else: + raise NotImplementedError(f"Get device mesh ndim={device_mesh.ndim}, but only support 1 or 2") + return sharding_strategy diff --git a/verl/verl/workers/engine/megatron/__init__.py b/verl/verl/workers/engine/megatron/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..51d3a1949e9170f5ff2d0997b55076d97b0c4918 --- /dev/null +++ b/verl/verl/workers/engine/megatron/__init__.py @@ -0,0 +1,29 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +# HACK Avoid cpu worker trigger cuda jit error +import os + +from verl.utils.device import is_cuda_available + +if not is_cuda_available and "TORCH_CUDA_ARCH_LIST" not in os.environ: + os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0" + +from .transformer_impl import MegatronEngine, MegatronEngineWithLMHead, MegatronEngineWithValueHead # noqa: E402 + +if not is_cuda_available: + del os.environ["TORCH_CUDA_ARCH_LIST"] + +__all__ = ["MegatronEngine", "MegatronEngineWithLMHead", "MegatronEngineWithValueHead"] diff --git a/verl/verl/workers/engine/megatron/__pycache__/__init__.cpython-312.pyc b/verl/verl/workers/engine/megatron/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c547597c7146f64a884fb391fccf1fdd3772ab96 Binary files /dev/null and b/verl/verl/workers/engine/megatron/__pycache__/__init__.cpython-312.pyc differ diff --git a/verl/verl/workers/engine/megatron/__pycache__/transformer_impl.cpython-312.pyc b/verl/verl/workers/engine/megatron/__pycache__/transformer_impl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..24408d09299134635afd1d247b06486d85ff2305 Binary files /dev/null and b/verl/verl/workers/engine/megatron/__pycache__/transformer_impl.cpython-312.pyc differ diff --git a/verl/verl/workers/engine/megatron/transformer_impl.py b/verl/verl/workers/engine/megatron/transformer_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..eeccb2083362a21390c27872d3857bf92a656605 --- /dev/null +++ b/verl/verl/workers/engine/megatron/transformer_impl.py @@ -0,0 +1,1031 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +import logging +import os +from functools import partial +from typing import Any, Callable, ContextManager, Iterator, Optional + +import torch +import torch.distributed +from megatron.core import parallel_state as mpu +from megatron.core.pipeline_parallel import get_forward_backward_func +from omegaconf import OmegaConf +from tensordict import TensorDict + +import verl.utils.torch_functional as verl_F +from verl.models.mcore import get_mcore_weight_converter +from verl.trainer.config import CheckpointConfig +from verl.utils import tensordict_utils as tu +from verl.utils.checkpoint.megatron_checkpoint_manager import MegatronCheckpointManager +from verl.utils.dataset.dataset_utils import DatasetPadMode +from verl.utils.debug import log_gpu_memory_usage +from verl.utils.device import get_device_id, get_device_name +from verl.utils.megatron.pipeline_parallel import make_batch_generator +from verl.utils.megatron.router_replay_patch import RouterReplay, RouterReplayAction, apply_router_replay_patch +from verl.utils.megatron.router_replay_utils import ( + RouterReplayHelper, + merge_router_topk_indices, + pp_gather, + reorder_and_merge_vpp_layers, + set_router_replay_data, +) +from verl.utils.megatron.tensor_parallel import ( + vocab_parallel_entropy, + vocab_parallel_log_probs_from_logits, + vocab_parallel_sum_pi_squared, +) +from verl.utils.megatron_peft_utils import add_base_layer_suffix, build_peft_config_for_vllm +from verl.utils.megatron_utils import ( + check_mtp_config, + get_megatron_module_device, + get_megatron_mtp_loss, + load_megatron_model_to_gpu, + load_megatron_optimizer, + offload_megatron_model_to_cpu, + offload_megatron_optimizer, + patch_engine_mtp, + register_megatron_training_hooks, + unwrap_model, +) +from verl.utils.model import extract_multi_modal_inputs, load_mcore_dist_weights +from verl.utils.seqlen_balancing import restore_dynamic_batch +from verl.workers.config import HFModelConfig, McoreEngineConfig, McoreOptimizerConfig + +from ..base import BaseEngine, BaseEngineCtx, EngineRegistry +from ..utils import postprocess_batch_func, prepare_micro_batches +from .utils import set_random_seed + +logger = logging.getLogger(__file__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + + +class MegatronEngine(BaseEngine): + def __init__( + self, + model_config: HFModelConfig, + engine_config: McoreEngineConfig, + optimizer_config: McoreOptimizerConfig, + checkpoint_config: CheckpointConfig, + ): + super().__init__() + + self.model_config = model_config + self.engine_config = engine_config + self.optimizer_config = optimizer_config + self.checkpoint_config = checkpoint_config + assert self.engine_config.use_mbridge, "use_mbridge must be True" + self._init_device_mesh() + + set_random_seed(seed=self.engine_config.seed) + + self._is_offload_param = self.engine_config.param_offload + self._is_offload_grad = self.engine_config.grad_offload + self._is_offload_optimizer = self.engine_config.optimizer_offload + + self.mode = None + + self.layer_name_mapping = { + "qkv_layer_name": "self_attention.linear_qkv.", + "gate_proj_layer_name": "linear_fc1.", + } + self.weight_converter = None + + # QAT configuration + self._qat_config = getattr(self.engine_config, "qat", None) + self._qat_enabled = self._qat_config is not None and getattr(self._qat_config, "enable", False) + if self._qat_enabled: + if self.engine_config.vanilla_mbridge: + raise ValueError( + "QAT requires non-vanilla Megatron bridge. " + "Please set 'use_mbridge=True' and 'vanilla_mbridge=False'." + ) + logger.info(f"QAT enabled in MegatronEngine: mode={self._qat_config.mode}") + + # Router replay configuration for MoE models + self.enable_routing_replay = self.engine_config.router_replay.mode != "disabled" + logger.info(f"enable_routing_replay in MegatronEngine: {self.enable_routing_replay}") + if self.enable_routing_replay: + apply_router_replay_patch() + self.mini_layer_topk_idx_list = [] + # Apply checkpoint patch for MoE models + from verl.utils.device import is_cuda_available + + if is_cuda_available: + from verl.models.mcore.patch import apply_patch_megatron_recomputation_backward + + apply_patch_megatron_recomputation_backward() + + def _init_device_mesh(self): + # TODO: set different parallelism for actor, critic, ref + if mpu.is_initialized(): + return + + extra_args = dict() + + if self.engine_config.dynamic_context_parallel: + assert "dynamic_context_parallel" in inspect.signature(mpu.initialize_model_parallel).parameters, ( + "dynamic_context_parallel is not supported in your megatron version, " + + "please update your megatron version to the latest version" + ) + assert self.engine_config.max_seqlen_per_dp_cp_rank is not None, ( + "max_seqlen_per_dp_cp_rank is required when dynamic_context_parallel is enabled" + ) + extra_args["dynamic_context_parallel"] = self.engine_config.dynamic_context_parallel + + mpu.initialize_model_parallel( + tensor_model_parallel_size=self.engine_config.tensor_model_parallel_size, + pipeline_model_parallel_size=self.engine_config.pipeline_model_parallel_size, + virtual_pipeline_model_parallel_size=self.engine_config.virtual_pipeline_model_parallel_size, + use_sharp=False, + context_parallel_size=self.engine_config.context_parallel_size, + expert_model_parallel_size=self.engine_config.expert_model_parallel_size, + expert_tensor_parallel_size=self.engine_config.expert_tensor_parallel_size, + nccl_communicator_config_path=None, + **extra_args, + ) + + def _build_tf_config(self): + from verl.utils.megatron_utils import mapping_string_to_attn_backend + from verl.utils.torch_dtypes import PrecisionType + + check_mtp_config(self.model_config, self.engine_config) + + self.param_dtype = PrecisionType.to_dtype(self.engine_config.dtype) + self.dtype = PrecisionType.to_dtype(self.param_dtype) + + override_transformer_config = mapping_string_to_attn_backend({**self.engine_config.override_transformer_config}) + + self.provider = None + self.vanilla_bridge = self.engine_config.vanilla_mbridge + + if self.vanilla_bridge: + from verl.models.mcore.mbridge import AutoBridge + + bridge = AutoBridge.from_config(self.model_config.hf_config, dtype=self.param_dtype) + if self.engine_config.dynamic_context_parallel: + override_transformer_config["max_seqlen_per_dp_cp_rank"] = self.engine_config.max_seqlen_per_dp_cp_rank + # note(baiyan): we must set the transformer_config.dynamic_context_parallel to False + # because of the bad coupling design in Megatron-LM + # https://github.com/xiaoyao0115/Megatron-LM/blob/88733ab6614e3e91b9d095172f41e7d8b5d8e9d4/megatron/core/pipeline_parallel/dynamic_cp_schedule.py#L552-L553 + # but it does not affect the functionality of dynamic CP, so we can use it to avoid the coupling. + override_transformer_config["dynamic_context_parallel"] = False + override_transformer_config["context_parallel_size"] = mpu.get_data_parallel_world_size() + bridge.set_extra_args(**override_transformer_config) + tf_config = bridge.config + tf_config.fp16 = self.param_dtype == torch.float16 + tf_config.bf16 = self.param_dtype == torch.bfloat16 + else: + from verl.models.mcore.bridge import AutoBridge + + # Use Megatron-Bridge to convert HF config to Megatron config + bridge = AutoBridge.from_hf_pretrained( + self.model_config.local_path, trust_remote_code=self.model_config.trust_remote_code + ) + # Get Megatron provider and configure it + provider = bridge.to_megatron_provider(load_weights=False) + + # In case of invalid overrides, we need to make sure some critical params are set correctly + provider.params_dtype = self.param_dtype + + # Ensure dtype settings propagate to Megatron-Bridge/TE + provider.fp16 = self.param_dtype == torch.float16 + provider.bf16 = self.param_dtype == torch.bfloat16 + + # Pass distributed info + provider.tensor_model_parallel_size = self.engine_config.tensor_model_parallel_size + provider.pipeline_model_parallel_size = self.engine_config.pipeline_model_parallel_size + provider.expert_model_parallel_size = self.engine_config.expert_model_parallel_size + provider.expert_tensor_parallel_size = self.engine_config.expert_tensor_parallel_size + provider.virtual_pipeline_model_parallel_size = self.engine_config.virtual_pipeline_model_parallel_size + provider.context_parallel_size = self.engine_config.context_parallel_size + provider.sequence_parallel = self.engine_config.sequence_parallel + + # Match verl implementation (need variable_seq_lengths) + from megatron.core.transformer.enums import AttnBackend + + provider.attention_backend = AttnBackend.flash + provider.variable_seq_lengths = True + provider.moe_token_dispatcher_type = "alltoall" + provider.moe_router_load_balancing_type = "none" + + # Apply QAT: set quantization layer spec and patch Megatron-Bridge + if self._qat_enabled: + from verl.utils.modelopt import patch_provider_for_qat + + patch_provider_for_qat(provider) + + # Apply transformer config overrides + for key, value in override_transformer_config.items(): + setattr(provider, key, value) + + if self.enable_routing_replay: + provider.enable_routing_replay = True + + provider.finalize() + self.provider = provider + tf_config = None # Will be set after model creation + self.bridge = bridge + + if not self.bridge: + self.weight_converter = get_mcore_weight_converter(self.model_config.hf_config, self.dtype) + + # Set enable_routing_replay directly on tf_config instead of passing through + # override_transformer_config, because dataclass subclasses like MLATransformerConfig + # generate their own __init__ and don't inherit the patched TransformerConfig.__init__ + # that accepts this kwarg. + if self.enable_routing_replay and tf_config is not None: + tf_config.enable_routing_replay = True + + if torch.distributed.get_rank() == 0: + if tf_config is not None: + print(f"TF config: {tf_config}") + self.tf_config = tf_config + + from verl.workers.config.megatron_peft import get_peft_cls + + self.peft_cls = get_peft_cls( + model_config=self.model_config, bridge=self.bridge, provider=self.provider, dtype=self.param_dtype + ) + + def _build_megatron_module(self): + from verl.utils.megatron_utils import McoreModuleWrapperConfig, make_megatron_module + from verl.utils.model import print_model_size + + self.is_value_model = self.model_config.model_type == "value_model" + if self.engine_config.forward_only: + wrap_with_ddp = False + else: + wrap_with_ddp = True + + wrap_config = McoreModuleWrapperConfig( + is_value_model=self.is_value_model, + wrap_with_ddp=wrap_with_ddp, + use_distributed_optimizer=self.engine_config.use_distributed_optimizer, + use_megatron_fsdp=self.engine_config.use_megatron_fsdp, + ) + if self.is_value_model: + self.model_config.hf_config.tie_word_embeddings = False + + module, updated_tf_config = make_megatron_module( + wrap_config=wrap_config, + tf_config=self.tf_config, + hf_config=self.model_config.hf_config, + bridge=self.bridge, + provider=self.provider, + override_model_config=self.engine_config.override_mcore_model_config, + override_ddp_config=self.engine_config.override_ddp_config, + peft_cls=self.peft_cls, + peft_config=self.model_config.get("lora", None), + ) + self.tf_config = updated_tf_config + print(f"module: {len(module)}") + + if self.engine_config.use_dist_checkpointing: + load_mcore_dist_weights( + module, self.engine_config.dist_checkpointing_path, is_value_model=self.is_value_model + ) + else: + if self.vanilla_bridge: + self.bridge.load_weights(module, self.model_config.local_path) + else: + allowed_mismatched_params = [] + if self.is_value_model: + allowed_mismatched_params = ["output_layer.weight"] + self.bridge.load_hf_weights( + module, self.model_config.local_path, allowed_mismatched_params=allowed_mismatched_params + ) + + if torch.distributed.get_rank() == 0: + print_model_size(module[0]) + + if self.enable_routing_replay: + print(f"routing replay layers: {len(RouterReplay.router_instances)}") + + return module + + def _maybe_enable_fused_kernels(self): + if not self.engine_config.use_fused_kernels: + return + + if self.is_value_model or self.model_config.mtp.enable: + logger.warning_once( + "Fused kernels are not supported for value models or when MTP is enabled in Megatron engine; disabling." + ) + self.engine_config.use_fused_kernels = False + return + + from verl.models.mcore.model_forward_fused import patch_fused_forward + + for model in self.module: + patch_fused_forward(model) + + def _build_optimizer(self): + from verl.utils.megatron.optimizer import get_megatron_optimizer, init_megatron_optim_config + + optim_config_megatron = init_megatron_optim_config( + self.optimizer_config, + use_distributed_optimizer=self.engine_config.use_distributed_optimizer, + fp16=self.param_dtype == torch.float16, + ) + optimizer = get_megatron_optimizer(model=self.module, config=optim_config_megatron) + register_megatron_training_hooks(self.module, optimizer) + return optimizer + + def _build_lr_scheduler(self): + from verl.utils.megatron.optimizer import get_megatron_optimizer_param_scheduler + + optimizer_scheduler = get_megatron_optimizer_param_scheduler( + optimizer=self.optimizer, config=self.optimizer_config + ) + return optimizer_scheduler + + @property + def is_param_offload_enabled(self) -> bool: + return self._is_offload_param + + @property + def is_optimizer_offload_enabled(self) -> bool: + return self._is_offload_optimizer + + def is_mp_src_rank_with_outputs(self): + return ( + mpu.get_tensor_model_parallel_rank() == 0 + and mpu.get_pipeline_model_parallel_rank() == mpu.get_pipeline_model_parallel_world_size() - 1 + and mpu.get_context_parallel_rank() == 0 + ) + + def initialize(self): + self._build_tf_config() + + self.module = self._build_megatron_module() + + if self._qat_enabled and not self.engine_config.forward_only: + from verl.utils.modelopt import apply_qat_to_modules + + self.module = apply_qat_to_modules(self.module, self._qat_config) + + self._maybe_enable_fused_kernels() + + if self.model_config.mtp.enable: + patch_engine_mtp(self.module, self.model_config) + + # For forward_only, we don't need optimizer, lr_scheduler, checkpoint_mananager + if self.engine_config.forward_only: + self.optimizer = None + self.lr_scheduler = None + self.to(device="cpu", model=self._is_offload_param, optimizer=False, grad=False) + log_gpu_memory_usage("After offload model during init (forward_only)", logger=logger) + return + + self.optimizer = self._build_optimizer() + self.lr_scheduler = self._build_lr_scheduler() + + full_reshardable = self.engine_config.dist_ckpt_optim_fully_reshardable + mem_eff = self.engine_config.distrib_optim_fully_reshardable_mem_efficient + + tmp_config = OmegaConf.create( + { + "model": {"path": self.model_config.local_path}, + "megatron": { + "dist_ckpt_optim_fully_reshardable": full_reshardable, + "distrib_optim_fully_reshardable_mem_efficient": mem_eff, + }, + } + ) + + role = "actor" if not self.is_value_model else "critic" + + self.checkpoint_mananager = MegatronCheckpointManager( + config=tmp_config, + checkpoint_config=self.checkpoint_config, + model_config=self.model_config.hf_config, + transformer_config=self.tf_config, + role=role, + model=self.module, + arch=self.model_config.architectures[0], + hf_config=self.model_config.hf_config, + param_dtype=self.param_dtype, + share_embeddings_and_output_weights=self.model_config.share_embeddings_and_output_weights, + processing_class=self.model_config.get_processor(), + optimizer=self.optimizer, + optimizer_scheduler=self.lr_scheduler, + use_distributed_optimizer=self.engine_config.use_distributed_optimizer, + use_checkpoint_opt_param_scheduler=self.optimizer_config.use_checkpoint_opt_param_scheduler, + bridge=self.bridge, + provider=self.provider, + peft_cls=self.peft_cls, + use_dist_checkpointing=self.engine_config.use_dist_checkpointing, + use_megatron_fsdp=self.engine_config.use_megatron_fsdp, + ) + + self.to( + device="cpu", + model=self._is_offload_param, + optimizer=self._is_offload_optimizer, + grad=self._is_offload_param, + ) + + log_gpu_memory_usage("After offload model/optimizer/grad during init", logger=logger) + + def train_mode(self, **kwargs): + """ + Context manager entry for switching the engine and model into training mode. + + Usage: + with engine.train_mode(): + # runs in training mode + """ + return EngineTrainModeCtx(self, **kwargs) + + def eval_mode(self, **kwargs): + """ + Context manager entry for switching the engine and model into evaluation mode. + + Usage: + with engine.eval_mode(): + # runs in evaluation mode + """ + return EngineEvalModeCtx(self, **kwargs) + + def optimizer_zero_grad(self): + """ + Zero out gradients of all parameters before starting a new backward pass. + """ + self.optimizer.zero_grad() + # use use_contiguous_buffers_in_local_ddp and no overlap_dp_param_comm + for chunk in self.module: + # if use distributed optimizer, zero grad buffer will be handled by optimizer + chunk.zero_grad_buffer() + + def optimizer_step(self): + """ + Perform an optimization step to update model parameters based on accumulated gradients. + + Returns: + grad_norm (float): The norm of the gradients before clipping or update. + """ + update_successful, grad_norm, num_zeros_in_grad = self.optimizer.step() + + if update_successful: + # allgather already execute in optimizer.step in new megatron + pass + else: + raise NotImplementedError("Megatron optimizer step failed. This should not happen") + + return grad_norm + + def lr_scheduler_step(self): + """ + Advance the learning rate scheduler by one step. + + Returns: + current_lr (float or list[float]): Updated learning rate(s). + """ + from verl.utils.megatron.optimizer import get_megatron_last_lr + + self.lr_scheduler.step(1) + return get_megatron_last_lr(self.optimizer) + + def to(self, device: str, model: bool = True, optimizer: bool = True, grad: bool = True): + """ + Move model parameters, optimizer states, or both to the specified device. + Note that this function executes irrespective of offload config. It serves as manual control + + Args: + device: Target device identifier. + model: If True, move the model. + optimizer: If True, move the optimizer states. + """ + super().to(device=device, model=model, optimizer=optimizer, grad=grad) + + device_name = get_device_name() + + assert device in (device_name, "cpu") + if device == device_name: + if model: + load_megatron_model_to_gpu(self.module, load_grad=grad) + if optimizer and self.optimizer is not None: + load_megatron_optimizer(self.optimizer) + elif device == "cpu": + if model: + offload_megatron_model_to_cpu(self.module) + if optimizer and self.optimizer is not None: + offload_megatron_optimizer(self.optimizer) + else: + raise ValueError(f"Invalid device type: {device}") + + def get_data_parallel_rank(self): + if self.engine_config.dynamic_context_parallel: + # in order to let every dp-cp group has full data to split, we set dp=1 + return 0 + return mpu.get_data_parallel_rank() + + def get_data_parallel_size(self): + if self.engine_config.dynamic_context_parallel: + # in order to let every dp-cp group has full data to split, we set dp=1 + return 1 + return mpu.get_data_parallel_world_size() + + def get_data_parallel_group(self): + return mpu.get_data_parallel_group() + + def get_model_parallel_group(self): + return mpu.get_model_parallel_group() + + def get_context_parallel_group(self): + return mpu.get_context_parallel_group() + + def save_checkpoint( + self, + local_path: str, + hdfs_path: Optional[str] = None, + global_step: int = 0, + max_ckpt_to_keep: Optional[int] = None, + **kwargs, + ) -> None: + """ + Save model, optimizer, and scheduler states to a checkpoint. + + Args: + local_path: Local filesystem path to save checkpoint. + hdfs_path: Optional HDFS path to copy checkpoint. + global_step: Integer training step number for naming. + max_ckpt_to_keep: Maximum number of recent checkpoints to retain. + """ + origin_module_device = get_megatron_module_device(self.module) + if self._is_offload_param or origin_module_device == "cpu": + load_megatron_model_to_gpu(self.module, load_grad=True) + self.checkpoint_mananager.save_checkpoint( + local_path=local_path, hdfs_path=hdfs_path, global_step=global_step, max_ckpt_to_keep=max_ckpt_to_keep + ) + torch.distributed.barrier() + if self._is_offload_param: + offload_megatron_model_to_cpu(self.module) + + def load_checkpoint( + self, local_path: str, hdfs_path: Optional[str] = None, del_local_after_load: bool = True, **kwargs + ) -> None: + """ + Load model, optimizer, and scheduler states from a checkpoint. + + Args: + local_path: Local filesystem path of the checkpoint. + hdfs_path: Optional HDFS path where checkpoint is stored. + del_local_after_load: Whether to delete local copy after loading. + """ + if self._is_offload_param: + load_megatron_model_to_gpu(self.module) + self.checkpoint_mananager.load_checkpoint( + local_path=local_path, hdfs_path=hdfs_path, del_local_after_load=del_local_after_load + ) + if self._is_offload_param: + offload_megatron_model_to_cpu(self.module) + if self._is_offload_optimizer: + offload_megatron_optimizer(self.optimizer) + + def forward_backward_batch(self, data: TensorDict, loss_function: Callable, forward_only=False) -> Any: + tu.assign_non_tensor(data, sp_size=self.engine_config.context_parallel_size) + + # compute num_tokens in global batch for loss normalization + batch_num_tokens = data["loss_mask"].sum().to(get_device_id()) + torch.distributed.all_reduce( + batch_num_tokens, op=torch.distributed.ReduceOp.SUM, group=self.get_data_parallel_group() + ) + tu.assign_non_tensor(data, batch_num_tokens=batch_num_tokens.item()) + tu.assign_non_tensor(data, dp_size=self.get_data_parallel_size()) + + vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size() + if vpp_size is not None and vpp_size > 1: + num_batches_divided_by = self.tf_config.microbatch_group_size_per_vp_stage + else: + num_batches_divided_by = None + + micro_batches, indices = prepare_micro_batches( + data=data, + dp_group=self.get_data_parallel_group(), + num_batches_divided_by=num_batches_divided_by, + same_micro_num_in_dp=True, + min_num_micro_batch=None, + ) + + if num_batches_divided_by is not None: + assert len(micro_batches) % num_batches_divided_by == 0, ( + f"micro_batches {micro_batches} must be divisible by num_batches_divided_by " + f"{num_batches_divided_by} for megatron backend" + ) + + # compute input shapes for pp stages + n_micro_batch = len(micro_batches) + + for micro_batch in micro_batches: + tu.assign_non_tensor(micro_batch, num_micro_batch=n_micro_batch) + + forward_backward_func = get_forward_backward_func() + + postprocess_micro_batch_func = partial( + self.postprocess_micro_batch_func, + forward_only=forward_only, + loss_function=loss_function, + ) + + tu.assign_non_tensor(data, num_micro_batch=n_micro_batch) + + forward_step = partial( + self.forward_step, + logits_processor_func=loss_function, + postprocess_micro_batch_func=postprocess_micro_batch_func, + ) + + enable_routing_replay = tu.get_non_tensor_data(data, key="enable_routing_replay", default=False) + + if enable_routing_replay: + # Set to REPLAY mode: for R3 mode or actor update phase in R2 mode + RouterReplay.set_global_router_replay_action(RouterReplayAction.REPLAY_FORWARD) + if forward_only and self.engine_config.router_replay.mode == "R2": + # In R2 mode, forward_only calls (e.g., compute_log_probs) need to record routing information + RouterReplay.set_global_router_replay_action(RouterReplayAction.RECORD) + + # batch should be a list of batches inside micro-batches + batch_generator = make_batch_generator(micro_batches, vpp_size=len(self.module)) + + # TODO: we may use the new schedule instead + # for flash-attn: (seq_len, batch_size, hidden_size) = (mbs*seq_len, 1, hidden_size) + losses_reduced = forward_backward_func( + forward_step_func=forward_step, + data_iterator=batch_generator, + model=self.module, + num_microbatches=n_micro_batch, + seq_length=1, # the communication shape is obtained via p2p comm + micro_batch_size=1, # the communication shape is obtained via p2p comm + forward_only=forward_only, + ) + + if self.model_config.mtp.enable and mpu.is_pipeline_last_stage(ignore_virtual=True): + # All CP ranks must participate in the all_reduce inside get_megatron_mtp_loss, + # because save_loss_to_tracker uses avg_group=DP+CP group. + # Only collect metrics on the src rank afterward. + metrics = get_megatron_mtp_loss(n_micro_batch) + if self.is_mp_src_rank_with_outputs(): + if "metrics" not in losses_reduced[0]: + losses_reduced[0]["metrics"] = {} + losses_reduced[0]["metrics"].update(metrics) + + if RouterReplayHelper.is_r2_record_action(self.tf_config): + if self.tf_config.virtual_pipeline_model_parallel_size is not None: + # config = self.actor_module[0].module.module.config + vp_size = len(self.module) + microbatch_group_size_per_vp_stage = self.tf_config.microbatch_group_size_per_vp_stage + bs = n_micro_batch + topk_idx_td = reorder_and_merge_vpp_layers( + self.mini_layer_topk_idx_list, bs, vp_size, microbatch_group_size_per_vp_stage + ) + else: + tensors = [tensor for nt in self.mini_layer_topk_idx_list for tensor in nt.unbind()] + topk_idx_td = torch.nested.as_nested_tensor(tensors, layout=torch.jagged) + self.mini_layer_topk_idx_list = [] + + layers_topk_idx = pp_gather(topk_idx_td.to(torch.uint8), self.tf_config) + use_dynamic_bsz = tu.get_non_tensor_data(data=data, key="use_dynamic_bsz", default=True) + if use_dynamic_bsz and indices is not None: + layers_topk_idx = restore_dynamic_batch(layers_topk_idx, indices) + + output = {} + if mpu.is_pipeline_last_stage(ignore_virtual=True): + output = postprocess_batch_func(output_lst=losses_reduced, indices=indices, data=data) + if RouterReplayHelper.is_r2_record_action(self.tf_config): + output["model_output"]["routed_experts"] = layers_topk_idx + if enable_routing_replay: + RouterReplay.clear_global_indices() + RouterReplay.clear_global_router_replay_action() + return output + + def get_per_tensor_param(self, base_sync_done=False, **kwargs): + peft_config = None + non_merge_lora_sync = self.peft_cls is not None and not self.model_config.lora.get("merge", False) + adapter_only = base_sync_done and non_merge_lora_sync + if non_merge_lora_sync: + peft_config = build_peft_config_for_vllm(self.model_config.lora) + # when lora adapter only, we only load adapter weights when base sync is done, otherwise load all weights + load_megatron_model_to_gpu(self.module, load_grad=False, load_frozen_params=not adapter_only) + if self.vanilla_bridge: + per_tensor_param = self.bridge.export_weights(self.module) + elif adapter_only: + per_tensor_param = self.bridge.export_adapter_weights(self.module) + else: + per_tensor_param = ( + self.bridge.export_hf_weights(self.module, merge_adapter_weights=False) + if non_merge_lora_sync + else self.bridge.export_hf_weights(self.module) + ) + if non_merge_lora_sync: + per_tensor_param = add_base_layer_suffix( + per_tensor_param, model_type=self.model_config.hf_config.model_type + ) + + # QAT: process weights through QATWeightExporter for quantized weight sync to vLLM + if self._qat_enabled: + from verl.utils.modelopt import export_qat_weights + + per_tensor_param = export_qat_weights(per_tensor_param, self.module, self._qat_config.mode, self.bridge) + + return per_tensor_param, peft_config + + def disable_adapter(self) -> ContextManager: + return self.peft_cls.disable_adapter(self.module) + + def forward_step(self, batch_iter, model, logits_processor_func, postprocess_micro_batch_func): + raise NotImplementedError("forward_step must be implemented in subclass") + + def postprocess_micro_batch_func(self, output, data: TensorDict, forward_only: bool, loss_function): + raise NotImplementedError("postprocess_micro_batch_func must be implemented in subclass") + + +class EngineEvalModeCtx(BaseEngineCtx): + def __init__(self, engine: MegatronEngine, **kwargs): + super().__init__(engine=engine, mode="eval", **kwargs) + + def __enter__(self): + assert isinstance(self.engine, MegatronEngine) + super().__enter__() + # mcore module is a list of model chunk in each vpp stage + for module in self.engine.module: + module.eval() + + def __exit__(self, exc_type, exc_value, traceback): + assert isinstance(self.engine, MegatronEngine) + super().__exit__(exc_type, exc_value, traceback) + + +class EngineTrainModeCtx(BaseEngineCtx): + def __init__(self, engine: MegatronEngine, **kwargs): + super().__init__(engine=engine, mode="train", **kwargs) + + def __enter__(self): + assert isinstance(self.engine, MegatronEngine) + super().__enter__() + # mcore module is a list of model chunk in each vpp stage + for module in self.engine.module: + module.train() + + def __exit__(self, exc_type, exc_value, traceback): + assert isinstance(self.engine, MegatronEngine) + self.engine.optimizer_zero_grad() + super().__exit__(exc_type, exc_value, traceback) + + +@EngineRegistry.register(model_type="language_model", backend="megatron") +class MegatronEngineWithLMHead(MegatronEngine): + def prepare_model_inputs(self, batch: TensorDict): + input_ids = batch["input_ids"] + loss_mask = batch["loss_mask"].to(bool) + multi_modal_inputs = extract_multi_modal_inputs(batch.get("multi_modal_inputs", [])) + + routed_experts = batch.get("routed_experts", None) + + return { + "input_ids": input_ids, + "loss_mask": loss_mask, + "multi_modal_inputs": multi_modal_inputs, + "routed_experts": routed_experts, + } + + def prepare_model_outputs(self, output: dict, data: TensorDict): + return output + + def forward_step( + self, batch_iter: Iterator[TensorDict], model, logits_processor_func, postprocess_micro_batch_func + ): + batch: TensorDict = next(batch_iter) + + if self.engine_config.dynamic_context_parallel: + # split the batch and give the sub-batches to each dp-cp group + from verl.utils.megatron_utils import dynamic_cp_split_batch + + batch = dynamic_cp_split_batch( + batch=batch, + engine_config=self.engine_config, + dp_size=mpu.get_data_parallel_world_size(), + dp_rank=mpu.get_data_parallel_rank(), + ) + + batch = batch.to(get_device_id()) + use_fused_kernels = tu.get_non_tensor_data(batch, key="use_fused_kernels", default=False) + calculate_entropy = tu.get_non_tensor_data(batch, key="calculate_entropy", default=False) + calculate_sum_pi_squared = tu.get_non_tensor_data(batch, key="calculate_sum_pi_squared", default=False) + distillation_use_topk = tu.get_non_tensor_data(batch, key="distillation_use_topk", default=False) + + if calculate_sum_pi_squared and use_fused_kernels: + raise NotImplementedError( + "calculate_sum_pi_squared=True is not supported with use_fused_kernels=True: " + "fused kernels do not materialize the full logits tensor needed for Σπ²." + ) + pad_mode = tu.get_non_tensor_data(batch, key="pad_mode", default=DatasetPadMode.NO_PADDING) + temperature = batch["temperature"] + model_inputs = self.prepare_model_inputs(batch) + input_ids = model_inputs["input_ids"] + multi_modal_inputs = model_inputs["multi_modal_inputs"] + local_cp_size = tu.get_non_tensor_data(data=batch, key="local_cp_size", default=None) + loss_mask = model_inputs["loss_mask"] + + unwrapped_model = unwrap_model(model) + if hasattr(unwrapped_model, "vp_stage"): + vp_rank = unwrapped_model.vp_stage + else: + vp_rank = 0 + + if RouterReplayHelper.is_replay_backward_action(self.tf_config, vp_rank): + router_instance_list = RouterReplayHelper.get_micro_batch_router_list(self.tf_config, vp_rank) + for router in router_instance_list: + router.set_router_replay_action(RouterReplayAction.REPLAY_FORWARD) + + if RouterReplayHelper.is_replay_forward_action(self.tf_config, vp_rank): + layers_topk_idx = model_inputs["routed_experts"] + set_router_replay_data(layers_topk_idx, None, self.tf_config, vp_rank) + + if pad_mode == DatasetPadMode.NO_PADDING: + label = input_ids.clone() + else: + raise NotImplementedError(f"Pad mode {pad_mode} is not supported for megatron engine") + + if use_fused_kernels: + if not self.engine_config.use_remove_padding: + logger.warning_once( + "Fused kernels require `use_remove_padding=True` for Megatron engine. Falling back to non-fused." + ) + use_fused_kernels = False + elif isinstance(temperature, torch.Tensor): + if temperature.numel() != 1: + logger.warning_once( + "Fused kernels do not support per-sample temperature. Falling back to non-fused." + ) + use_fused_kernels = False + else: + temperature_value = float(temperature.item()) + else: + temperature_value = float(temperature) + + if use_fused_kernels: + from verl.models.mcore import get_mcore_forward_fused_model_engine_fn + + fused_forward_fn = get_mcore_forward_fused_model_engine_fn(self.model_config.hf_config) + output = fused_forward_fn( + model=model, + input_ids=input_ids, + labels=label, + multi_modal_inputs=multi_modal_inputs, + temperature=temperature_value, + calculate_entropy=calculate_entropy, + pad_token_id=self.model_config.tokenizer.pad_token_id, + ) + else: + if not isinstance(temperature, torch.Tensor): + temperature = torch.tensor([temperature] * input_ids.shape[0], device=input_ids.device) + + temperature = temperature.to(torch.float32) + assert temperature.shape[0] == input_ids.shape[0] + temperature = verl_F.expand_as_nested(temperature, input_ids) # (bsz, j1) + from verl.models.mcore import get_mcore_engine_forward_fn + + forward_fn = get_mcore_engine_forward_fn(self.model_config.hf_config) + data_format = "thd" if self.engine_config.use_remove_padding else "bshd" + + def logits_processor(logits, label, temperature): + assert logits.shape[:2] == label.shape[:2] + # avoid non-positive temperature such as padding + temperature[temperature <= 0] = 1e-8 + assert torch.all(temperature > 0).item(), f"temperature tensor must be positive. Got {temperature}" + logits.div_(temperature.unsqueeze(dim=-1).to(logits.dtype)) + ret = {} + # sum_pi_squared is non-destructive — must run before vocab_parallel_entropy. + if calculate_sum_pi_squared: + ret["sum_pi_squared"] = vocab_parallel_sum_pi_squared(logits) + if calculate_entropy: + logits_bak = logits.clone() + # # disable the hint until the fused_kernel is optimized for triton>=3.3 + # if torch.distributed.get_rank() == 0: + # logger.warning_once( + # "For memory-efficient computation, enable fused kernels via " + # "`actor_rollout_ref.model.use_fused_kernels=True`. " + # "The current `clone()` operation ensures correctness but increases memory usage." + # ) + entropy = vocab_parallel_entropy(logits) + ret["entropy"] = entropy + else: + logits_bak = logits + + # logits_processor_func return tensors with shape (1, total_nnz/cp_size) + if distillation_use_topk: + ret.update(logits_processor_func(student_logits=logits_bak, data=batch, data_format=data_format)) + log_probs = vocab_parallel_log_probs_from_logits(logits_bak, label) + ret["log_probs"] = log_probs + return ret + + logits_processor_args = {"label": label, "temperature": temperature, "loss_mask": loss_mask} + + output = forward_fn( + model, + input_ids, + multi_modal_inputs, + logits_processor=logits_processor, + logits_processor_args=logits_processor_args, + vision_model=hasattr(self.model_config.hf_config, "vision_config"), + pad_token_id=self.model_config.tokenizer.pad_token_id, + data_format=data_format, + mtp_enable_train=self.model_config.mtp.enable and self.model_config.mtp.enable_train, + local_cp_size=local_cp_size, + ) + + # Router replay: record routing decisions for R2 mode + if RouterReplayHelper.is_r2_record_action(self.tf_config, vp_rank): + merge_router_topk_indices(None, input_ids, self.mini_layer_topk_idx_list, self.tf_config, vp_rank) + + # Router replay: switch to backward replay mode for next backward pass + if RouterReplayHelper.is_replay_forward_action(self.tf_config, vp_rank): + router_instance_list = RouterReplayHelper.get_micro_batch_router_list(self.tf_config, vp_rank) + for router in router_instance_list: + router.set_router_replay_action(RouterReplayAction.REPLAY_BACKWARD) + + return output, partial(postprocess_micro_batch_func, data=batch, local_cp_size=local_cp_size) + + def postprocess_micro_batch_func( + self, output, data: TensorDict, forward_only: bool, loss_function, local_cp_size=None + ): + # For memory efficiency + # We move calculation of entropy to compute_log_probs, forward_only == True + device = data["input_ids"].device + model_output = self.prepare_model_outputs(output, data) + + if loss_function is not None: + # TODO(baiyan): How to support hybrid context parallel with dp_group, + # now the dp_group is not used, so just leave it as is, but what if we need to use it? + loss, metrics = loss_function(model_output=model_output, data=data, dp_group=self.get_data_parallel_group()) + # scale loss by num_micro_batch because megatron will scale loss + # by n_micro_batch inside pp schedule + scaled_loss = loss * data["num_micro_batch"] + else: + assert forward_only, "forward_only must be True when loss_function is None" + loss = torch.tensor(1.0, device=device) + scaled_loss = loss + metrics = {} + if local_cp_size is not None: + # aggregate model_output by DP-CP groups + from verl.utils.megatron_utils import dynamic_cp_merge_output + + model_output = dynamic_cp_merge_output( + model_output, + dp_size=mpu.get_data_parallel_world_size(), + dp_rank=mpu.get_data_parallel_rank(), + local_cp_size=local_cp_size, + ) + + output = { + "model_output": model_output, + "loss": loss.detach().item(), + "metrics": metrics, + } + + # return loss and stats + return scaled_loss, output + + +@EngineRegistry.register(model_type="value_model", backend="megatron") +class MegatronEngineWithValueHead(MegatronEngineWithLMHead): + # for value head + def forward_step(self, batch_iter, model, logits_processor_func, postprocess_micro_batch_func): + batch: TensorDict = next(batch_iter) + batch = batch.to(get_device_id()) + model_inputs = self.prepare_model_inputs(batch) + input_ids = model_inputs["input_ids"] + multi_modal_inputs = model_inputs["multi_modal_inputs"] + + from verl.models.mcore import get_mcore_engine_forward_fn + + forward_fn = get_mcore_engine_forward_fn(self.model_config.hf_config) + + output = forward_fn( + model, + input_ids, + multi_modal_inputs, + value_model=True, + vision_model=hasattr(self.model_config.hf_config, "vision_config"), + pad_token_id=self.model_config.tokenizer.pad_token_id, + data_format="thd" if self.engine_config.use_remove_padding else "bshd", + ) + + return output, partial(postprocess_micro_batch_func, data=batch) + + def prepare_model_outputs(self, output: dict | torch.Tensor, data: TensorDict): + return {"values": output} diff --git a/verl/verl/workers/engine/megatron/utils.py b/verl/verl/workers/engine/megatron/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0d9f3b8aadfffe82e17c04094805baa3751c4e7c --- /dev/null +++ b/verl/verl/workers/engine/megatron/utils.py @@ -0,0 +1,35 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from verl.utils.device import get_torch_device + + +def set_random_seed(seed): + import random + + import numpy as np + import torch + + torch.manual_seed(seed) + np.random.seed(seed) + random.seed(seed) + if get_torch_device().device_count() > 0: + from megatron.core import tensor_parallel + + tensor_parallel.model_parallel_cuda_manual_seed(seed) + # FIXME: torch cumsum not support deterministic (used in vllm sampler), + # https://github.com/pytorch/pytorch/issues/89492 + # torch.use_deterministic_algorithms(True, warn_only=True) + # os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8' diff --git a/verl/verl/workers/engine/mindspeed/__init__.py b/verl/verl/workers/engine/mindspeed/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..793b5d650969bcb53745b6541652d9dd64632311 --- /dev/null +++ b/verl/verl/workers/engine/mindspeed/__init__.py @@ -0,0 +1,17 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .transformer_impl import MindspeedEngineWithLMHead, MindspeedEngineWithValueHead, MindSpeedLLMEngineWithLMHead + +__all__ = ["MindspeedEngineWithLMHead", "MindspeedEngineWithValueHead", "MindSpeedLLMEngineWithLMHead"] diff --git a/verl/verl/workers/engine/mindspeed/__pycache__/__init__.cpython-312.pyc b/verl/verl/workers/engine/mindspeed/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fd790ddfdf30c839d9775247c6e61dc18ec53620 Binary files /dev/null and b/verl/verl/workers/engine/mindspeed/__pycache__/__init__.cpython-312.pyc differ diff --git a/verl/verl/workers/engine/mindspeed/__pycache__/transformer_impl.cpython-312.pyc b/verl/verl/workers/engine/mindspeed/__pycache__/transformer_impl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d9aa1d65703f8d75089b39615a5cefea18e8ecd0 Binary files /dev/null and b/verl/verl/workers/engine/mindspeed/__pycache__/transformer_impl.cpython-312.pyc differ diff --git a/verl/verl/workers/engine/mindspeed/transformer_impl.py b/verl/verl/workers/engine/mindspeed/transformer_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..b34f54cd34f11d7e761dfd681a1dc94f7b8b4be7 --- /dev/null +++ b/verl/verl/workers/engine/mindspeed/transformer_impl.py @@ -0,0 +1,139 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +import os + +try: + from mindspeed.megatron_adaptor import repatch +except ImportError: + repatch = None + +from verl.trainer.config import CheckpointConfig +from verl.utils.megatron.router_replay_patch import RouterReplay +from verl.utils.model import print_model_size +from verl.workers.config import ( + HFModelConfig, + McoreEngineConfig, + McoreOptimizerConfig, + MindSpeedEngineConfig, +) + +from ..base import EngineRegistry +from ..megatron import MegatronEngineWithLMHead, MegatronEngineWithValueHead +from .utils import ( + apply_patch, + gpt_model_provider, +) + +logger = logging.getLogger(__file__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + + +def _mindspeed_repatch(engine_config): + if repatch is not None: + repatch_config = dict(engine_config.get("override_transformer_config", {})) + repatch_config.setdefault("use_flash_attn", True) + if engine_config.context_parallel_size > 1: + repatch_config["context_parallel_size"] = engine_config.context_parallel_size + repatch(repatch_config) + + +@EngineRegistry.register(model_type="language_model", backend="megatron", device="npu") +class MindspeedEngineWithLMHead(MegatronEngineWithLMHead): + def __init__( + self, + model_config: HFModelConfig, + engine_config: McoreEngineConfig, + optimizer_config: McoreOptimizerConfig, + checkpoint_config: CheckpointConfig, + ): + super().__init__(model_config, engine_config, optimizer_config, checkpoint_config) + + def _init_device_mesh(self): + # repatch must happen before initialize_model_parallel so that + # initialize_model_parallel_cp_wrapper is in effect when the call is made. + # The initial MindSpeed patch pass sees context_parallel_size=1 (default) because + # verl passes CP size via hydra config rather than --context-parallel-size CLI arg, + # so the CP ring-rank initialization wrapper is not registered on the first pass. + _mindspeed_repatch(self.engine_config) + super()._init_device_mesh() + + +@EngineRegistry.register(model_type="value_model", backend="megatron", device="npu") +class MindspeedEngineWithValueHead(MegatronEngineWithValueHead): + def __init__( + self, + model_config: HFModelConfig, + engine_config: McoreEngineConfig, + optimizer_config: McoreOptimizerConfig, + checkpoint_config: CheckpointConfig, + ): + super().__init__(model_config, engine_config, optimizer_config, checkpoint_config) + + def _init_device_mesh(self): + # repatch must happen before initialize_model_parallel so that + # initialize_model_parallel_cp_wrapper is in effect when the call is made. + # The initial MindSpeed patch pass sees context_parallel_size=1 (default) because + # verl passes CP size via hydra config rather than --context-parallel-size CLI arg, + # so the CP ring-rank initialization wrapper is not registered on the first pass. + _mindspeed_repatch(self.engine_config) + super()._init_device_mesh() + + +@EngineRegistry.register(model_type="language_model", backend="mindspeed_llm", device="npu") +class MindSpeedLLMEngineWithLMHead(MegatronEngineWithLMHead): + def __init__( + self, + model_config: HFModelConfig, + engine_config: MindSpeedEngineConfig, + optimizer_config: McoreOptimizerConfig, + checkpoint_config: CheckpointConfig, + ): + super().__init__(model_config, engine_config, optimizer_config, checkpoint_config) + + def _init_device_mesh(self): + apply_patch(self.model_config, self.engine_config, self.optimizer_config) + super()._init_device_mesh() + + def _build_megatron_module(self): + is_value_model = ( + "ForTokenClassification" in self.model_config.architectures[0] + or "ForSequenceClassification" in self.model_config.architectures[0] + ) + + self.is_value_model = is_value_model + + import torch.distributed + from megatron.core.enums import ModelType + from megatron.training.training import get_model + + # For forward_only, we don't need optimizer, lr_scheduler, checkpoint_mananager + if self.engine_config.forward_only: + module = get_model(gpt_model_provider, ModelType.encoder_or_decoder, wrap_with_ddp=False) + return module + + module = get_model(gpt_model_provider, ModelType.encoder_or_decoder, wrap_with_ddp=True) + if self.vanilla_bridge: + self.bridge.load_weights(module, self.model_config.local_path) + else: + raise ValueError(f"vanilla_bridge should be true now, but got {self.vanilla_bridge}") + + if torch.distributed.get_rank() == 0: + print_model_size(module[0]) + + if self.enable_routing_replay: + print(f"routing replay layers: {len(RouterReplay.router_instances)}") + + return module diff --git a/verl/verl/workers/engine/mindspeed/utils.py b/verl/verl/workers/engine/mindspeed/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9b600e30aba6720a077f2e014d1a0ea8b33fa61d --- /dev/null +++ b/verl/verl/workers/engine/mindspeed/utils.py @@ -0,0 +1,223 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import argparse + +import torch + +from verl.workers.config import HFModelConfig, McoreOptimizerConfig, MindSpeedEngineConfig + + +def get_base_mcore_config_from_model_config(model_config: HFModelConfig) -> dict: + """ + Create a base TransformerConfig with common parameters across different model architectures. + + Args: + model_config: HuggingFace model configuration + + Returns: + TransformerConfig with common parameters + """ + + hf_config = model_config.hf_config + base_config = { + "num_layers": hf_config.num_hidden_layers, + "hidden_size": hf_config.hidden_size, + "num_attention_heads": hf_config.num_attention_heads, + "num_query_groups": hf_config.num_key_value_heads, + "ffn_hidden_size": hf_config.intermediate_size, + "attention_dropout": hf_config.attention_dropout, + "hidden_dropout": getattr(hf_config, "hidden_dropout", 0.0), + "kv_channels": getattr(hf_config, "head_dim", None), + "norm_topk_prob": getattr(hf_config, "norm_topk_prob", False), + "layernorm_epsilon": hf_config.rms_norm_eps, + "max_position_embeddings": hf_config.max_position_embeddings, + "tie_word_embeddings": hf_config.tie_word_embeddings, + "torch_dtype": hf_config.torch_dtype, + "bf16": hf_config.dtype is torch.bfloat16, + "rotary_base": int(hf_config.rope_theta), + "num_experts": getattr(hf_config, "num_experts", None), + "moe_router_topk": getattr(hf_config, "num_experts_per_tok", None), + "moe_ffn_hidden_size": getattr(hf_config, "moe_intermediate_size", None), + "padded_vocab_size": hf_config.vocab_size, + "make_vocab_size_divisible_by": 1, + "untie_embeddings_and_output_weights": True, + } + + tokenizer_config = { + "tokenizer_name_or_path": model_config.tokenizer_path, + "tokenizer_type": "PretrainedFromHF", + } + base_config.update(tokenizer_config) + return base_config + + +def get_base_mcore_config_from_engine_config(engine_config: MindSpeedEngineConfig) -> dict: + """ + Create a base TransformerConfig with common parameters across different model architectures. + + Args: + engine_config: mindspeed engine configuration + + Returns: + TransformerConfig with common parameters + """ + + base_config = { + "tensor_model_parallel_size": engine_config.tensor_model_parallel_size, + "expert_model_parallel_size": engine_config.expert_model_parallel_size, + "expert_tensor_parallel_size": engine_config.expert_tensor_parallel_size, + "pipeline_model_parallel_size": engine_config.pipeline_model_parallel_size, + "virtual_pipeline_model_parallel_size": engine_config.virtual_pipeline_model_parallel_size, + "context_parallel_size": engine_config.context_parallel_size, + "sequence_parallel": engine_config.sequence_parallel, + "use_distributed_optimizer": engine_config.use_distributed_optimizer, + "seed": engine_config.seed, + } + if engine_config.strategy == "mindspeed_llm": + base_config.update(engine_config.llm_kwargs) + elif engine_config.strategy == "mindspeed_mm": + base_config.update(engine_config.mm_kwargs) + return base_config + + +def get_base_mcore_config_from_optim_config(optim_config: McoreOptimizerConfig) -> dict: + """ + Create a base TransformerConfig with common parameters across different model architectures. + + Args: + optim_config: megatron optimizer configuration + + Returns: + TransformerConfig with common parameters + """ + + base_config = { + "lr": optim_config.lr, + "lr_decay_style": optim_config.lr_decay_style, + "min_lr": optim_config.min_lr, + "weight_decay": optim_config.weight_decay, + "lr_warmup_fraction": optim_config.lr_warmup_steps_ratio, + "clip_grad": optim_config.clip_grad, + "adam_beta1": optim_config.betas[0], + "adam_beta2": optim_config.betas[1], + } + + base_config.update(optim_config.override_optimizer_config) + return base_config + + +def set_global_config(config): + from megatron.training.arguments import parse_args, validate_args + from megatron.training.global_vars import set_global_variables + + args = parse_args(ignore_unknown_args=True) + for key, value in config.items(): + setattr(args, key, value) + + validate_args(args) + try: + set_global_variables(args) + except AssertionError: + print("megatron args already set") + + +def add_mcore_arguments(all_config: dict) -> dict: + mcore_config_dict = {} + mcore_config_list = [] + for key, value in all_config.items(): + if value is None: + continue + mcore_config_dict[key] = value + if isinstance(value, bool): + mcore_config_list.append(f"--{key.replace('_', '-')}") + + from megatron.training.arguments import add_megatron_arguments + + parser = argparse.ArgumentParser(description="Megatron-LM Arguments", allow_abbrev=False) + parser = add_megatron_arguments(parser) + args, _ = parser.parse_known_args(mcore_config_list) + return {**vars(args), **mcore_config_dict} + + +def apply_patch(model_config, engine_config, optimizer_config): + model_config = get_base_mcore_config_from_model_config(model_config) + optimizer_config = get_base_mcore_config_from_optim_config(optimizer_config) + engine_config = get_base_mcore_config_from_engine_config(engine_config) + all_config = {**model_config, **optimizer_config, **engine_config} + mcore_config = add_mcore_arguments(all_config) + from mindspeed_llm.tasks.megatron_adaptor_v2 import repatch + + repatch(mcore_config) + set_global_config(mcore_config) + + +def gpt_model_provider(pre_process=True, post_process=True): + """ + Builds the model. + + If you set the use_mcore_models to True, it will return the mcore GPT model and if not the legacy GPT model. + + Args: + pre_process (bool, optional): Set to true if you need to compute embedings. Defaults to True. + post_process (bool, optional): Set to true if you need to want to compute output logits/loss. + Defaults to True. + + + Returns: + Union[GPTModel, megatron.legacy.model.GPTModel]: The returned model + """ + from megatron.core.models.gpt import GPTModel + from megatron.core.models.gpt.gpt_layer_specs import ( + get_gpt_layer_local_spec, + get_gpt_layer_with_transformer_engine_spec, + ) + from megatron.core.transformer.spec_utils import import_module + from megatron.training import get_args + from megatron.training.arguments import core_transformer_config_from_args + + args = get_args() + use_te = args.transformer_impl == "transformer_engine" + # Experimental loading arguments from configs + config = core_transformer_config_from_args(args) + + if args.spec is not None: + transformer_layer_spec = import_module(args.spec) + else: + if use_te: + transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec( + args.num_experts, args.moe_grouped_gemm, qk_layernorm=args.qk_layernorm + ) + else: + transformer_layer_spec = get_gpt_layer_local_spec( + args.num_experts, args.moe_grouped_gemm, qk_layernorm=args.qk_layernorm + ) + + model = GPTModel( + config=config, + transformer_layer_spec=transformer_layer_spec, + vocab_size=args.padded_vocab_size, + max_sequence_length=args.max_position_embeddings, + pre_process=pre_process, + post_process=post_process, + fp16_lm_cross_entropy=args.fp16_lm_cross_entropy, + parallel_output=True, + share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, + position_embedding_type=args.position_embedding_type, + rotary_percent=args.rotary_percent, + seq_len_interpolation_factor=args.rotary_seq_len_interpolation_factor, + ) + + return model diff --git a/verl/verl/workers/engine/torchtitan/__init__.py b/verl/verl/workers/engine/torchtitan/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..345757277af1504e14f8921143b0afc04d897554 --- /dev/null +++ b/verl/verl/workers/engine/torchtitan/__init__.py @@ -0,0 +1,16 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from .transformer_impl import TorchTitanEngine, TorchTitanEngineWithLMHead + +__all__ = ["TorchTitanEngine", "TorchTitanEngineWithLMHead"] diff --git a/verl/verl/workers/engine/torchtitan/__pycache__/__init__.cpython-312.pyc b/verl/verl/workers/engine/torchtitan/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..725ed98bcc65a72472a33ee8a7ae112a88da868a Binary files /dev/null and b/verl/verl/workers/engine/torchtitan/__pycache__/__init__.cpython-312.pyc differ diff --git a/verl/verl/workers/engine/torchtitan/__pycache__/transformer_impl.cpython-312.pyc b/verl/verl/workers/engine/torchtitan/__pycache__/transformer_impl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f8a0d66b13813ea7e08a04a0a9da22a6f1522c9f Binary files /dev/null and b/verl/verl/workers/engine/torchtitan/__pycache__/transformer_impl.cpython-312.pyc differ diff --git a/verl/verl/workers/engine/torchtitan/transformer_impl.py b/verl/verl/workers/engine/torchtitan/transformer_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..6db178b00536670696371149293286acf399b2fa --- /dev/null +++ b/verl/verl/workers/engine/torchtitan/transformer_impl.py @@ -0,0 +1,735 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +The concrete Engine implementation using PyTorch TorchTitan parallelism (FSDP2 + TP + PP) +""" + +import gc +import importlib +import logging +import os +import re +from contextlib import nullcontext +from typing import Any, Callable, Optional + +import torch +import torch.distributed +from tensordict import TensorDict +from torch.distributed.checkpoint.state_dict import get_model_state_dict +from torch.distributed.tensor import DTensor +from torchtitan.components.checkpoint import CheckpointManager +from torchtitan.components.loss import CrossEntropyLoss +from torchtitan.components.lr_scheduler import LRSchedulersContainer +from torchtitan.components.optimizer import OptimizersContainer +from torchtitan.config import CompileConfig, ParallelismConfig, TrainingConfig +from torchtitan.distributed import utils as dist_utils +from torchtitan.distributed.context_parallel import prepare_context_parallel_input +from torchtitan.distributed.parallel_dims import ParallelDims +from torchtitan.train import Trainer + +import verl.utils.torch_functional as verl_F +from verl.trainer.config import CheckpointConfig +from verl.utils import tensordict_utils as tu +from verl.utils.dataset.dataset_utils import DatasetPadMode +from verl.utils.debug import log_gpu_memory_usage +from verl.utils.device import get_device_id, get_device_name +from verl.utils.fsdp_utils import ( + load_fsdp_model_to_gpu, + load_fsdp_optimizer, + offload_fsdp_model_to_cpu, + offload_fsdp_optimizer, +) +from verl.utils.model import extract_multi_modal_inputs +from verl.utils.torch_functional import logprobs_from_logits +from verl.workers.config import HFModelConfig, TorchtitanEngineConfig, TorchtitanOptimizerConfig +from verl.workers.engine.torchtitan.utils import ( + NoOpDataLoader, + derive_torchtitan_name_and_flavor, + enable_fsdp_gradient_division, + get_attention_masks, + iter_per_tensor_params_ep, +) + +from ..base import BaseEngine, BaseEngineCtx, EngineRegistry +from ..utils import enable_full_determinism, postprocess_batch_func, prepare_micro_batches + +logger = logging.getLogger(__file__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + +device_name = get_device_name() + + +class TorchTitanEngine(BaseEngine): + """ + Concrete Engine implementation using PyTorch TorchTitan parallelism. + + Supports model sharding with FSDP2, tensor parallelism, activation/optimizer offloading, + LoRA, and sequence parallelism following the TorchTitan design. + """ + + def __init__( + self, + model_config: HFModelConfig, + engine_config: TorchtitanEngineConfig, + optimizer_config: TorchtitanOptimizerConfig, + checkpoint_config: CheckpointConfig, + ): + """ + Initialize the TorchTitanEngine. + + Sets up distributed device meshes for tensor and data parallelism, LoRA, and offload policies. + + Args: + model_config: Configuration for HuggingFace model. + engine_config: Configuration for FSDP/TorchTitan engine (uses FSDP2). + optimizer_config: Configuration for optimizer. + checkpoint_config: Configuration for checkpointing. + """ + super().__init__() + + self.model_config = model_config + self.engine_config = engine_config + self.optimizer_config = optimizer_config + self.checkpoint_config = checkpoint_config + + # Derive torchtitan model name and flavor from HF config + torchtitan_name, torchtitan_flavor = derive_torchtitan_name_and_flavor(self.model_config.hf_config) + + # Get ModelSpec from model registry + model_module = importlib.import_module(f"torchtitan.models.{torchtitan_name}") + model_spec = model_module.model_registry(torchtitan_flavor, attn_backend=self.engine_config.attn_type) + + optimizer = OptimizersContainer.Config( + name=self.optimizer_config.name, + lr=self.optimizer_config.lr, + eps=self.optimizer_config.eps, + beta1=self.optimizer_config.betas[0], + beta2=self.optimizer_config.betas[1], + weight_decay=self.optimizer_config.weight_decay, + ) + + total_steps = self.optimizer_config.total_training_steps + lr_warmup_steps = self.optimizer_config.lr_warmup_steps + if lr_warmup_steps is None or lr_warmup_steps <= 0: + lr_warmup_steps = int(self.optimizer_config.lr_warmup_steps_ratio * total_steps) + + lr_scheduler = LRSchedulersContainer.Config( + warmup_steps=lr_warmup_steps, + decay_type=self.optimizer_config.decay_type, + min_lr_factor=self.optimizer_config.min_lr_factor, + ) + parallelism = ParallelismConfig( + data_parallel_replicate_degree=self.engine_config.data_parallel_replicate_size, + data_parallel_shard_degree=self.engine_config.data_parallel_shard_size, + fsdp_reshard_after_forward=self.engine_config.reshard_after_forward, + tensor_parallel_degree=self.engine_config.tensor_parallel_size, + pipeline_parallel_degree=self.engine_config.pipeline_parallel_size, + context_parallel_degree=self.engine_config.context_parallel_size, + expert_parallel_degree=self.engine_config.expert_parallel_size, + ) + checkpoint = CheckpointManager.Config( + enable=True, + initial_load_in_hf=True, + initial_load_model_only=True, + initial_load_path=model_config.path, + ) + compile_config = CompileConfig(enable=self.engine_config.use_torch_compile) + training_kwargs = {} + if self.engine_config.max_seq_len is not None: + training_kwargs["seq_len"] = self.engine_config.max_seq_len + if self.engine_config.offload_policy or self.engine_config.forward_only: + training = TrainingConfig(enable_cpu_offload=True, **training_kwargs) + else: + training = TrainingConfig(**training_kwargs) + + # Construct Torchtitan's Trainer.Config + self.config = Trainer.Config( + model_spec=model_spec, + hf_assets_path=self.model_config.path, + optimizer=optimizer, + lr_scheduler=lr_scheduler, + parallelism=parallelism, + checkpoint=checkpoint, + compile=compile_config, + training=training, + # Use a no-op dataloader since verl has its own data loading + dataloader=NoOpDataLoader.Config(), + # Provide a concrete loss so Trainer.__init__ can build it; + # verl uses its own loss function and ignores this one. + loss=CrossEntropyLoss.Config(), + ) + self.trainer = Trainer(self.config) + + self._init_device_mesh() + + # Re-enable FSDP's gradient division for verl's loss scaling. + # TorchTitan disables gradient division by default (for global token normalization), + # but verl's loss function multiplies by dp_size to compensate for gradient averaging. + if self.engine_config.data_parallel_shard_size > 1: + dp_size = self.get_data_parallel_size() + for model_part in self.trainer.model_parts: + enable_fsdp_gradient_division(model_part, dp_size) + + if self.engine_config.full_determinism: + enable_full_determinism(seed=self.engine_config.seed) + + # set FSDP offload params + self._is_offload_param = self.engine_config.param_offload + self._is_offload_optimizer = self.engine_config.optimizer_offload + + if self.engine_config.entropy_from_logits_with_chunking: + entropy_from_logits = verl_F.entropy_from_logits_with_chunking + else: + entropy_from_logits = verl_F.entropy_from_logits + + self.compute_entropy_from_logits = ( + torch.compile(entropy_from_logits, dynamic=True) + if self.engine_config.use_torch_compile + else entropy_from_logits + ) + + @property + def is_param_offload_enabled(self) -> bool: + return self._is_offload_param + + @property + def is_optimizer_offload_enabled(self) -> bool: + return self._is_offload_optimizer + + def is_mp_src_rank_with_outputs(self): + """ + Whether the current rank is the first rank in model parallel group that contains model outputs + """ + is_collect = True + # TP: outputs are on TP rank 0 + if self.parallel_dims.tp > 1: + tp_mesh = self.parallel_dims.get_optional_mesh("tp") + is_collect = is_collect and (tp_mesh.get_local_rank() == 0) + # PP: outputs are on the last PP rank + if self.parallel_dims.pp > 1: + pp_mesh = self.parallel_dims.get_optional_mesh("pp") + is_collect = is_collect and (pp_mesh.get_local_rank() == self.parallel_dims.pp - 1) + # CP: outputs are on CP rank 0 + if self.parallel_dims.cp > 1: + cp_mesh = self.parallel_dims.get_optional_mesh("cp") + is_collect = is_collect and (cp_mesh.get_local_rank() == 0) + return is_collect + + def initialize(self): + """ + Build the model, optimizer, and learning rate scheduler with TorchTitan parallelism. + + Applies device, dtype, and precision configurations, including mixed precision. + Sets up checkpoint manager. + """ + self.module = self.trainer.model_parts + self.checkpointer = self.trainer.checkpointer + # load initial HF weights + self.checkpointer.load() + + if not self.engine_config.forward_only: + self.optimizer = self.trainer.optimizers + self.lr_scheduler = self.trainer.lr_schedulers + else: + self.optimizer = None + self.lr_scheduler = None + + self.to( + device="cpu", + model=self._is_offload_param, + optimizer=self._is_offload_optimizer, + grad=self._is_offload_param, + ) + + log_gpu_memory_usage("After offload model/optimizer/grad during init", logger=logger) + + def _init_device_mesh(self): + """Initialize the device mesh for TorchTitan style parallelism.""" + world_size = torch.distributed.get_world_size() + self.parallel_dims = ParallelDims( + dp_shard=self.engine_config.data_parallel_shard_size, + dp_replicate=self.engine_config.data_parallel_replicate_size, + cp=self.engine_config.context_parallel_size, + tp=self.engine_config.tensor_parallel_size, + pp=self.engine_config.pipeline_parallel_size, + ep=self.engine_config.expert_parallel_size, + world_size=world_size, + ) + self.device_mesh = self.parallel_dims.build_mesh() + + def train_mode(self, **kwargs): + """Return a context manager for training mode.""" + return EngineTrainModeCtx(self, **kwargs) + + def eval_mode(self, **kwargs): + """Return a context manager for evaluation mode.""" + return EngineEvalModeCtx(self, **kwargs) + + def get_data_parallel_rank(self): + mesh = self._get_data_parallel_mesh() + return 0 if mesh is None else mesh.get_local_rank() + + def get_data_parallel_size(self): + return self.engine_config.data_parallel_shard_size * self.engine_config.data_parallel_replicate_size + + def get_data_parallel_group(self): + mesh = self._get_data_parallel_mesh() + if mesh is not None: + return mesh.get_group() + # If world_size == dp_size (e.g. single GPU, or all ranks are DP), + # return WORLD so that collective ops in _postprocess_output + # (allgather_dict_into_dict, all_reduce) still run and produce the + # correct metric aggregation format. + if torch.distributed.get_world_size() == self.get_data_parallel_size(): + return torch.distributed.group.WORLD + return None + + def get_model_parallel_group(self): + raise NotImplementedError + + def get_context_parallel_group(self): + raise NotImplementedError + + def _get_data_parallel_mesh(self): + """Get the data parallel mesh, handling hybrid/fully/replicate shard modes.""" + mesh = self.parallel_dims.get_optional_mesh(["dp_replicate", "fsdp"]) + if mesh is None: + mesh = self.parallel_dims.get_optional_mesh("fsdp") + if mesh is None: + mesh = self.parallel_dims.get_optional_mesh("dp_replicate") + return mesh + + def forward_backward_batch(self, data: TensorDict, loss_function: Callable, forward_only=False): + """Perform forward and optionally backward pass on a batch.""" + tu.assign_non_tensor(data, sp_size=self.engine_config.tensor_parallel_size) + + # Compute num_tokens in global batch for loss normalization + batch_num_tokens = data["loss_mask"].sum().to(get_device_id()) + dp_group = self.get_data_parallel_group() + if dp_group is not None: + torch.distributed.all_reduce(batch_num_tokens, op=torch.distributed.ReduceOp.SUM, group=dp_group) + tu.assign_non_tensor(data, batch_num_tokens=batch_num_tokens.item()) + tu.assign_non_tensor(data, dp_size=self.get_data_parallel_size()) + + micro_batches, indices = prepare_micro_batches( + data=data, + dp_group=self.get_data_parallel_group(), + same_micro_num_in_dp=True, + ) + + output_lst = [] + + ctx = torch.no_grad() if forward_only else nullcontext() + + for micro_batch in micro_batches: + with ctx: + loss, output = self.forward_step(micro_batch, loss_function=loss_function, forward_only=forward_only) + if not forward_only: + loss.backward() + output_lst.append(output) + + return postprocess_batch_func(output_lst=output_lst, indices=indices, data=data) + + def model_forward_step( + self, + *, + inputs: torch.Tensor, + extra_inputs: dict[str, torch.Tensor] | None = None, + extra_kwargs: dict[str, torch.Tensor] | None = None, + ) -> torch.Tensor: + """ + Perform a forward pass through the trainer model without backward. + """ + model_parts = self.module + parallel_dims = self.parallel_dims + + if parallel_dims.pp_enabled: + raise NotImplementedError( + "Pipeline parallelism is not yet supported in model_forward_step. " + "This will be implemented in a follow-up PR." + ) + else: + # Non-PP forward + assert len(model_parts) == 1 + with self.trainer.train_context(): + pred = model_parts[0](inputs, **extra_inputs, **extra_kwargs) + + if isinstance(pred, DTensor): + pred = pred.full_tensor() + return pred + + def forward_step(self, micro_batch: TensorDict, loss_function, forward_only): + raise NotImplementedError("forward_step must be implemented in subclass") + + def optimizer_zero_grad(self): + """Zero gradients.""" + self.optimizer.zero_grad() + + def optimizer_step(self): + """Perform optimizer step with gradient clipping.""" + grad_norm = dist_utils.clip_grad_norm_( + [p for m in self.module for p in m.parameters()], + self.config.training.max_norm, + foreach=True, + pp_mesh=self.parallel_dims.get_optional_mesh("pp"), + ep_enabled=self.parallel_dims.ep_enabled, + ) + + # if grad_norm is not finite, skip the update + if not torch.isfinite(grad_norm): + logger.warning(f"grad_norm is not finite: {grad_norm}") + self.optimizer.zero_grad() + else: + self.optimizer.step() + return grad_norm.item() + + def lr_scheduler_step(self): + """Advance learning rate scheduler.""" + self.lr_scheduler.step() + lr = self.lr_scheduler.schedulers[0].get_last_lr()[0] + return lr + + def to(self, device: str, model: bool = True, optimizer: bool = True, grad: bool = True): + """Move model and/or optimizer to CPU or GPU.""" + super().to(device=device, model=model, optimizer=optimizer, grad=grad) + + if self.engine_config.forward_only: + return + + device_name = get_device_name() + assert device in (device_name, "cpu") + if device == device_name: + if model: + for module in self.module: + load_fsdp_model_to_gpu(module) + if optimizer and self.optimizer is not None: + load_fsdp_optimizer(self.optimizer, device) + gc.collect() + elif device == "cpu": + if model: + for module in self.module: + offload_fsdp_model_to_cpu(module) + if optimizer and self.optimizer is not None: + offload_fsdp_optimizer(self.optimizer) + else: + raise ValueError(f"Invalid device type: {device}") + + def save_checkpoint( + self, + local_path: str, + hdfs_path: Optional[str] = None, + global_step: int = 0, + max_ckpt_to_keep: Optional[int] = None, + **kwargs, + ) -> None: + """Save checkpoint.""" + if self._is_offload_param: + for module in self.module: + load_fsdp_model_to_gpu(module) + + # Override TorchTitan's folder to use verl's path + parent_dir = os.path.dirname(local_path) + self.checkpointer.folder = parent_dir + + if max_ckpt_to_keep is not None: + self.checkpointer.keep_latest_k = max_ckpt_to_keep + + self.checkpointer.save(curr_step=global_step) + + torch.distributed.barrier() + if self._is_offload_param: + for module in self.module: + offload_fsdp_model_to_cpu(module) + + def load_checkpoint( + self, local_path: str, hdfs_path: Optional[str] = None, del_local_after_load: int = True, **kwargs + ) -> None: + """Load checkpoint.""" + if self._is_offload_param: + for module in self.module: + load_fsdp_model_to_gpu(module) + + # Override TorchTitan's folder to use verl's path + parent_dir = os.path.dirname(local_path) + self.checkpointer.folder = parent_dir + + # Extract step number from path (verl uses global_step_N format) + match = re.search(r"global_step_(\d+)", local_path) + if match: + step = int(match.group(1)) + self.checkpointer.load(step=step) + else: + # Fallback to latest + self.checkpointer.load(step=-1) + + torch.distributed.barrier() + if self._is_offload_param: + for module in self.module: + offload_fsdp_model_to_cpu(module) + + if self._is_offload_optimizer: + offload_fsdp_optimizer(self.optimizer) + + def get_per_tensor_param(self, **kwargs): + for module in self.module: + load_fsdp_model_to_gpu(module) + + # Collect state dicts from all model parts + params = {} + for module in self.module: + module_params = get_model_state_dict(module) + params.update(module_params) + + if self._is_offload_param: + for module in self.module: + offload_fsdp_model_to_cpu(module) + + # Convert TorchTitan key names to HuggingFace key names (expected by vLLM) + sd_adapter = self.checkpointer.sd_adapter + if sd_adapter is not None: + params = sd_adapter.to_hf(params) + + # When weight tying is enabled, the sd_adapter skips lm_head.weight during + # to_hf() conversion (since it's the same tensor as embed_tokens.weight in + # the torchtitan model). But vLLM needs lm_head.weight explicitly, so we + # add it back as a reference to embed_tokens.weight. + if "model.embed_tokens.weight" in params and "lm_head.weight" not in params: + params["lm_head.weight"] = params["model.embed_tokens.weight"] + + device = get_device_id() # used when fsdp2 set cpu_offload_policy + + # When Expert Parallel (EP) is used, sd_adapter.to_hf() only produces + # individual expert weights for the locally-owned experts (e.g., 16 out of + # 128 with EP=8). vLLM needs ALL experts. We gather the missing experts + # by all-gathering each expert weight across the EP process group. + if self.parallel_dims.ep_enabled: + ep_mesh = self.parallel_dims.get_optional_mesh("ep") + ep_group = ep_mesh.get_group() + ep_size = self.parallel_dims.ep + per_tensor_param = iter_per_tensor_params_ep(params, device, ep_group, ep_size) + else: + # TODO: cast fp32 to bf16 to reduce weight sync overhead, need more fine-grained control, e.g MoE gate + per_tensor_param = ( + ( + name, + param.to(device, non_blocking=True).full_tensor().to(torch.bfloat16, non_blocking=True) + if isinstance(param, DTensor) + else param, + ) + for name, param in params.items() + ) + # TODO: support Torchtitan PEFT + return per_tensor_param, None + + +class EngineEvalModeCtx(BaseEngineCtx): + def __init__(self, engine: TorchTitanEngine, **kwargs): + super().__init__(engine=engine, mode="eval", **kwargs) + + def __enter__(self): + assert isinstance(self.engine, TorchTitanEngine) + super().__enter__() + for module in self.engine.module: + module.eval() + + def __exit__(self, exc_type, exc_value, traceback): + assert isinstance(self.engine, TorchTitanEngine) + + # Reshard the root FSDP module + if self.engine.engine_config.data_parallel_shard_size > 1: + for module in self.engine.module: + module.reshard() + + super().__exit__(exc_type, exc_value, traceback) + + +class EngineTrainModeCtx(BaseEngineCtx): + def __init__(self, engine: TorchTitanEngine, **kwargs): + super().__init__(engine=engine, mode="train", **kwargs) + + def __enter__(self): + assert isinstance(self.engine, TorchTitanEngine) + super().__enter__() + for module in self.engine.module: + module.train() + + def __exit__(self, exc_type, exc_value, traceback): + assert isinstance(self.engine, TorchTitanEngine) + self.engine.optimizer_zero_grad() + super().__exit__(exc_type, exc_value, traceback) + + +@EngineRegistry.register(model_type="language_model", backend=["torchtitan"], device=["cuda", "npu"]) +class TorchTitanEngineWithLMHead(TorchTitanEngine): + """TorchTitan engine implementation for language models with LM head.""" + + def prepare_model_inputs(self, micro_batch: TensorDict): + use_remove_padding = tu.get_non_tensor_data(data=micro_batch, key="use_remove_padding", default=True) + pad_mode = tu.get_non_tensor_data(data=micro_batch, key="pad_mode", default=DatasetPadMode.NO_PADDING) + assert pad_mode == DatasetPadMode.NO_PADDING, f"pad_mode {pad_mode} not supported" + + multi_modal_inputs = extract_multi_modal_inputs(micro_batch.get("multi_modal_inputs", [])) + input_ids = micro_batch["input_ids"] + position_ids = micro_batch["position_ids"] + output_args = {} + + if use_remove_padding: + input_ids = input_ids.values().unsqueeze(0) + if position_ids.dim() == 3: + position_ids = position_ids.values().unsqueeze(1) + else: + position_ids = position_ids.values().unsqueeze(0) + + labels = torch.roll(input_ids, shifts=-1, dims=1) + attn_type = self.engine_config.attn_type + attention_mask = get_attention_masks( + input_batch=input_ids, + positions=position_ids, + attn_type=attn_type, + ) + else: + loss_mask = micro_batch["loss_mask"] + pad_token_id = tu.get_non_tensor_data(data=micro_batch, key="pad_token_id", default=0) + batch_size = micro_batch.batch_size[0] + max_seq_len = max(input_ids.offsets().diff()) + + labels = torch.roll(input_ids.values(), shifts=-1, dims=0) + input_ids = torch.nested.to_padded_tensor( + input_ids, padding=pad_token_id, output_size=(batch_size, max_seq_len) + ) + + if position_ids.dim() == 3: + position_ids = torch.nested.to_padded_tensor( + position_ids, padding=0, output_size=(batch_size, 4, max_seq_len) + ).transpose(0, 1) + else: + position_ids = torch.nested.to_padded_tensor( + position_ids, padding=0, output_size=(batch_size, max_seq_len) + ) + + attention_mask_list = [torch.ones_like(t, dtype=torch.int32) for t in loss_mask] + attention_mask = torch.nested.as_nested_tensor(attention_mask_list, layout=torch.jagged) + attention_mask = torch.nested.to_padded_tensor( + attention_mask, padding=0, output_size=(batch_size, max_seq_len) + ) + + extra_inputs = { + "positions": position_ids, + } + # For arguments, like attention_masks, we have to put them in a separate + # dict as extra_inputs are not forwarded to other stages in PP, but + # extra_kwargs are. + extra_kwargs: dict[str, Any] = {"attention_masks": attention_mask} + if self.parallel_dims.cp_enabled: + input_ids, labels, extra_kwargs = prepare_context_parallel_input( + input_ids, + labels, + extra_kwargs, + self.parallel_dims.get_mesh("cp"), + self.trainer.device, + self.trainer.config.parallelism.context_parallel_load_balancer, + ) + + # TODO(jessicazhong): multimodal is not yet supported for Torchtitan engine + extra_inputs.update(multi_modal_inputs) + output_args["labels"] = labels + return input_ids, extra_inputs, extra_kwargs, output_args + + def prepare_model_outputs(self, logits, output_args, micro_batch: TensorDict): + use_remove_padding = tu.get_non_tensor_data(data=micro_batch, key="use_remove_padding", default=True) + pad_mode = tu.get_non_tensor_data(data=micro_batch, key="pad_mode", default=DatasetPadMode.NO_PADDING) + assert pad_mode == DatasetPadMode.NO_PADDING, f"pad_mode {pad_mode} not supported" + + temperature = micro_batch["temperature"] + calculate_entropy = tu.get_non_tensor_data(data=micro_batch, key="calculate_entropy", default=False) + labels = output_args["labels"] + model_output = {} + + input_ids = micro_batch["input_ids"] + cu_seqlens = input_ids.offsets() + if use_remove_padding: + labels = labels.squeeze(0) + logits_rmpad = logits.squeeze(0) + # PyTorch's autograd doesn't allow in-place modification of views when gradients need to flow back + logits_rmpad = logits_rmpad / temperature + + inplace_backward = True + if calculate_entropy: + inplace_backward = False + log_probs = logprobs_from_logits( + logits=logits_rmpad, + labels=labels, + inplace_backward=inplace_backward, + ) + + if calculate_entropy: + if not self.engine_config.entropy_checkpointing: + entropy_rmpad = self.compute_entropy_from_logits(logits_rmpad) + else: + entropy_rmpad = torch.utils.checkpoint.checkpoint(self.compute_entropy_from_logits, logits_rmpad) + + log_probs = torch.nested.nested_tensor_from_jagged(log_probs.squeeze(0), cu_seqlens) + if calculate_entropy: + entropy = torch.nested.nested_tensor_from_jagged(entropy_rmpad, cu_seqlens) + else: + logits.div_(temperature) + if calculate_entropy: + if not self.engine_config.entropy_checkpointing: + entropy = verl_F.entropy_from_logits(logits) + else: + entropy = torch.utils.checkpoint.checkpoint(verl_F.entropy_from_logits, logits) + + seq_lengths = cu_seqlens.diff() + starts = torch.zeros_like(seq_lengths, dtype=torch.int64) + logits = torch.nested.narrow(logits, 1, starts, seq_lengths, layout=torch.jagged) + logits_rmpad = torch.cat([t for t in logits.unbind()]) + log_probs = logprobs_from_logits(logits=logits_rmpad, labels=output_args["labels"]) + log_probs = torch.nested.nested_tensor_from_jagged(log_probs, cu_seqlens) + if calculate_entropy: + entropy = torch.nested.narrow(entropy, 1, starts, seq_lengths, layout=torch.jagged) + entropy_rmpad = torch.cat([t for t in entropy.unbind()]) + entropy = torch.nested.nested_tensor_from_jagged(entropy_rmpad, cu_seqlens) + + model_output["log_probs"] = log_probs + if calculate_entropy: + model_output["entropy"] = entropy + + return model_output + + def forward_step(self, micro_batch: TensorDict, loss_function, forward_only): + device_name = get_device_name() + micro_batch = micro_batch.to(get_device_id()) + input_ids, extra_inputs, extra_kwargs, output_args = self.prepare_model_inputs(micro_batch=micro_batch) + + with torch.autocast(device_type=device_name, dtype=torch.bfloat16): + logits = self.model_forward_step(inputs=input_ids, extra_inputs=extra_inputs, extra_kwargs=extra_kwargs) + + model_output = self.prepare_model_outputs(logits=logits, output_args=output_args, micro_batch=micro_batch) + + if loss_function is not None: + loss, metrics = loss_function( + model_output=model_output, data=micro_batch, dp_group=self.get_data_parallel_group() + ) + else: + assert forward_only, "forward_only must be True when loss_function is None" + loss = torch.tensor(1.0, device=device_name) + metrics = {} + + output = { + "model_output": model_output, + "loss": loss.detach().item(), + "metrics": metrics, + } + + return loss, output diff --git a/verl/verl/workers/engine/torchtitan/utils.py b/verl/verl/workers/engine/torchtitan/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..32b10542f41a95e42a458acfe38513cfaaac8f38 --- /dev/null +++ b/verl/verl/workers/engine/torchtitan/utils.py @@ -0,0 +1,357 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import importlib +import logging +import re +from collections import defaultdict +from collections.abc import Generator, Iterator +from dataclasses import dataclass +from typing import Any + +import torch +import torch.distributed +import torch.nn as nn +from torch.distributed._composable.fsdp import FSDPModule +from torch.distributed.tensor import DTensor +from torch.nn.attention.flex_attention import _mask_mod_signature, and_masks +from torchtitan.components.dataloader import BaseDataLoader +from torchtitan.models.common.attention import ( + AttentionMasksType, + VarlenMetadata, + create_attention_mask, + get_causal_mask_mod, +) + +logger = logging.getLogger(__name__) + + +class NoOpDataLoader(BaseDataLoader): + """A no-op dataloader for use when verl manages its own data loading. + + Satisfies the BaseDataLoader interface required by torchtitan's Trainer + but does nothing. Its __iter__ yields nothing, and state_dict / + load_state_dict are no-ops. + """ + + @dataclass(kw_only=True, slots=True) + class Config(BaseDataLoader.Config): + pass + + def __init__(self, **kwargs): + pass + + def __iter__(self) -> Iterator[tuple[dict[str, torch.Tensor], torch.Tensor]]: + return iter([]) + + def state_dict(self): + return {} + + def load_state_dict(self, state_dict): + pass + + +# Mapping from HuggingFace model_type to torchtitan model name. +# Torchtitan models not mapped here: +# - flux: diffusion model, not applicable to verl's RL/SFT workflows +# - llama3_ft: fault-tolerant variant of llama3, same HF models (mapped via "llama") +_HF_MODEL_TYPE_TO_TORCHTITAN_NAME = { + "qwen2": "qwen3", + "qwen3": "qwen3", + "qwen2_moe": "qwen3", + "qwen3_moe": "qwen3", + "llama": "llama3", + "llama4": "llama4", + "deepseek_v3": "deepseek_v3", + "gpt_oss": "gpt_oss", +} + + +def derive_torchtitan_name_and_flavor(hf_config) -> tuple[str, str]: + """Derive torchtitan model name and flavor from a HuggingFace config. + + The name is mapped from ``hf_config.model_type``. The flavor is found by + matching architecture parameters (dim, n_layers, vocab_size) against the + known flavors registered in the torchtitan model package. + + Args: + hf_config: A HuggingFace AutoConfig object. + + Returns: + A ``(name, flavor)`` tuple. + + Raises: + ValueError: If model_type is unsupported or no matching flavor is found. + """ + model_type = getattr(hf_config, "model_type", None) + if model_type is None: + raise ValueError("HuggingFace config does not have 'model_type' field") + + name = _HF_MODEL_TYPE_TO_TORCHTITAN_NAME.get(model_type) + if name is None: + raise ValueError( + f"Cannot derive torchtitan model name from HF model_type '{model_type}'. " + f"Supported types: {list(_HF_MODEL_TYPE_TO_TORCHTITAN_NAME.keys())}." + ) + + # Import the model package and use model_registry to build each flavor's config. + # model_registry has sensible defaults for all optional params (attn_backend, etc.). + model_module = importlib.import_module(f"torchtitan.models.{name}") + model_registry = model_module.model_registry + + # The configs dict name isn't derivable from the model name + # (e.g. gpt_oss -> gptoss_configs), so we find it by convention. + flavor_names = None + for attr, obj in vars(model_module).items(): + if attr.endswith("_configs") and isinstance(obj, dict): + flavor_names = list(obj.keys()) + break + + if flavor_names is None: + raise ValueError( + f"Could not find model configs dict in torchtitan.models.{name}. " + f"Expected a dict attribute ending with '_configs'." + ) + + hidden_size = hf_config.hidden_size + num_layers = hf_config.num_hidden_layers + vocab_size = hf_config.vocab_size + + for flavor_name in flavor_names: + cfg = model_registry(flavor_name).model + n_layers = getattr(cfg, "n_layers", None) or len(getattr(cfg, "layers", [])) + if ( + getattr(cfg, "dim", None) == hidden_size + and n_layers == num_layers + and getattr(cfg, "vocab_size", None) == vocab_size + ): + logger.info( + f"Auto-derived torchtitan name='{name}', flavor='{flavor_name}' from HF model_type='{model_type}'" + ) + return name, flavor_name + + raise ValueError( + f"No matching torchtitan flavor found for model_type='{model_type}' " + f"(hidden_size={hidden_size}, num_hidden_layers={num_layers}, " + f"vocab_size={vocab_size}). " + f"Available flavors for '{name}': {flavor_names}." + ) + + +def enable_fsdp_gradient_division(model: nn.Module, dp_size: int) -> None: + """ + Re-enable FSDP's automatic gradient division. + + TorchTitan calls disable_fsdp_gradient_division() which sets gradient_divide_factor=1.0. + This re-enables it by setting the factor to the specified dp_size, so gradients are + averaged across FSDP ranks. This is needed for verl's loss scaling (loss * dp_size) + to work correctly. + + Args: + model: The model (or model part) to enable gradient division on. + dp_size: The data parallel size to use as the gradient divide factor. + """ + + for module in model.modules(): + if isinstance(module, FSDPModule): + module.set_gradient_divide_factor(float(dp_size)) + + +def get_attention_masks( + input_batch: torch.Tensor, + positions: torch.Tensor, + attn_type: str, +) -> AttentionMasksType: + match attn_type: + case "flex": + return _get_flex_attention_masks( + input_batch, + positions, + ) + case "varlen": + return _create_varlen_metadata_for_document( + input_batch, + positions, + ) + case _: + raise TypeError("Only varlen and flex attn masks are supported") + + +def _get_document_mask_mod(positions: torch.Tensor) -> _mask_mod_signature: + # Detect boundaries from position resets + first_dummy_value = positions[:, :1] - 1 + position_diff = torch.diff(positions, prepend=first_dummy_value, dim=-1) + sequence_indices = (position_diff != 1).cumsum(-1) # [batch, seq] + + def document_mask(b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor) -> torch.Tensor: + return sequence_indices[b, q_idx] == sequence_indices[b, kv_idx] + + return document_mask + + +def _get_flex_attention_masks( + input_batch: torch.Tensor, + positions: torch.Tensor, +) -> AttentionMasksType: + mask_mods = [get_causal_mask_mod()] + B = input_batch.shape[0] + mask_mods.append(_get_document_mask_mod(positions=positions)) + return create_attention_mask(and_masks(*mask_mods), B, None, input_batch.shape[1], input_batch.shape[1]) + + +def _create_varlen_metadata_for_document(input_batch: torch.Tensor, positions: torch.Tensor) -> VarlenMetadata: + """ + Creates cumulative sequence length indices needed for variable length attention + + Args: + input_batch: Input token IDs with shape [batch, seq]. + positions: Position IDs with shape [batch, seq]. Boundaries detected where + position diff != 1 (i.e., position resets). + + Returns: + VarlenMetadata containing cumulative sequence length indices for q, k, and max_seq_len + """ + batch_size, seq_len = input_batch.shape + device = input_batch.device + + # Detect boundaries from position resets (where diff != 1) + first_dummy_value = positions[:, :1] - 1 + position_diff = torch.diff(positions, prepend=first_dummy_value, dim=-1) + # boundary_mask[b, i] is True if position i starts a new document + boundary_mask = position_diff != 1 # [batch, seq] + boundary_mask[:, 0] = True + + cu_seqlens_list, all_seq_lengths = [], [] + offset = 0 + + for b in range(batch_size): + # Find positions where new documents start + boundary_positions = boundary_mask[b].nonzero(as_tuple=True)[0].to(torch.int32) + sample_cu_seqlens = torch.cat( + [ + boundary_positions, + torch.tensor([seq_len], dtype=torch.int32, device=device), + ] + ) + sample_cu_seqlens = torch.unique_consecutive(sample_cu_seqlens) + + seq_lengths = torch.diff(sample_cu_seqlens) + all_seq_lengths.append(seq_lengths) + + cu_seqlens_adjusted = sample_cu_seqlens[:-1] + offset + cu_seqlens_list.append(cu_seqlens_adjusted) + + offset += seq_len + + packed_cu_seqlens = torch.cat(cu_seqlens_list + [torch.tensor([offset], dtype=torch.int32, device=device)]) + + max_seqlen = 0 + if len(all_seq_lengths) > 0: + all_seq_lengths = torch.cat(all_seq_lengths) + # device to host sync but only done once per model forward + max_seqlen = all_seq_lengths.max().item() + + return VarlenMetadata( + cu_seq_q=packed_cu_seqlens, + cu_seq_k=packed_cu_seqlens, + max_q=max_seqlen, + max_k=max_seqlen, + ) + + +# Regex to parse: model.layers.{L}.mlp.experts.{E}.{weight_suffix} +_EXPERT_PATTERN = re.compile(r"\.layers\.(\d+)\..*\.experts\.(\d+)\.(.*)") + + +def _parse_expert_name(name: str) -> tuple[int, int, str] | None: + """Parse layer_id, expert_id, weight_suffix from expert param name.""" + match = _EXPERT_PATTERN.search(name) + if match: + return int(match.group(1)), int(match.group(2)), match.group(3) + return None + + +def _make_expert_name_template(name: str) -> str: + """Convert 'model.layers.0.mlp.experts.3.w1' -> 'model.layers.0.mlp.experts.{}.w1'""" + return _EXPERT_PATTERN.sub(lambda m: f".layers.{m.group(1)}.mlp.experts.{{}}.{m.group(3)}", name) + + +def iter_per_tensor_params_ep( + params: dict[str, Any], + device: int, + ep_group: torch.distributed.ProcessGroup, + ep_size: int, +) -> Generator[tuple[str, torch.Tensor], None, None]: + """Yield (name, tensor) pairs for weight sync with Expert Parallel. + + Gathers expert weights across EP ranks one (layer, weight_type) group + at a time to avoid OOM from materializing all experts simultaneously. + + Non-expert params are yielded first (with FSDP full_tensor() if needed), + then expert params are all-gathered per group and yielded individually. + + Args: + params: HF-format state dict with per-expert keys. Expert keys must + follow the pattern ``model.layers.{L}.mlp.experts.{E}.{suffix}``. + device: device ID to place tensors on. + ep_group: The EP process group for all-gather. + ep_size: Number of EP ranks. + """ + expert_params: dict[tuple[int, str], dict[int, tuple[str, Any]]] = defaultdict(dict) + non_expert_params: list[tuple[str, Any]] = [] + + for name, param in params.items(): + parsed = _parse_expert_name(name) if "mlp.experts." in name else None + if parsed is None: + non_expert_params.append((name, param)) + else: + layer_id, expert_id, weight_suffix = parsed + expert_params[(layer_id, weight_suffix)][expert_id] = (name, param) + + params.clear() + + # Yield non-expert params + for name, param in non_expert_params: + if isinstance(param, DTensor): + yield name, param.to(device, non_blocking=True).full_tensor().to(torch.bfloat16, non_blocking=True) + else: + yield name, param + del non_expert_params + + # Yield expert params with all-gather + for (layer_id, weight_suffix), experts_dict in sorted(expert_params.items()): + sorted_expert_ids = sorted(experts_dict.keys()) + + # Stack local expert weights + local_weights = [] + for eid in sorted_expert_ids: + _, param = experts_dict[eid] + if isinstance(param, DTensor): + param = param.to(device, non_blocking=True).full_tensor() + else: + param = param.to(device, non_blocking=True) + local_weights.append(param) + + name_template = _make_expert_name_template(experts_dict[sorted_expert_ids[0]][0]) + local_stacked = torch.stack(local_weights, dim=0) + + # All-gather across EP ranks + gathered_list = [torch.empty_like(local_stacked) for _ in range(ep_size)] + torch.distributed.all_gather(gathered_list, local_stacked, group=ep_group) + all_experts = torch.cat(gathered_list, dim=0) + + for expert_id in range(all_experts.shape[0]): + yield name_template.format(expert_id), all_experts[expert_id].to(torch.bfloat16).clone() + + del local_weights, local_stacked, gathered_list, all_experts + torch.cuda.empty_cache() diff --git a/verl/verl/workers/engine/utils.py b/verl/verl/workers/engine/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..47a35bcb0148f7b92bdf769b1ecad52596cca15c --- /dev/null +++ b/verl/verl/workers/engine/utils.py @@ -0,0 +1,160 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import random + +import numpy as np +import torch +from tensordict import TensorDict + +from verl.utils import tensordict_utils as tu +from verl.utils.dataset.dataset_utils import DatasetPadMode +from verl.utils.device import is_npu_available +from verl.utils.py_functional import append_to_dict +from verl.utils.seqlen_balancing import rearrange_micro_batches, restore_dynamic_batch + + +def enable_full_determinism(seed: int): + """ + Helper function for reproducibility in distributed training. + See https://pytorch.org/docs/stable/notes/randomness.html for details. + """ + + os.environ["PYTHONHASHSEED"] = str(seed) + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8" + os.environ["FLASH_ATTENTION_DETERMINISTIC"] = "1" + if is_npu_available: + # The environment variable required to enable deterministic mode on Ascend NPUs. + os.environ["HCCL_DETERMINISTIC"] = "true" + os.environ["CLOSE_MATMUL_K_SHIFT"] = "1" + + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.use_deterministic_algorithms(True, warn_only=True) + # Enable CUDNN deterministic mode + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + torch.backends.cudnn.enabled = False + if is_npu_available: + torch.npu.manual_seed(seed) + torch.npu.manual_seed_all(seed) + + +def prepare_micro_batches( + data: TensorDict, + dp_group=None, + num_batches_divided_by=None, + same_micro_num_in_dp=True, + min_num_micro_batch=None, + use_dynamic_bsz_balance=True, +): + """ + Prepare micro batches from data. + """ + use_dynamic_bsz = tu.get_non_tensor_data(data=data, key="use_dynamic_bsz", default=True) + sp_size = tu.get_non_tensor_data(data=data, key="sp_size", default=1) + + force_group_size = tu.get_non_tensor_data(data=data, key="force_group_size", default=1) + + if use_dynamic_bsz: + assert "max_token_len_per_gpu" in data.keys(), "max_token_len_per_gpu must be set when use_dynamic_bsz is True" + max_token_len_per_gpu = data["max_token_len_per_gpu"] + max_token_len = max_token_len_per_gpu * sp_size + micro_batches, batch_idx_list = rearrange_micro_batches( + data, + max_token_len=max_token_len, + dp_group=dp_group, + num_batches_divided_by=num_batches_divided_by, + same_micro_num_in_dp=same_micro_num_in_dp, + min_num_micro_batch=min_num_micro_batch, + use_dynamic_bsz_balance=use_dynamic_bsz_balance, + force_group_size=force_group_size, + ) + else: + total_data_size = len(data) + micro_batch_size_per_gpu = data["micro_batch_size_per_gpu"] + assert total_data_size % (force_group_size * micro_batch_size_per_gpu) == 0, ( + "data size must be divisible by force_group_size * micro_batch_size_per_gpu" + ) + micro_batches = tu.chunk_tensordict(data, total_data_size // (micro_batch_size_per_gpu * force_group_size)) + batch_idx_list = None + return micro_batches, batch_idx_list + + +def postprocess_batch_func(output_lst, indices, data: TensorDict): + """postprocess the output of a forward_backward_batch. + output_lst is a list of dict containing outputs for each micro-batch + reorder entropy and outputs. Return None for other pp ranks + only on last rank. It should be on every tp rank + + each losses_reduced contains 1. model_output, 2. loss, 3. metrics. + """ + + use_dynamic_bsz = tu.get_non_tensor_data(data=data, key="use_dynamic_bsz", default=True) + pad_mode = tu.get_non_tensor_data(data=data, key="pad_mode", default=DatasetPadMode.NO_PADDING) + assert pad_mode == DatasetPadMode.NO_PADDING, "postprocess_batch_func only support NO_PADDING pad_mode" + + # losses_reduced is a list of dict containing outputs for each micro-batch + # reorder entropy and outputs. Return None for other pp ranks + # only on last rank. It should be on every tp rank + + # losses_reduced contains 1. model_output, 2. loss, 3. metrics. + # We perform reverse + + model_output = {} + losses = [] + aggregated_metrics = {} + + # model output + for o in output_lst: + if "model_output" in o: + for key, val in o["model_output"].items(): + if key not in model_output: + model_output[key] = [] + model_output[key].append(val) + + # concat results from micro batches + for key, val in model_output.items(): + if pad_mode == DatasetPadMode.NO_PADDING: + tensors = [tensor for nt in model_output[key] for tensor in nt.unbind()] + model_output[key] = torch.nested.as_nested_tensor(tensors, layout=torch.jagged) + else: + raise NotImplementedError(f"pad_mode {pad_mode} not implemented") + + # reverse with dynamic bsz + if use_dynamic_bsz: + model_output[key] = restore_dynamic_batch(model_output[key], indices) + + # loss + for o in output_lst: + if "loss" in o: + losses.append(o["loss"]) + + # metrics + for o in output_lst: + if "metrics" in o: + metrics = o["metrics"] + append_to_dict(aggregated_metrics, metrics) + + output = { + "model_output": model_output, + "loss": losses, + "metrics": aggregated_metrics, + } + + return output diff --git a/verl/verl/workers/engine/veomni/__init__.py b/verl/verl/workers/engine/veomni/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..057facf27d33d3af60526eed5ef95784ee199878 --- /dev/null +++ b/verl/verl/workers/engine/veomni/__init__.py @@ -0,0 +1,17 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .transformer_impl import VeOmniEngine, VeOmniEngineWithLMHead + +__all__ = ["VeOmniEngine", "VeOmniEngineWithLMHead"] diff --git a/verl/verl/workers/engine/veomni/__pycache__/__init__.cpython-312.pyc b/verl/verl/workers/engine/veomni/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7e6e921c5ff3b8a6937a0aaa2364f7bdeaeeaad7 Binary files /dev/null and b/verl/verl/workers/engine/veomni/__pycache__/__init__.cpython-312.pyc differ diff --git a/verl/verl/workers/engine/veomni/__pycache__/transformer_impl.cpython-312.pyc b/verl/verl/workers/engine/veomni/__pycache__/transformer_impl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e8b9a33e56559e0e47abb4931f20e9a39f8a053e Binary files /dev/null and b/verl/verl/workers/engine/veomni/__pycache__/transformer_impl.cpython-312.pyc differ diff --git a/verl/verl/workers/engine/veomni/transformer_impl.py b/verl/verl/workers/engine/veomni/transformer_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..7476c0014c5c6374c6a58c9bfcb2f5dc96163413 --- /dev/null +++ b/verl/verl/workers/engine/veomni/transformer_impl.py @@ -0,0 +1,664 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import logging +from dataclasses import dataclass, field +from typing import Any, Callable, Optional, Sequence + +import torch +import torch.distributed as dist +from tensordict import TensorDict +from torch.distributed.tensor import DTensor +from veomni.arguments import OpsImplementationConfig +from veomni.distributed import parallel_state +from veomni.distributed.offloading import build_activation_offloading_context +from veomni.distributed.torch_parallelize import build_parallelize_model +from veomni.models.auto import build_foundation_model +from veomni.optim import build_lr_scheduler, build_optimizer +from veomni.utils.seqlen_pos_transform_utils import prepare_fa_kwargs_from_position_ids + +import verl.utils.torch_functional as verl_F +from verl.trainer.config import CheckpointConfig +from verl.utils import tensordict_utils as tu +from verl.utils.checkpoint.fsdp_checkpoint_manager import FSDPCheckpointManager +from verl.utils.device import get_device_id, get_device_name +from verl.utils.fsdp_utils import fsdp_version +from verl.utils.model import convert_weight_keys +from verl.utils.profiler import log_gpu_memory_usage +from verl.utils.ulysses import ( + get_ulysses_sequence_parallel_group, + set_ulysses_sequence_parallel_group, +) +from verl.workers.config import HFModelConfig, VeOmniEngineConfig, VeOmniOptimizerConfig + +from ..base import BaseEngineCtx, EngineRegistry +from ..fsdp.transformer_impl import FSDPEngine, FSDPEngineWithLMHead +from ..utils import enable_full_determinism, postprocess_batch_func, prepare_micro_batches +from .utils import ( + MOE_PARAM_HANDERS, + VL_TYPE2INDEX, + load_veomni_model_to_gpu, + load_veomni_optimizer, + offload_veomni_model_to_cpu, + offload_veomni_optimizer, +) + +logger = logging.getLogger(__file__) + + +class VeOmniEngine(FSDPEngine): + def __init__( + self, + model_config: HFModelConfig, + engine_config: VeOmniEngineConfig, + optimizer_config: VeOmniOptimizerConfig, + checkpoint_config: CheckpointConfig, + **kwargs, + ): + """ + Initialize the VeOmniEngine. + + Sets up distributed device meshes, LoRA, and offload policies based on config. + + Args: + config: Configuration object with VeOmni and model settings. + """ + + self.model_config = model_config + self.engine_config = engine_config + self.optimizer_config = optimizer_config + self.checkpoint_config = checkpoint_config + # VeOmniEngine only supports fsdp2. + self.data_parallel_mode = "fsdp2" + self.rank = dist.get_rank() + + fsdp_size = self.engine_config.fsdp_size + world_size = dist.get_world_size() + dp_size = world_size // self.engine_config.ulysses_parallel_size + + if fsdp_size < 0 or fsdp_size >= dp_size: + data_parallel_replicate_size = 1 + data_parallel_shard_size = dp_size + else: + if dp_size % fsdp_size != 0: + raise ValueError( + f"Data parallel size ({dp_size}) must be divisible by fsdp_size ({fsdp_size}). " + "Please adjust your parallel configuration." + ) + data_parallel_replicate_size = dp_size // fsdp_size + data_parallel_shard_size = fsdp_size + + parallel_state.init_parallel_state( + dp_size=dp_size, + dp_replicate_size=data_parallel_replicate_size, + dp_shard_size=data_parallel_shard_size, + extra_parallel_sizes=(self.engine_config.expert_parallel_size,), + ulysses_size=self.engine_config.ulysses_parallel_size, + dp_mode=self.data_parallel_mode, + ) + + if self.engine_config.full_determinism: + enable_full_determinism(seed=self.engine_config.seed) + + self.use_remove_padding = self.model_config.use_remove_padding + + self._is_offload_param = self.engine_config.param_offload + self._is_offload_optimizer = self.engine_config.optimizer_offload + self._is_lora = self.model_config.lora_rank > 0 + + self.use_ulysses_sp = parallel_state.get_parallel_state().sp_enabled + self.ulysses_sequence_parallel_size = self.engine_config.ulysses_parallel_size + + if self.use_ulysses_sp: + self.ulysses_parallel_group = parallel_state.get_parallel_state().device_mesh["sp"].get_group() + else: + self.ulysses_parallel_group = None + + if self.engine_config.entropy_from_logits_with_chunking: + entropy_from_logits = verl_F.entropy_from_logits_with_chunking + else: + entropy_from_logits = verl_F.entropy_from_logits + + self.compute_entropy_from_logits = ( + torch.compile(entropy_from_logits, dynamic=True) + if self.engine_config.use_torch_compile # use torch compile by default + else entropy_from_logits + ) + + def initialize(self): + """ + Build the model, optimizer, and learning rate scheduler under VeOmni. + + Applies device, dtype, and precision configurations, including mixed precision. + Sets up checkpoint manager and FLOPs counter. + """ + self._build_model_optimizer() + + self.checkpoint_manager = FSDPCheckpointManager( + model=self.module, + optimizer=self.optimizer, + lr_scheduler=self.lr_scheduler, + processing_class=self.model_config.get_processor(), + checkpoint_config=self.checkpoint_config, + trust_remote_code=self.model_config.trust_remote_code, + ) + + self.to( + device="cpu", + model=self._is_offload_param, + optimizer=self._is_offload_optimizer, + grad=self._is_offload_optimizer, + ) + + log_gpu_memory_usage("After offload model/optimizer/grad during init", logger=logger) + + def _build_optimizer(self, module): + optimizer = build_optimizer( + module, + lr=self.optimizer_config.lr, + betas=self.optimizer_config.betas, + weight_decay=self.optimizer_config.weight_decay, + optimizer_type=self.optimizer_config.optimizer, + ) + get_optimizer_pre_hook = getattr(module, "get_optimizer_pre_hook", None) + if get_optimizer_pre_hook is not None: + optimizer_pre_hook = get_optimizer_pre_hook(module, module.config, self.data_parallel_mode) + optimizer.register_step_pre_hook(optimizer_pre_hook) + + return optimizer + + def _build_lr_scheduler(self, optimizer): + optim_config = self.optimizer_config + lr_scheduler = build_lr_scheduler( + optimizer, + train_steps=optim_config.total_training_steps, + lr=optim_config.lr, + lr_min=optim_config.lr_min, + lr_decay_style=optim_config.lr_scheduler_type, + lr_decay_ratio=optim_config.lr_decay_ratio, + lr_warmup_ratio=optim_config.lr_warmup_steps_ratio, + lr_start=optim_config.lr_start, + ) + + return lr_scheduler + + def _build_model_optimizer(self): + # build_foundation_model runs apply_ops_config(ops_implementation) + # before constructing the model, so per-model device_patch files see + # the resolved kernel backends. + ops_implementation = OpsImplementationConfig( + attn_implementation=self.engine_config.attn_implementation, + moe_implementation=self.engine_config.moe_implementation, + cross_entropy_loss_implementation=self.engine_config.cross_entropy_loss_implementation, + rms_norm_implementation=self.engine_config.rms_norm_implementation, + swiglu_mlp_implementation=self.engine_config.swiglu_mlp_implementation, + rotary_pos_emb_implementation=self.engine_config.rotary_pos_emb_implementation, + load_balancing_loss_implementation=self.engine_config.load_balancing_loss_implementation, + ) + + # Load base model with specified configuration and dtype + module = build_foundation_model( + config_path=self.model_config.local_hf_config_path, + weights_path=self.model_config.local_path, + torch_dtype="float32" if self.engine_config.mixed_precision else "bfloat16", + attn_implementation=self.engine_config.attn_implementation, + ops_implementation=ops_implementation, + init_device=self.engine_config.init_device, + ) + log_gpu_memory_usage("After load base model", logger=logger) + + # Applies parallel strategies to the model. + log_gpu_memory_usage("Before parallelize model", logger=logger) + module = build_parallelize_model( + module, + init_device=self.engine_config.init_device, + weights_path=self.model_config.local_path, + enable_full_shard=self.engine_config.enable_full_shard, + enable_mixed_precision=self.engine_config.mixed_precision, + enable_gradient_checkpointing=self.model_config.enable_gradient_checkpointing, + enable_fsdp_offload=self.engine_config.enable_fsdp_offload, + basic_modules=list( + set(getattr(module, "_no_split_modules", None) or []) | set(self.engine_config.basic_modules) + ), + enable_reentrant=self.engine_config.enable_reentrant, + enable_forward_prefetch=self.engine_config.forward_prefetch, + ) + log_gpu_memory_usage("After parallelize model", logger=logger) + + if not self.engine_config.forward_only: + # Initialize optimizer with model parameters and config settings + optimizer = self._build_optimizer(module) + # Create learning rate scheduler with warmup and decay settings + lr_scheduler = self._build_lr_scheduler(optimizer) + else: + optimizer = None + lr_scheduler = None + + self.module = module + self.optimizer = optimizer + self.lr_scheduler = lr_scheduler + self.model_fwd_context, self.model_bwd_context = build_activation_offloading_context( + self.model_config.enable_activation_offload, + self.model_config.enable_gradient_checkpointing, + self.engine_config.activation_gpu_limit, + ) + + def optimizer_step(self): + """ + Perform an optimization step using the optimizer. + """ + if hasattr(self.module, "clip_grad_norm_"): + grad_norm = self.module.clip_grad_norm_(self.optimizer_config.clip_grad) + else: + grad_norm = torch.nn.utils.clip_grad_norm_(self.module.parameters(), self.optimizer_config.clip_grad) + + if isinstance(grad_norm, DTensor): + grad_norm = grad_norm.full_tensor() + + # if grad_norm is not finite, skip the update + if not torch.isfinite(grad_norm): + print(f"WARN: grad_norm is not finite: {grad_norm}") + self.optimizer.zero_grad() + else: + self.optimizer.step() + return grad_norm.item() + + def forward_backward_batch(self, data: TensorDict, loss_function: Callable, forward_only=False) -> Any: + """ + Perform a forward pass and optionally a backward pass on a batch of data. + + Args: + data: The input data for the forward pass, typically containing tensors and metadata. + loss_function: The loss function to optimize. See `verl.workers.roles.utils.losses` for examples. + forward_only: If True, perform only the forward pass. If False, perform forward and backward pass. + + Returns: + Any: The output of the forward pass, which can be used for loss computation or other purposes. + """ + tu.assign_non_tensor(data, sp_size=parallel_state.get_parallel_state().ulysses_size) + + # compute num_tokens in global batch for loss normalization + batch_num_tokens = data["loss_mask"].sum().to(get_device_id()) + torch.distributed.all_reduce( + batch_num_tokens, op=torch.distributed.ReduceOp.SUM, group=self.get_data_parallel_group() + ) + tu.assign_non_tensor(data, batch_num_tokens=batch_num_tokens.item()) + tu.assign_non_tensor(data, dp_size=self.get_data_parallel_size()) + + micro_batches, indices = prepare_micro_batches( + data=data, dp_group=self.get_data_parallel_group(), same_micro_num_in_dp=True + ) + + output_lst = [] + + for micro_batch in micro_batches: + with self.model_fwd_context: + loss, meta_info = self.forward_step(micro_batch, loss_function=loss_function, forward_only=forward_only) + if not forward_only: + with self.model_bwd_context: + loss.backward() + + output_lst.append(meta_info) + + return postprocess_batch_func(output_lst=output_lst, indices=indices, data=data) + + def get_data_parallel_rank(self): + return parallel_state.get_parallel_state().device_mesh.get_local_rank("dp") + + def get_data_parallel_size(self): + return torch.distributed.get_world_size() // parallel_state.get_parallel_state().ulysses_size + + def get_data_parallel_group(self): + if parallel_state.get_parallel_state().ulysses_size > 1: + return parallel_state.get_parallel_state().device_mesh.get_group(mesh_dim="dp") + else: + return torch.distributed.group.WORLD + + def get_model_parallel_group(self): + raise NotImplementedError + + def get_context_parallel_group(self): + raise NotImplementedError + + def is_mp_src_rank_with_outputs(self): + """ + Whether the current rank is the first rank in model parallel group that contains model outputs + """ + if parallel_state.get_parallel_state().ulysses_size > 1: + is_collect = parallel_state.get_parallel_state().device_mesh["ulysses"].get_local_rank() == 0 + else: + is_collect = True + return is_collect + + def train_mode(self, **kwargs): + """ + Return a context manager that switches to training mode with VeOmni-specific handling. + + Includes parameter and optimizer offload entry/exit. + """ + return EngineTrainModeCtx(self, **kwargs) + + def eval_mode(self, **kwargs): + """ + Return a context manager that switches to evaluation mode with VeOmni-specific handling. + + Includes activation offload entry/exit. + """ + return EngineEvalModeCtx(self, **kwargs) + + def to(self, device: str, model: bool = True, optimizer: bool = True, grad: bool = True): + """ + Move model parameters, optimizer states, or both to the specified device. + Note that this function executes irrespective of offload config. It serves as manual control. + + Args: + device: Target device identifier. + model: If True, move the model. + optimizer: If True, move the optimizer states. + """ + super(FSDPEngine, self).to(device=device, model=model, optimizer=optimizer, grad=grad) + + device_name = get_device_name() + + assert device in (device_name, "cpu") + if device == device_name: + if model: + load_veomni_model_to_gpu(self.module) + if optimizer and self.optimizer is not None: + load_veomni_optimizer(self.optimizer, device) + elif device == "cpu": + if model: + offload_veomni_model_to_cpu(self.module) + if optimizer and self.optimizer is not None: + offload_veomni_optimizer(self.optimizer) + else: + raise ValueError(f"Invalid device type: {device}") + + def save_checkpoint( + self, + local_path: str, + hdfs_path: Optional[str] = None, + global_step: int = 0, + max_ckpt_to_keep: Optional[int] = None, + **kwargs, + ) -> None: + """ + Save VeOmni checkpoint, handling parameter offload as needed. + """ + origin_module_device = next(self.module.parameters()).device.type + if self._is_offload_param or origin_module_device == "cpu": + load_veomni_model_to_gpu(self.module) + + self.checkpoint_manager.save_checkpoint( + local_path=local_path, hdfs_path=hdfs_path, global_step=global_step, max_ckpt_to_keep=max_ckpt_to_keep + ) + + torch.distributed.barrier() + if self._is_offload_param: + offload_veomni_model_to_cpu(self.module) + + def load_checkpoint( + self, local_path: str, hdfs_path: Optional[str] = None, del_local_after_load: int = True, **kwargs + ) -> None: + """ + Load VeOmni checkpoint, restoring parameters and optimizer state. + """ + if self._is_offload_param: + load_veomni_model_to_gpu(self.module) + + self.checkpoint_manager.load_checkpoint( + local_path=local_path, hdfs_path=hdfs_path, del_local_after_load=del_local_after_load + ) + + torch.distributed.barrier() + if self._is_offload_param: + offload_veomni_model_to_cpu(self.module) + + if self._is_offload_optimizer: + offload_veomni_optimizer(self.optimizer) + + def get_per_tensor_param(self, **kwargs): + load_veomni_model_to_gpu(self.module) + + params = self.module.state_dict() + params = convert_weight_keys(params, getattr(self.module, "_fsdp_wrapped_module", self.module)) + + if self._is_offload_param: + offload_veomni_model_to_cpu(self.module) + + device = get_device_id() + ps = parallel_state.get_parallel_state() + model_type = getattr(self.module.config, "model_type", "default") + process_func = MOE_PARAM_HANDERS.get(model_type, lambda n, t: iter([(n, t)])) + + def param_generator(): + for name, param in params.items(): + unsharded_tensor = param.full_tensor() if isinstance(param, DTensor) else param + + is_expert_layer = "mlp.experts." in name + is_proj = any(p in name for p in ["down_proj", "gate_proj", "up_proj", "gate_up_proj"]) + + if is_expert_layer and is_proj and ps.ep_enabled: + output_shape = list(unsharded_tensor.shape) + output_shape[0] *= ps.extra_parallel_sizes["ep"] + stacked_tensor = torch.empty(output_shape, dtype=unsharded_tensor.dtype, device=device) + + # all gather expert tensors [32, H, I] -> [128, H, I] + torch.distributed.all_gather_into_tensor(stacked_tensor, unsharded_tensor, group=ps.ep_group) + yield from process_func(name, stacked_tensor) + + del stacked_tensor + else: + if is_expert_layer: + yield from process_func(name, unsharded_tensor) + else: + yield name, unsharded_tensor + + # TODO: support VeOmni LoRA + return param_generator(), None + + +class EngineEvalModeCtx(BaseEngineCtx): + def __init__(self, engine: VeOmniEngine, **kwargs): + super().__init__(engine=engine, mode="eval", **kwargs) + + def __enter__(self): + assert isinstance(self.engine, VeOmniEngine) + super().__enter__() + self.prev_sp_group = get_ulysses_sequence_parallel_group() + set_ulysses_sequence_parallel_group(self.engine.ulysses_parallel_group) + self.engine.module.train() + + def __exit__(self, exc_type, exc_value, traceback): + assert isinstance(self.engine, VeOmniEngine) + set_ulysses_sequence_parallel_group(self.prev_sp_group) + + # https://pytorch.org/docs/stable/notes/fsdp.html#fsdp-notes + # unshard the root FSDP module + if parallel_state.get_parallel_state().dp_shard_size > 1: + if fsdp_version(self.engine.module) == 1: + self.engine.module._handle.reshard(True) + elif fsdp_version(self.engine.module) == 2: + self.engine.module.reshard() + + super().__exit__(exc_type, exc_value, traceback) + + +class EngineTrainModeCtx(BaseEngineCtx): + def __init__(self, engine: VeOmniEngine, **kwargs): + super().__init__(engine=engine, mode="train", **kwargs) + + def __enter__(self): + assert isinstance(self.engine, VeOmniEngine) + super().__enter__() + self.prev_sp_group = get_ulysses_sequence_parallel_group() + set_ulysses_sequence_parallel_group(self.engine.ulysses_parallel_group) + # TODO: Switch to eval mode after Integrating the CI environment + # VeOmni (ref: https://github.com/ByteDance-Seed/VeOmni/pull/421) + self.engine.module.train() + + def __exit__(self, exc_type, exc_value, traceback): + assert isinstance(self.engine, VeOmniEngine) + set_ulysses_sequence_parallel_group(self.prev_sp_group) + self.engine.optimizer_zero_grad() + super().__exit__(exc_type, exc_value, traceback) + + +@dataclass +class OmniSequenceShardCollator: + """ + Data collator to chunk inputs along the sequence length. + """ + + # features to slice sequence dimension + sp_slice_features: dict[str, int] = field( + default_factory=lambda: { + "input_ids": -1, + "labels": -1, + "pixel_values": 0, + "pixel_values_videos": 0, + }, + metadata={"help": "features to slice sequence dimension."}, + ) + + # features to padding sequence dimension + padding_features: dict[str, int] = field( + default_factory=lambda: { + "pixel_values": 0, + "pixel_values_videos": 0, + }, + metadata={"help": "features to padding sequence dimension."}, + ) + + # padding scale for padding features + padding_scale: dict[str, int] = field( + default_factory=lambda: {"pixel_values": 4, "pixel_values_videos": 4}, + metadata={"help": "padding scale for padding features."}, + ) + + def __post_init__(self): + self.sp_size = parallel_state.get_parallel_state().sp_size + self.sp_rank = parallel_state.get_parallel_state().sp_rank + + def sp_slice(self, feature: torch.Tensor, dim: int = -1) -> dict[str, "torch.Tensor"]: + seq_length = feature.size(dim) + sp_chunk_size = (seq_length + self.sp_size - 1) // self.sp_size + return feature.narrow(dim, self.sp_rank * sp_chunk_size, sp_chunk_size) + + def sp_padding( + self, tensor: "torch.Tensor", dim: int = -1, pad_value: int = 0, pad_scale: int = 1 + ) -> "torch.Tensor": + """ + Pads a tensor with pad_length to aligns tensor with sp size. + """ + seq_length = tensor.size(dim) + scale_sp_size = self.sp_size * pad_scale + + sp_chunk_size = (seq_length + scale_sp_size - 1) // scale_sp_size + pad_size = sp_chunk_size * scale_sp_size - seq_length + if pad_size == 0: + return tensor + + pad_shape = list(tensor.shape) + pad_shape[dim] = pad_size + pad = torch.full(pad_shape, fill_value=pad_value, dtype=tensor.dtype, device=tensor.device) + return torch.cat((tensor, pad), dim=dim) + + def __call__(self, batch: Sequence[dict[str, "torch.Tensor"]]) -> dict[str, "torch.Tensor"]: + for key in batch.keys(): + if key in self.padding_features.keys(): + batch[key] = self.sp_padding( + batch[key], + dim=self.sp_slice_features.get(key, -1), + pad_value=self.padding_features[key], + pad_scale=self.padding_scale.get(key, 1), + ) + + # sp slice + for key in batch.keys(): + if key in self.sp_slice_features.keys(): + batch[key] = self.sp_slice(batch[key], dim=self.sp_slice_features[key]) + + return batch + + +def _prepare_veomni_flash_attention_kwargs(position_ids: torch.Tensor) -> dict[str, torch.Tensor | int]: + """Normalize packed position_ids layout and derive varlen FlashAttention kwargs. + + Supported formats for use_remove_padding=true: + - 2D: (1, total_nnz) - standard packed format + - 3D: (rope_dim, 1, total_nnz) - VeRL mRoPE packed format + """ + if position_ids.dim() == 2: + # (1, total_nnz) - standard packed format + fa_position_ids = position_ids + elif position_ids.dim() == 3: + # (rope_dim, 1, total_nnz) - VeRL mRoPE packed format + if position_ids.shape[1] == 1: + fa_position_ids = position_ids[0] + else: + raise ValueError( + f"Unsupported 3D position_ids shape: {tuple(position_ids.shape)}, expected (rope_dim, 1, total_nnz)" + ) + else: + raise ValueError( + f"Unsupported position_ids rank: {position_ids.dim()}, " + f"expected 2 (1, total_nnz) or 3 (rope_dim, 1, total_nnz)" + ) + + (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = prepare_fa_kwargs_from_position_ids(fa_position_ids) + return { + "cu_seq_lens_q": cu_seq_lens_q, + "cu_seq_lens_k": cu_seq_lens_k, + "max_length_q": max_length_q, + "max_length_k": max_length_k, + } + + +@EngineRegistry.register(model_type="language_model", backend=["veomni"], device=["cuda", "npu"]) +class VeOmniEngineWithLMHead(VeOmniEngine, FSDPEngineWithLMHead): + def prepare_model_inputs(self, micro_batch: TensorDict): + model_inputs, output_args = super().prepare_model_inputs(micro_batch) + input_ids_rmpad = model_inputs["input_ids"] + sp_enabled = parallel_state.get_parallel_state().sp_enabled + sp_shard_collator = OmniSequenceShardCollator() if sp_enabled else None + + if self.module.config.model_type in VL_TYPE2INDEX.keys(): + image_mask = input_ids_rmpad == VL_TYPE2INDEX[self.module.config.model_type]["IMAGE_INPUT_INDEX"] + video_mask = input_ids_rmpad == VL_TYPE2INDEX[self.module.config.model_type]["VIDEO_INPUT_INDEX"] + model_inputs.update({"image_mask": image_mask, "video_mask": video_mask}) + + if sp_enabled: + sp_shard_collator(model_inputs) + + use_remove_padding = tu.get_non_tensor_data(data=micro_batch, key="use_remove_padding", default=True) + if use_remove_padding and model_inputs.get("position_ids", None) is not None: + model_inputs.update(_prepare_veomni_flash_attention_kwargs(model_inputs["position_ids"])) + if sp_enabled: + model_inputs["position_ids"] = sp_shard_collator.sp_slice(model_inputs["position_ids"], dim=-1) + + # Activate VeOmni's chunk_logprobs path: ForCausalLMLoss short-circuits + # to per-token log_probs/entropy on return_log_probs=True. Pass the + # already-rolled labels as shift_labels so chunk_logprobs skips its + # internal causal shift and the output seq length matches the input — + # prepare_model_outputs().squeeze(0) then lands at (total_nnz,). + use_fused_kernels = tu.get_non_tensor_data(data=micro_batch, key="use_fused_kernels", default=False) + if use_fused_kernels and use_remove_padding: + shift_labels = output_args["input_ids_rmpad_rolled"].unsqueeze(0) + model_inputs["labels"] = input_ids_rmpad + model_inputs["shift_labels"] = shift_labels + model_inputs["return_log_probs"] = True + + return model_inputs, output_args diff --git a/verl/verl/workers/engine/veomni/utils.py b/verl/verl/workers/engine/veomni/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..b1ac2095c0c12ce74ad321aa1a88b946f37e5658 --- /dev/null +++ b/verl/verl/workers/engine/veomni/utils.py @@ -0,0 +1,120 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from verl.utils.device import get_device_id, get_torch_device + +VL_TYPE2INDEX = { + "qwen2_5_vl": { + "IMAGE_INPUT_INDEX": 151655, + "VIDEO_INPUT_INDEX": 151656, + }, + "qwen3_vl": { + "IMAGE_INPUT_INDEX": 151655, + "VIDEO_INPUT_INDEX": 151656, + }, + "qwen3_vl_moe": { + "IMAGE_INPUT_INDEX": 151655, + "VIDEO_INPUT_INDEX": 151656, + }, + "qwen3_5": { + "IMAGE_INPUT_INDEX": 248056, + "VIDEO_INPUT_INDEX": 248057, + }, + "qwen3_5_moe": { + "IMAGE_INPUT_INDEX": 248056, + "VIDEO_INPUT_INDEX": 248057, + }, +} + + +@torch.no_grad() +def offload_veomni_model_to_cpu(model, empty_cache: bool = True): + from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState + from torch.distributed.fsdp._fully_shard._fsdp_state import _get_module_fsdp_state + + for module in model.modules(): + state = _get_module_fsdp_state(module) + if state is None: + continue + fsdp_param_group = state._fsdp_param_group + + if fsdp_param_group is None: + continue + + fsdp_param_group._training_state = TrainingState.IDLE + + model.reshard() + model.cpu() + if empty_cache: + get_torch_device().empty_cache() + + +@torch.no_grad() +def load_veomni_model_to_gpu(model): + device = get_device_id() + model.to(device) + + +@torch.no_grad() +def offload_veomni_optimizer(optimizer): + optimizers = [] + # Check if this is a MultiOptimizer (for ep and non-ep parameters when ep+fsdp2 is enabled) + if hasattr(optimizer, "_is_multi_optimizer") and optimizer._is_multi_optimizer: + optimizers.extend(optimizer.optimizers_dict.values()) + else: + optimizers.append(optimizer) + + for opt in optimizers: + if not opt.state: + continue + for param_group in opt.param_groups: + for param in param_group["params"]: + state = opt.state[param] + for key, value in state.items(): + if isinstance(value, torch.Tensor): + state[key] = value.to("cpu", non_blocking=True) + + +@torch.no_grad() +def load_veomni_optimizer(optimizer, device_id): + optimizers = [] + # Check if this is a MultiOptimizer (for ep and non-ep parameters when ep+fsdp2 is enabled) + if hasattr(optimizer, "_is_multi_optimizer") and optimizer._is_multi_optimizer: + optimizers.extend(optimizer.optimizers_dict.values()) + else: + optimizers.append(optimizer) + + for opt in optimizers: + if not opt.state: + continue + for param_group in opt.param_groups: + for param in param_group["params"]: + state = opt.state[param] + for key, value in state.items(): + if isinstance(value, torch.Tensor): + state[key] = value.to(device_id, non_blocking=True) + + +def _map_moe_params_qwen3_moe(name, tensor): + for i in range(tensor.size(0)): + new_key = name.replace("mlp.experts.", f"mlp.experts.{i}.") + ".weight" + yield new_key, tensor[i].to(get_device_id(), non_blocking=True) + + +MOE_PARAM_HANDERS = { + "qwen3_moe": _map_moe_params_qwen3_moe, + "deepseek_v3": _map_moe_params_qwen3_moe, +} diff --git a/verl/verl/workers/engine_workers.py b/verl/verl/workers/engine_workers.py new file mode 100644 index 0000000000000000000000000000000000000000..b792534a9fca6aa1a68a19e1285608c37a4bc896 --- /dev/null +++ b/verl/verl/workers/engine_workers.py @@ -0,0 +1,750 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import functools +import logging +import os +from contextlib import nullcontext +from copy import deepcopy +from functools import partial +from itertools import chain +from typing import Optional + +import psutil +import torch +from codetiming import Timer +from omegaconf import DictConfig, open_dict +from tensordict import NonTensorData, TensorDict +from torch.distributed.device_mesh import init_device_mesh + +from verl.checkpoint_engine import CheckpointEngineRegistry +from verl.single_controller.base import Worker +from verl.single_controller.base.decorator import Dispatch, make_nd_compute_dataproto_dispatch_fn, register +from verl.trainer.distillation import distillation_ppo_loss, is_distillation_enabled +from verl.utils import tensordict_utils as tu +from verl.utils.config import omega_conf_to_dataclass +from verl.utils.device import get_device_name, get_torch_device, set_expandable_segments +from verl.utils.distributed import initialize_global_process_group_ray, set_numa_affinity +from verl.utils.flops_counter import FlopsCounter +from verl.utils.import_utils import import_external_libs +from verl.utils.memory_utils import aggressive_empty_cache +from verl.utils.metric.utils import Metric +from verl.utils.profiler import DistProfiler, DistProfilerExtension, ProfilerConfig, log_gpu_memory_usage +from verl.utils.py_functional import append_to_dict +from verl.utils.tensordict_utils import maybe_fix_3d_position_ids +from verl.utils.torch_functional import allgather_dict_into_dict +from verl.workers.config import ( + ActorConfig, + DistillationConfig, + HFModelConfig, + MtpConfig, + RolloutConfig, + TrainingWorkerConfig, +) +from verl.workers.rollout.base import BaseRollout, get_rollout_class +from verl.workers.utils.losses import ppo_loss + +logger = logging.getLogger(__file__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + + +def _with_routing_replay_flag(enabled: bool): + """Decorator to set 'enable_routing_replay' flag on the data TensorDict.""" + + def decorator(func): + @functools.wraps(func) + def wrapper(self, data: TensorDict, *args, **kwargs): + if self.enable_routing_replay: + tu.assign_non_tensor_data(data, "enable_routing_replay", enabled) + return func(self, data, *args, **kwargs) + + return wrapper + + return decorator + + +class TrainingWorker(Worker, DistProfilerExtension): + """ + TrainingWorker provides a Tinker-like API (https://thinkingmachines.ai/tinker/) as a RayWorkerGroup + to a single controller. Currently, we only provide more coarse grained APIs, + and do not provide exact APIs as Tinker does. But this can be added in the future. + """ + + def __init__(self, config: TrainingWorkerConfig): + Worker.__init__(self) + + from verl.workers.engine import BaseEngine, EngineRegistry + + initialize_global_process_group_ray(timeout_second=None) + + set_numa_affinity() + + self.config = config + self.model_config = self.config.model_config + self.engine_config = self.config.engine_config + self.optimizer_config = self.config.optimizer_config + self.checkpoint_config = self.config.checkpoint_config + self.device_name = get_device_name() + + if self.engine_config is None: + assert self.optimizer_config is None + if self.config.auto_select_engine_optim_fn is None: + raise ValueError( + "engine_config is not provided and auto_select_engine_optim_fn is not set. " + "Cannot determine engine backend." + ) + # Support automatically select engine backend given model config + self.engine_config, self.optimizer_config = self.config.auto_select_engine_optim_fn( + self.model_config, self.device_name + ) + + # we use the one defined in model + # TODO: this is not elegant and should refactor later + self.engine_config.use_remove_padding = self.model_config.get("use_remove_padding", False) + self.engine_config.use_fused_kernels = self.model_config.get("use_fused_kernels", False) + + self.profiler_config = self.config.profiler_config + if self.profiler_config is not None: + self.profiler_tool_config = self.profiler_config.tool_config.get(self.profiler_config.tool, {}) + else: + self.profiler_tool_config = None + + DistProfilerExtension.__init__( + self, DistProfiler(rank=self.rank, config=self.profiler_config, tool_config=self.profiler_tool_config) + ) + + self.model_config.model_type = self.config.model_type + self.engine: BaseEngine = EngineRegistry.new( + model_type=self.config.model_type, + backend=self.engine_config.strategy, + model_config=self.model_config, + engine_config=self.engine_config, + optimizer_config=self.optimizer_config, + checkpoint_config=self.checkpoint_config, + ) + + # build dispatch info + self._register_dispatch_collect_info( + mesh_name="train", + dp_rank=self.engine.get_data_parallel_rank(), + is_collect=self.engine.is_mp_src_rank_with_outputs(), + ) + + if hasattr(self.model_config, "hf_config"): + self.flops_counter = FlopsCounter(self.model_config.hf_config) + else: + self.flops_counter = None + + self.loss_fn = None + + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def to(self, device, model=True, optimizer=True, grad=True): + """Manual control of load/offload""" + assert device in ["cpu", "device"] + + if device == "device": + device = get_device_name() + + self.engine.to(device=device, model=model, optimizer=optimizer, grad=grad) + + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def set_loss_fn(self, loss_fn): + self.loss_fn = loss_fn + + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def reset(self): + """ + Reset the model engine to the initial state. If the engine is not initialized, + we initialize it. Otherwise, reload ckpt and reset states + """ + self.engine.initialize() + + def _postprocess_output(self, output, *, global_token_num, delta_time, forward_only, images_seqlens): + """ + + Args: + output: a dictionary containing loss, model_outputs and metrics + + Returns: + + """ + + metrics: dict = output.pop("metrics") + # perform all gather in dp group to ensure that it's correct. + # Here each metric in metrics can be a list (micro-batch metrics) or a singleton + # we should always sum the loss of each micro-batch as we scale by global_bsz/global_token + loss = torch.sum(torch.tensor(output.pop("loss"), device=self.device_name)) + dp_group = self.engine.get_data_parallel_group() + if dp_group is not None: + torch.distributed.all_reduce(loss, op=torch.distributed.ReduceOp.AVG, group=dp_group) + loss = loss.item() + + # For grad_norm, we do not perform all reduce because it is already been done when clipping grad + grad_norm = metrics.pop("grad_norm", None) + if isinstance(grad_norm, torch.Tensor): + grad_norm = grad_norm.detach().item() + lr = metrics.pop("lr", None) + + # For other metrics, we perform all gather in dp group (only if DP > 1) + if dp_group is not None: + final_metrics = allgather_dict_into_dict(data=metrics, group=dp_group) + else: + final_metrics = metrics + final_metrics["loss"] = loss + if grad_norm is not None: + final_metrics["grad_norm"] = grad_norm + if lr is not None: + final_metrics["lr"] = lr + + # log memory + final_metrics["perf/max_memory_allocated_gb"] = get_torch_device().max_memory_allocated() / (1024**3) + final_metrics["perf/max_memory_reserved_gb"] = get_torch_device().max_memory_reserved() / (1024**3) + final_metrics["perf/cpu_memory_used_gb"] = psutil.virtual_memory().used / (1024**3) + + # TODO: confirm the mtp loss IS same across dp + for k, v in final_metrics.items(): + if k.startswith("mtp_losses"): + flatten_v = [sublist[0] for sublist in v] # sublist should be single element + final_metrics[k] = sum(flatten_v) / len(flatten_v) + # compute mfu + if global_token_num is not None and self.flops_counter is not None: + estimated_flops, promised_flops = self.flops_counter.estimate_flops( + global_token_num, delta_time, images_seqlens=images_seqlens + ) + final_metrics["mfu"] = estimated_flops / promised_flops / torch.distributed.get_world_size() + if forward_only: + final_metrics["mfu"] /= 3.0 + # model outputs + model_output = output.pop("model_output", {}) + # We only return final_metrics + final_output = tu.get_tensordict(tensor_dict=model_output, non_tensor_dict={"metrics": final_metrics}) + return final_output + + @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="train"), blocking=False) + def train_mini_batch(self, data: TensorDict) -> TensorDict: + """Split a batch into N mini-batches run for multiple epochs + + Args: + data: + + Returns: + + """ + maybe_fix_3d_position_ids(data) + batch_size_per_dp = data.shape[0] + disable_auto_offload = tu.pop(data, key="disable_auto_offload", default=False) + mini_batch_size = tu.pop(data, key="mini_batch_size", default=None) + num_mini_batch = tu.pop(data, key="num_mini_batch", default=None) + epochs = tu.pop(data, key="epochs", default=1) + seed = tu.pop(data, key="seed", default=42) + dataloader_kwargs = tu.pop(data, key="dataloader_kwargs", default={}) + + assert mini_batch_size is not None or num_mini_batch is not None + + if mini_batch_size is None: + assert batch_size_per_dp % num_mini_batch == 0, f"Got {batch_size_per_dp=} and {num_mini_batch=}" + mini_batch_size_per_gpu = batch_size_per_dp // num_mini_batch + else: + assert mini_batch_size % self.engine.get_data_parallel_size() == 0, ( + f"Got {mini_batch_size=} and {self.engine.get_data_parallel_size()=}" + ) + mini_batch_size_per_gpu = mini_batch_size // self.engine.get_data_parallel_size() + + # make iterator + dataloader = tu.make_iterator( + data, + mini_batch_size=mini_batch_size_per_gpu, + epochs=epochs, + seed=seed + self.engine.get_data_parallel_rank(), + dataloader_kwargs=dataloader_kwargs, + ) + + with ( + self.engine.train_mode(disable_auto_offload=disable_auto_offload), + Timer(name="train_batch", logger=None), + ): + # update + output_lst = [] + total_num_iterations = data.shape[0] // mini_batch_size_per_gpu * epochs + + for batch_idx, mini_batch_td in enumerate(dataloader): + # add global token num + if "input_ids" in mini_batch_td: + global_token_num = mini_batch_td["input_ids"].offsets().diff().tolist() # (total_nnz,) + # allgather from dp rank + global_token_num_output = [None] * torch.distributed.get_world_size( + self.engine.get_data_parallel_group() + ) + torch.distributed.all_gather_object( + global_token_num_output, global_token_num, self.engine.get_data_parallel_group() + ) + global_token_num = [x for xs in global_token_num_output for x in xs] + else: + global_token_num = None + + tu.assign_non_tensor( + mini_batch_td, + global_token_num=NonTensorData(global_token_num), + update_lr_scheduler=batch_idx == total_num_iterations - 1, + disable_auto_offload=True, + ) + actor_output = self.train_batch(mini_batch_td) + output_lst.append(actor_output) + + if self.engine.is_mp_src_rank_with_outputs(): + actor_output = [tu.get(output, "metrics") for output in output_lst] + metrics = {} + for output in actor_output: + for key, val in output.items(): + # flattn dp and micro batch + if isinstance(val, list): + output[key] = ( + Metric.aggregate_dp(val) + if isinstance(val[0], Metric) + else list(chain.from_iterable(val)) + ) + append_to_dict(metrics, output) + + output = tu.get_tensordict(tensor_dict={}, non_tensor_dict={"metrics": metrics}).cpu() + else: + output = None + return output + + @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="train"), blocking=False) + @DistProfiler.annotate(color="red", role="train_batch") + def train_batch(self, data: TensorDict) -> TensorDict: + assert self.loss_fn is not None, "loss function can't be None when calling train_batch" + assert not self.engine_config.forward_only, "Can't run `train_batch` when forward_only is in the engine config." + # global_token_num should be a list of number of tokens of each seq in this batch + global_token_num = tu.get(data, key="global_token_num") + disable_auto_offload = tu.get(data, key="disable_auto_offload", default=False) + images_seqlens = tu.get(data, key="images_seqlens", default=None) + + # inject engineering parameters if not specified + default_keys = dict( + use_remove_padding=self.model_config.get("use_remove_padding", False), + use_dynamic_bsz=self.engine_config.use_dynamic_bsz, + max_token_len_per_gpu=self.engine_config.max_token_len_per_gpu, + micro_batch_size_per_gpu=self.engine_config.micro_batch_size_per_gpu, + use_fused_kernels=self.engine_config.use_fused_kernels, + ) + + for key, val in default_keys.items(): + if key not in data.keys(): + tu.assign_non_tensor(data, **{key: val}) + + with ( + self.engine.train_mode(disable_auto_offload=disable_auto_offload), + Timer(name="train_batch", logger=None) as timer, + ): + output = self.engine.train_batch(data, loss_function=self.loss_fn) + # containing loss, model_output and metrics + # for training, we only care about loss and metrics + delta_time = timer.last + + update_lr_scheduler = tu.get(data, key="update_lr_scheduler", default=False) + # update lr scheduler + if update_lr_scheduler: + lr = self.engine.lr_scheduler_step() + else: + lr = None + + if self.engine.is_mp_src_rank_with_outputs(): + # we don't need model_output in training. Maybe we change out mind later + output.pop("model_output") + if lr is not None: + output["metrics"]["lr"] = lr + final_output = self._postprocess_output( + output, + global_token_num=global_token_num, + delta_time=delta_time, + forward_only=False, + images_seqlens=images_seqlens, + ).cpu() + else: + final_output = None + + return final_output + + @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="train"), blocking=False) + def infer_batch(self, data: TensorDict) -> TensorDict: + # add mfu calculator + global_token_num = tu.get(data, key="global_token_num") + compute_loss = tu.get(data, key="compute_loss", default=True) + disable_auto_offload = tu.get(data, key="disable_auto_offload", default=False) + no_lora_adapter = tu.pop(data, key="no_lora_adapter", default=False) + images_seqlens = tu.get(data, key="images_seqlens", default=None) + + default_keys = dict( + use_remove_padding=self.model_config.get("use_remove_padding", False), + use_dynamic_bsz=self.engine_config.use_dynamic_bsz, + max_token_len_per_gpu=self.engine_config.infer_max_token_len_per_gpu, + micro_batch_size_per_gpu=self.engine_config.infer_micro_batch_size_per_gpu, + use_fused_kernels=self.engine_config.use_fused_kernels, + ) + + for key, val in default_keys.items(): + if key not in data.keys(): + tu.assign_non_tensor(data, **{key: val}) + + # for sft training, we need to compute loss in eval + loss_function = self.loss_fn if compute_loss else None + + with ( + self.engine.eval_mode(disable_auto_offload=disable_auto_offload), + Timer(name="eval_batch", logger=None) as timer, + ): + adapter_ctx = self.engine.disable_adapter() if no_lora_adapter else nullcontext() + with adapter_ctx: + output = self.engine.infer_batch(data, loss_function=loss_function) + delta_time = timer.last + + if self.engine.is_mp_src_rank_with_outputs(): + final_output = self._postprocess_output( + output, + global_token_num=global_token_num, + delta_time=delta_time, + forward_only=True, + images_seqlens=images_seqlens, + ).cpu() + else: + final_output = None + + return final_output + + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def save_checkpoint(self, local_path, hdfs_path=None, global_step=0, max_ckpt_to_keep=None): + return self.engine.save_checkpoint(local_path, hdfs_path, global_step, max_ckpt_to_keep) + + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def load_checkpoint(self, local_path, hdfs_path=None, del_local_after_load=False): + return self.engine.load_checkpoint(local_path, hdfs_path, del_local_after_load) + + +class ActorRolloutRefWorker(Worker, DistProfilerExtension): + """Hybrid worker that includes actor model, rollout and optional ref model. + For standalone actor or rollout, use ActorWorker or BaseRollout respectively. + + NOTE: ActorRolloutRefWorker no longer support spmd mode and run native server mode. + """ + + def __init__( + self, config: DictConfig, role: str, distillation_config: Optional[DistillationConfig] = None, **kwargs + ): + Worker.__init__(self) + self.config = config + self.distillation_config = distillation_config + self.distillation_enabled = is_distillation_enabled(distillation_config) + self.role = role + self.actor: TrainingWorker = None + self.ref: TrainingWorker = None + self.rollout: BaseRollout = None + assert self.role in ["actor", "rollout", "ref", "actor_rollout", "actor_rollout_ref"] + self._is_actor = self.role in ["actor", "actor_rollout", "actor_rollout_ref"] + self._is_rollout = self.role in ["rollout", "actor_rollout", "actor_rollout_ref"] + self._is_ref = self.role in ["ref", "actor_rollout_ref"] + + if self._is_actor: + omega_profiler_config = config.actor.get("profiler", {}) + elif self._is_rollout: + # NOTE: In colocation mode, rollout config may not take effect (follow the actor config) + # This is for extendability in AsyncRL cases + omega_profiler_config = config.rollout.get("profiler", {}) + else: + omega_profiler_config = config.ref.get("profiler", {}) + + profiler_config = omega_conf_to_dataclass(omega_profiler_config, dataclass_type=ProfilerConfig) + if omega_profiler_config.get("tool", None) in ["npu", "nsys", "torch", "torch_memory", "precision_debugger"]: + tool_config = omega_conf_to_dataclass( + omega_profiler_config.get("tool_config", {}).get(omega_profiler_config.get("tool")) + ) + else: + tool_config = None + + self.enable_routing_replay = ( + self.config.actor.strategy == "megatron" and self.config.actor.megatron.router_replay.mode != "disabled" + ) + + DistProfilerExtension.__init__( + self, DistProfiler(rank=self.rank, config=profiler_config, tool_config=tool_config) + ) + + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def set_loss_fn(self, loss_fn): + self.actor.set_loss_fn(loss_fn=loss_fn) + + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def to(self, device, model=True, optimizer=True, grad=True): + """Manual control of load/offload""" + self.actor.to(device=device, model=model, optimizer=optimizer, grad=grad) + + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def init_model(self): + model_config: HFModelConfig = omega_conf_to_dataclass(self.config.model) + + # 1. build reference model + if "ref" in self.role: + # TODO: align ref config with actor config + with open_dict(self.config.ref): + self.config.ref.ppo_mini_batch_size = self.config.actor.ppo_mini_batch_size + self.config.ref.ppo_micro_batch_size = self.config.ref.pop("log_prob_micro_batch_size", None) + self.config.ref.ppo_micro_batch_size_per_gpu = self.config.ref.pop( + "log_prob_micro_batch_size_per_gpu", None + ) + self.config.ref.use_dynamic_bsz = self.config.ref.pop("log_prob_use_dynamic_bsz", False) + self.config.ref.ppo_max_token_len_per_gpu = self.config.ref.pop("log_prob_max_token_len_per_gpu", None) + ref_config: ActorConfig = omega_conf_to_dataclass(self.config.ref) + + # The ref model does not need to enable MTP; force it to false. + ref_config.model_config = deepcopy(model_config) + ref_config.model_config.mtp = MtpConfig(enable=False) + + # construct TrainingWorkerConfig + ref_training_config = TrainingWorkerConfig( + model_type=ref_config.model_config.get("model_type", "language_model"), + model_config=ref_config.model_config, + engine_config=ref_config.engine, + optimizer_config=ref_config.optim, + checkpoint_config=ref_config.checkpoint, + ) + + # assign engine configs + ref_training_config.engine_config.use_dynamic_bsz = self.config.ref.use_dynamic_bsz + ref_training_config.engine_config.infer_max_token_len_per_gpu = self.config.ref.ppo_max_token_len_per_gpu + ref_training_config.engine_config.infer_micro_batch_size_per_gpu = ( + self.config.ref.ppo_micro_batch_size_per_gpu + ) + ref_training_config.engine_config.use_remove_padding = model_config.get("use_remove_padding", False) + + self.ref = TrainingWorker(config=ref_training_config) + self.ref.reset() + self.set_dispatch_collect(mesh_name="ref", **self.ref.get_dispatch_collect()) + + # 2. build actor model + if "actor" in self.role: + actor_config: ActorConfig = omega_conf_to_dataclass(self.config.actor) + actor_config.model_config = model_config + distillation_config: Optional[DistillationConfig] = ( + omega_conf_to_dataclass(self.distillation_config) if self.distillation_enabled else None + ) + + actor_training_config = TrainingWorkerConfig( + model_type=actor_config.model_config.get("model_type", "language_model"), + model_config=actor_config.model_config, + engine_config=actor_config.engine, + optimizer_config=actor_config.optim, + checkpoint_config=actor_config.checkpoint, + ) + + assert self.config.actor.use_dynamic_bsz == self.config.rollout.log_prob_use_dynamic_bsz + + # assign engine configs + actor_training_config.engine_config.use_dynamic_bsz = self.config.actor.use_dynamic_bsz + actor_training_config.engine_config.infer_max_token_len_per_gpu = ( + self.config.rollout.log_prob_max_token_len_per_gpu + ) + actor_training_config.engine_config.infer_micro_batch_size_per_gpu = ( + self.config.rollout.log_prob_micro_batch_size_per_gpu + ) + actor_training_config.engine_config.max_token_len_per_gpu = self.config.actor.ppo_max_token_len_per_gpu + actor_training_config.engine_config.micro_batch_size_per_gpu = ( + self.config.actor.ppo_micro_batch_size_per_gpu + ) + actor_training_config.engine_config.use_remove_padding = model_config.get("use_remove_padding", False) + + if self.config.actor.use_dynamic_bsz: + assert self.config.rollout.log_prob_max_token_len_per_gpu is not None + assert self.config.actor.ppo_max_token_len_per_gpu is not None + else: + assert self.config.rollout.log_prob_micro_batch_size_per_gpu is not None + assert self.config.actor.ppo_micro_batch_size_per_gpu is not None + if self.distillation_enabled: + self.loss_fn = partial( + distillation_ppo_loss, config=actor_config, distillation_config=distillation_config + ) + else: + self.loss_fn = partial(ppo_loss, config=actor_config) + self.actor = TrainingWorker(config=actor_training_config) + self.actor.reset() + self.actor.set_loss_fn(self.loss_fn) + self.set_dispatch_collect(mesh_name="actor", **self.actor.get_dispatch_collect()) + + # 3. build rollout engine + if "rollout" in self.role: + rollout_config: RolloutConfig = omega_conf_to_dataclass(self.config.rollout) + + # TODO: move rollout_device_mesh into ServerAdapter + # 3.1 build rollout device mesh (sglang need only) + infer_tp = rollout_config.tensor_model_parallel_size * rollout_config.data_parallel_size + infer_pp = rollout_config.pipeline_model_parallel_size + infer_world_size = infer_tp * infer_pp + dp = self.world_size // infer_world_size + assert self.world_size % infer_world_size == 0, ( + f"rollout world_size: {self.world_size} is not divisible by infer_world_size: {infer_world_size}" + ) + rollout_device_mesh = init_device_mesh( + get_device_name(), mesh_shape=(dp, infer_tp, infer_pp), mesh_dim_names=["dp", "infer_tp", "infer_pp"] + ) + + # 3.2 initialize rollout engine + rollout_cls: type[BaseRollout] = get_rollout_class(rollout_config.name, rollout_config.mode) + self.rollout = rollout_cls( + config=rollout_config, model_config=model_config, device_mesh=rollout_device_mesh + ) + + # used for LoRA (base_sync_done is unused in merge-only mode but kept for Phase 2 adapter path) + self.base_sync_done: bool = "dummy" not in self.config.rollout.load_format + self.layered_summon = self.config.rollout.get("layered_summon", False) + self.peft_merge: bool = model_config.lora.get("merge", False) + + # 4. build checkpoint engine + if "actor" in self.role: + checkpoint_engine_config = omega_conf_to_dataclass(self.config.rollout.checkpoint_engine) + backend = checkpoint_engine_config.backend + bucket_size = checkpoint_engine_config.update_weights_bucket_megabytes << 20 + engine_kwargs = checkpoint_engine_config.engine_kwargs.get(backend, {}) + # If custom_backend_module is set, import it so plugins can register + # in CheckpointEngineRegistry before the backend is instantiated. + import_external_libs(checkpoint_engine_config.custom_backend_module or None) + self.checkpoint_engine = CheckpointEngineRegistry.new( + backend, is_master=(torch.distributed.get_rank() == 0), bucket_size=bucket_size, **engine_kwargs + ) + + # Free cached GPU memory so colocated vLLM processes can see it via cudaMemGetInfo + aggressive_empty_cache(force_sync=True) + + @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="ref")) + @DistProfiler.annotate(color="olive", role="ref_compute_log_prob") + @_with_routing_replay_flag(enabled=False) + def compute_ref_log_prob(self, data: TensorDict) -> TensorDict: + output = self.ref.infer_batch(data=data) + return output.cpu() if output is not None else None + + @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="actor")) + @DistProfiler.annotate(color="blue", role="actor_compute_log_prob") + @_with_routing_replay_flag(enabled=True) + def compute_log_prob(self, data: TensorDict) -> TensorDict: + output = self.actor.infer_batch(data) + + return output.cpu() if output is not None else None + + @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="actor")) + @DistProfiler.annotate(color="red", role="actor_update") + @_with_routing_replay_flag(enabled=True) + def update_actor(self, data: TensorDict) -> TensorDict: + output = self.actor.train_mini_batch(data=data) + return output.cpu() if output is not None else None + + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def load_checkpoint(self, local_path, hdfs_path=None, del_local_after_load=False): + assert "actor" in self.role, "load_checkpoint only support actor role" + self.actor.load_checkpoint(local_path, hdfs_path, del_local_after_load) + + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def save_checkpoint(self, local_path, hdfs_path=None, global_step=0, max_ckpt_to_keep=None): + assert "actor" in self.role, "save_checkpoint only support actor role" + self.actor.save_checkpoint(local_path, hdfs_path, global_step, max_ckpt_to_keep) + + @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False) + async def update_weights(self, global_steps: int = None, mode: str = "auto"): + """Update weights from trainer to rollout. + + 1. For sync training with colocated trainer and rollout, update rollout directly from model engine. + - before update_weights: rollout should be in sleep mode. + - after update_weights: rollout should be in wake_up mode. + 2. For async training with disaggregated trainer and rollout, send_weights only by checkpoint engine. + + LoRA handling: when model.lora.merge=True (peft_merge), LoRA is merged into + base weights before sync. The engine returns full HF-keyed params with + peft_config=None, so the rollout receives a standard weight update. + + Args: + global_steps: Current global training step count, passed to rollout for logging/tracking. + mode: Weight update strategy. Supported values: + - ``"auto"``: Automatically resolve to the backend configured in + ``config.rollout.checkpoint_engine.backend`` (default). + - ``"naive"``: Direct in-process weight sync between colocated trainer + and rollout. Used for synchronous training where both share the same + process. Rollout must be in sleep mode before this call. + - Any other value: Delegates to + :meth:`checkpoint_engine.send_weights` for asynchronous weight + transfer via checkpoint engine, suitable for disaggregated + trainer/rollout deployments. + """ + + # Resolve mode: "auto" falls back to config, explicit values take precedence + effective_mode = mode if mode != "auto" else self.config.rollout.checkpoint_engine.backend + + # 0. send_weights only for async training with disaggregated trainer and rollout + if effective_mode != "naive": + per_tensor_param, _ = self.actor.engine.get_per_tensor_param() + await self.checkpoint_engine.send_weights(per_tensor_param) + return + + set_expandable_segments(False) + log_gpu_memory_usage("Before resume weights", logger=logger) + + # 1. resume rollout memory (weights were released during sleep) + if self.config.rollout.free_cache_engine: + await self.rollout.resume(tags=["weights"]) + log_gpu_memory_usage("After resume weights", logger=logger) + + # 2. determine if we need a base weight sync (adapter path only) + per_tensor_param, peft_config = self.actor.engine.get_per_tensor_param( + layered_summon=self.layered_summon, base_sync_done=True + ) + + do_lora_base_sync = False + if not self.peft_merge and peft_config is not None: + self.rollout.sleep_level = 1 + do_lora_base_sync = not self.base_sync_done + + # 3. sync weights: For SGLang, we need base first (when needed), then adapter/merged + if do_lora_base_sync: + per_tensor_param_base, peft_config = self.actor.engine.get_per_tensor_param( + layered_summon=self.layered_summon, base_sync_done=False + ) + await self.rollout.update_weights( + per_tensor_param_base, peft_config=peft_config, base_sync_done=False, global_steps=global_steps + ) + + await self.rollout.update_weights( + per_tensor_param, peft_config=peft_config, base_sync_done=True, global_steps=global_steps + ) + + log_gpu_memory_usage("After update_weights", logger=logger) + + # 3. offload model to cpu + if self.actor.engine.is_param_offload_enabled: + self.actor.engine.to("cpu", model=True, optimizer=False, grad=False) + aggressive_empty_cache(force_sync=True) + + # 4. resume kv_cache + if self.config.rollout.free_cache_engine: + await self.rollout.resume(tags=["kv_cache"]) + log_gpu_memory_usage("After resume kv_cache", logger=logger) + + self.base_sync_done = True + set_expandable_segments(True) + + @register(dispatch_mode=Dispatch.DP_COMPUTE, blocking=False) + def execute_checkpoint_engine(self, method: str, *args, **kwargs): + """Execute checkpoint engine method. + + Args: + method (str): Checkpoint engine method name. + *args: Variable length argument list. + **kwargs: Arbitrary keyword arguments. + + """ + return getattr(self.checkpoint_engine, method)(*args, **kwargs) diff --git a/verl/verl/workers/reward_manager/__init__.py b/verl/verl/workers/reward_manager/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..06b693697f0383ff9bc1ace9a24a914b0fa1096c --- /dev/null +++ b/verl/verl/workers/reward_manager/__init__.py @@ -0,0 +1,37 @@ +# Copyright 2024 PRIME team and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .registry import get_reward_manager_cls, register # noqa: I001 +from .batch import BatchRewardManager +from .dapo import DAPORewardManager +from .naive import NaiveRewardManager +from .prime import PrimeRewardManager + +# Note(haibin.lin): no need to include all reward managers here in case of complicated dependencies +__all__ = [ + "BatchRewardManager", + "DAPORewardManager", + "NaiveRewardManager", + "PrimeRewardManager", + "register", + "get_reward_manager_cls", +] + +# Import experimental reward managers to ensure they are registered +try: + from verl.experimental.reward_loop.reward_manager.limited import RateLimitedRewardManager # noqa: F401 + + __all__.append("RateLimitedRewardManager") +except ImportError: + pass # Optional dependency, may not be available diff --git a/verl/verl/workers/reward_manager/abstract.py b/verl/verl/workers/reward_manager/abstract.py new file mode 100644 index 0000000000000000000000000000000000000000..494c87c5597e1e427712410e72bd040fd50204cf --- /dev/null +++ b/verl/verl/workers/reward_manager/abstract.py @@ -0,0 +1,72 @@ +# Copyright 2023-2025 SGLang Team +# Copyright Amazon.com, Inc. or its affiliates. +# Copyright 2025 ModelBest Inc. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from abc import ABC, abstractmethod +from typing import Any, Callable + +import torch + +from verl.protocol import DataProto + +RawRewardFn = Callable[..., Any] + + +class AbstractRewardManager(ABC): + @abstractmethod + def __init__( + self, + tokenizer: Any, + num_examine: int, + compute_score: RawRewardFn | None, + reward_fn_key: str = "data_source", + **kwargs: Any, + ): + pass + + @abstractmethod + def __call__( + self, + data: DataProto, + return_dict: bool = False, + ) -> torch.Tensor | dict[str, Any]: + pass + + def _extract_reward_from_rm_scores( + self, data: DataProto, return_dict: bool = False + ) -> torch.Tensor | dict[str, Any] | None: + """ + Extract reward from already-computed rm_scores if available. + This has been deprecated. + + Args: + data: DataProto object containing the batch data + return_dict: Whether to return a dictionary with reward_tensor and reward_extra_info + + Returns: + If rm_scores exists: + - If return_dict=True: dict with "reward_tensor" and "reward_extra_info" + - If return_dict=False: torch.Tensor of rm_scores + If rm_scores doesn't exist: None + """ + if "rm_scores" not in data.batch.keys(): + return None + + if return_dict: + reward_extra_keys = data.meta_info.get("reward_extra_keys", []) + reward_extra_info = {key: data.non_tensor_batch[key] for key in reward_extra_keys} + return {"reward_tensor": data.batch["rm_scores"], "reward_extra_info": reward_extra_info} + else: + return data.batch["rm_scores"] diff --git a/verl/verl/workers/reward_manager/batch.py b/verl/verl/workers/reward_manager/batch.py new file mode 100644 index 0000000000000000000000000000000000000000..078e301d45f93c4e4b1092245bd6c9c55476b8e0 --- /dev/null +++ b/verl/verl/workers/reward_manager/batch.py @@ -0,0 +1,128 @@ +# Copyright 2025 Individual Contributor: Mert Unsal +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections import defaultdict +from typing import Any + +import torch + +from verl import DataProto +from verl.workers.reward_manager import register +from verl.workers.reward_manager.abstract import AbstractRewardManager, RawRewardFn + + +@register("batch") +class BatchRewardManager(AbstractRewardManager): + """ + A batch reward manager that computes rewards for a batch of data. + + Args: + tokenizer (Tokenizer): The tokenizer to use for decoding the responses. + num_examine (int): The number of responses to examine. + compute_score (callable): The function to compute the rewards. + reward_fn_key (str): The key to use for the reward function. + reward_kwargs (dict): The keyword arguments to pass to the reward function. + """ + + def __init__( + self, tokenizer, num_examine, compute_score: RawRewardFn, reward_fn_key="data_source", **reward_kwargs + ): + self.tokenizer = tokenizer + self.num_examine = num_examine + self.compute_score = compute_score + self.reward_fn_key = reward_fn_key + self.reward_kwargs = reward_kwargs + + def verify(self, data): + prompt_ids = data.batch["prompts"] + response_ids = data.batch["responses"] + attention_mask = data.batch["attention_mask"] + + prompt_len = prompt_ids.shape[-1] + valid_response_lengths = attention_mask[:, prompt_len:].sum(dim=-1) + + responses_str = [] + for i in range(len(data)): + valid_len = valid_response_lengths[i] + valid_response_ids = response_ids[i][:valid_len] + response_str = self.tokenizer.decode(valid_response_ids, skip_special_tokens=True) + responses_str.append(response_str) + + ground_truths = [item.non_tensor_batch["reward_model"].get("ground_truth", None) for item in data] + data_sources = data.non_tensor_batch[self.reward_fn_key] + rollout_reward_scores = data.non_tensor_batch.get("reward_scores", [{} for _ in range(len(data))]) + extras = data.non_tensor_batch.get("extra_info", [{} for _ in range(len(data))]) + + for i in range(len(data)): + extras[i]["rollout_reward_scores"] = rollout_reward_scores[i] + + scores = self.compute_score( + data_sources=data_sources, + solution_strs=responses_str, + ground_truths=ground_truths, + extra_infos=extras, + **self.reward_kwargs, + ) + + return scores + + def __call__(self, data: DataProto, return_dict: bool = False) -> torch.Tensor | dict[str, Any]: + # If there is rm score, we directly return rm score. Otherwise, we compute via rm_score_fn + reward_from_rm_scores = self._extract_reward_from_rm_scores(data, return_dict) + if reward_from_rm_scores is not None: + return reward_from_rm_scores + + reward_tensor = torch.zeros_like(data.batch["responses"], dtype=torch.float32) + reward_extra_info = defaultdict(list) + prompt_ids = data.batch["prompts"] + prompt_len = prompt_ids.shape[-1] + attention_mask = data.batch["attention_mask"] + valid_response_lengths = attention_mask[:, prompt_len:].sum(dim=-1) + data_sources = data.non_tensor_batch[self.reward_fn_key] + + scores = self.verify(data) + rewards = [] + already_printed: dict[str, Any] = {} + + for i in range(len(data)): + length = valid_response_lengths[i].item() + score = scores[i] + + if isinstance(score, dict): + reward = score["score"] + for key, value in score.items(): + reward_extra_info[key].append(value) + else: + reward = score + + rewards.append(reward) + reward_tensor[i, length - 1] = reward + + data_source = data_sources[i] + if already_printed.get(data_source, 0) < self.num_examine: + response_str = self.tokenizer.decode(data.batch["responses"][i][:length], skip_special_tokens=True) + prompt_str = self.tokenizer.decode(data.batch["prompts"][i], skip_special_tokens=True) + ground_truth = data[i].non_tensor_batch["reward_model"].get("ground_truth", None) + print("[prompt]", prompt_str) + print("[response]", response_str) + print("[ground_truth]", ground_truth) + print("[score]", scores[i]) + already_printed[data_source] = already_printed.get(data_source, 0) + 1 + + data.batch["acc"] = torch.tensor(rewards, dtype=torch.float32, device=prompt_ids.device) + + if return_dict: + return {"reward_tensor": reward_tensor, "reward_extra_info": reward_extra_info} + else: + return reward_tensor diff --git a/verl/verl/workers/reward_manager/dapo.py b/verl/verl/workers/reward_manager/dapo.py new file mode 100644 index 0000000000000000000000000000000000000000..f44cb3d8f9d78ee2ebe3861d4fdb50c4c3e6aec9 --- /dev/null +++ b/verl/verl/workers/reward_manager/dapo.py @@ -0,0 +1,154 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections import defaultdict + +import torch + +from verl import DataProto +from verl.utils.reward_score import default_compute_score +from verl.workers.reward_manager import register +from verl.workers.reward_manager.abstract import AbstractRewardManager + + +@register("dapo") +class DAPORewardManager(AbstractRewardManager): + """The reward manager.""" + + def __init__( + self, + tokenizer, + num_examine, + compute_score=None, + reward_fn_key="data_source", + max_resp_len=None, + overlong_buffer_cfg=None, + ) -> None: + self.tokenizer = tokenizer + self.num_examine = num_examine # the number of batches of decoded responses to print to the console + self.compute_score = compute_score or default_compute_score + self.reward_fn_key = reward_fn_key + self.overlong_buffer_cfg = overlong_buffer_cfg + self.max_resp_len = max_resp_len + + if self.overlong_buffer_cfg is not None: + assert self.max_resp_len is not None, ( + f"max_resp_len must be provided if {overlong_buffer_cfg=}, but got None" + ) + assert self.max_resp_len >= self.overlong_buffer_cfg.len, ( + "max_resp_len must be larger than overlong_buffer.len" + ) + assert not self.overlong_buffer_cfg.enable or self.overlong_buffer_cfg.len > 0, ( + "overlong_buffer.len must be positive when overlong penalty is enabled," + f"but got {self.overlong_buffer_cfg.len}." + "To disable the overlong penalty, set overlong_buffer.enable = False" + ) + + def __call__(self, data: DataProto, return_dict: bool = False): + """We will expand this function gradually based on the available datasets""" + + # If there is rm score, we directly return rm score. Otherwise, we compute via rm_score_fn + reward_from_rm_scores = self._extract_reward_from_rm_scores(data, return_dict) + if reward_from_rm_scores is not None: + return reward_from_rm_scores + + reward_tensor = torch.zeros_like(data.batch["responses"], dtype=torch.float32) + reward_extra_info = defaultdict(list) + + already_print_data_sources = {} + + for i in range(len(data)): + data_item = data[i] # DataProtoItem + + prompt_ids = data_item.batch["prompts"] + + prompt_length = prompt_ids.shape[-1] + + valid_prompt_length = data_item.batch["attention_mask"][:prompt_length].sum() + valid_prompt_ids = prompt_ids[-valid_prompt_length:] + + response_ids = data_item.batch["responses"] + valid_response_length = data_item.batch["attention_mask"][prompt_length:].sum() + valid_response_ids = response_ids[:valid_response_length] + + # decode + prompt_str = self.tokenizer.decode(valid_prompt_ids, skip_special_tokens=True) + response_str = self.tokenizer.decode(valid_response_ids, skip_special_tokens=True) + eos_token = self.tokenizer.eos_token + if response_str.endswith(eos_token): + response_str = response_str[: -len(eos_token)] + + ground_truth = data_item.non_tensor_batch["reward_model"]["ground_truth"] + + data_source = data_item.non_tensor_batch[self.reward_fn_key] + + extra_info = data_item.non_tensor_batch.get("extra_info", {}) + + rollout_reward_scores = data_item.non_tensor_batch.get("reward_scores", {}) + + extra_info["rollout_reward_scores"] = rollout_reward_scores + + result = self.compute_score( + data_source=data_source, + solution_str=response_str, + ground_truth=ground_truth, + extra_info=extra_info, + ) + + score: float + if isinstance(result, dict): + score = result["score"] + # Store the information including original reward + for key, value in result.items(): + reward_extra_info[key].append(value) + else: + score = result + reward_extra_info["acc"].append(score) + + reward = score + + if self.overlong_buffer_cfg.enable: + overlong_buffer_len = self.overlong_buffer_cfg.len + expected_len = self.max_resp_len - overlong_buffer_len + exceed_len = valid_response_length - expected_len + overlong_penalty_factor = self.overlong_buffer_cfg.penalty_factor + overlong_reward = min(-exceed_len / overlong_buffer_len * overlong_penalty_factor, 0) + reward += overlong_reward + if self.overlong_buffer_cfg.log: + reward_extra_info["overlong_reward"].append(overlong_reward) + reward_extra_info["overlong"].append(overlong_reward < 0) + + reward_tensor[i, valid_response_length - 1] = reward + + if data_source not in already_print_data_sources: + already_print_data_sources[data_source] = 0 + + if already_print_data_sources[data_source] < self.num_examine: + already_print_data_sources[data_source] += 1 + print("[prompt]", prompt_str) + print("[response]", response_str) + print("[ground_truth]", ground_truth) + if isinstance(result, dict): + for key, value in result.items(): + print(f"[{key}]", value) + else: + print("[score]", score) + + if return_dict: + return { + "reward_tensor": reward_tensor, + "reward_extra_info": reward_extra_info, + } + else: + return reward_tensor diff --git a/verl/verl/workers/reward_manager/naive.py b/verl/verl/workers/reward_manager/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..f3ca122c2b6f3da12bda2ba08412529e78d21907 --- /dev/null +++ b/verl/verl/workers/reward_manager/naive.py @@ -0,0 +1,122 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections import defaultdict +from typing import Any + +import torch + +from verl import DataProto +from verl.utils.reward_score import default_compute_score +from verl.workers.reward_manager import register +from verl.workers.reward_manager.abstract import AbstractRewardManager + + +@register("naive") +class NaiveRewardManager(AbstractRewardManager): + """The reward manager.""" + + def __init__(self, tokenizer, num_examine, compute_score=None, reward_fn_key="data_source") -> None: + """ + Initialize the NaiveRewardManager instance. + + Args: + tokenizer: The tokenizer used to decode token IDs into text. + num_examine: The number of batches of decoded responses to print to the console for debugging purpose. + compute_score: A function to compute the reward score. If None, `default_compute_score` will be used. + reward_fn_key: The key used to access the data source in the non-tensor batch data. Defaults to + "data_source". + """ + self.tokenizer = tokenizer # Store the tokenizer for decoding token IDs + self.num_examine = num_examine # the number of batches of decoded responses to print to the console + self.compute_score = compute_score or default_compute_score + self.reward_fn_key = reward_fn_key # Store the key for accessing the data source + + def __call__(self, data: DataProto, return_dict: bool = False) -> torch.Tensor | dict[str, Any]: + """We will expand this function gradually based on the available datasets""" + + # If there is rm score, we directly return rm score. Otherwise, we compute via rm_score_fn + reward_from_rm_scores = self._extract_reward_from_rm_scores(data, return_dict) + if reward_from_rm_scores is not None: + return reward_from_rm_scores + + reward_tensor = torch.zeros_like(data.batch["responses"], dtype=torch.float32) + reward_extra_info = defaultdict(list) + + already_print_data_sources = {} + + for i in range(len(data)): + data_item = data[i] # DataProtoItem + + prompt_ids = data_item.batch["prompts"] + + prompt_length = prompt_ids.shape[-1] + + valid_prompt_length = data_item.batch["attention_mask"][:prompt_length].sum() + valid_prompt_ids = prompt_ids[-valid_prompt_length:] + + response_ids = data_item.batch["responses"] + valid_response_length = data_item.batch["attention_mask"][prompt_length:].sum() + valid_response_ids = response_ids[:valid_response_length] + + # decode + prompt_str = self.tokenizer.decode(valid_prompt_ids, skip_special_tokens=True) + response_str = self.tokenizer.decode(valid_response_ids, skip_special_tokens=True) + + ground_truth = data_item.non_tensor_batch["reward_model"]["ground_truth"] + data_source = data_item.non_tensor_batch[self.reward_fn_key] + extra_info = data_item.non_tensor_batch.get("extra_info", {}) + num_turns = data_item.non_tensor_batch.get("__num_turns__", None) + rollout_reward_scores = data_item.non_tensor_batch.get("reward_scores", {}) + extra_info["num_turns"] = num_turns + extra_info["rollout_reward_scores"] = rollout_reward_scores + + score = self.compute_score( + data_source=data_source, + solution_str=response_str, + ground_truth=ground_truth, + extra_info=extra_info, + ) + + if isinstance(score, dict): + reward = score["score"] + # Store the information including original reward + for key, value in score.items(): + reward_extra_info[key].append(value) + else: + reward = score + + reward_tensor[i, valid_response_length - 1] = reward + + if data_source not in already_print_data_sources: + already_print_data_sources[data_source] = 0 + + if already_print_data_sources[data_source] < self.num_examine: + already_print_data_sources[data_source] += 1 + print("[prompt]", prompt_str) + print("[response]", response_str) + print("[ground_truth]", ground_truth) + if isinstance(score, dict): + for key, value in score.items(): + print(f"[{key}]", value) + else: + print("[score]", score) + + if return_dict: + return { + "reward_tensor": reward_tensor, + "reward_extra_info": reward_extra_info, + } + else: + return reward_tensor diff --git a/verl/verl/workers/reward_manager/prime.py b/verl/verl/workers/reward_manager/prime.py new file mode 100644 index 0000000000000000000000000000000000000000..b15ed7c3fcb3931e8860b6fe32898ebb2cc5386c --- /dev/null +++ b/verl/verl/workers/reward_manager/prime.py @@ -0,0 +1,189 @@ +# Copyright 2024 PRIME team and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +from concurrent.futures import ProcessPoolExecutor +from functools import partial +from typing import Any, Callable, Optional + +import psutil +import torch +from transformers import PreTrainedTokenizer + +from verl import DataProto +from verl.utils.ray_utils import get_event_loop +from verl.utils.reward_score import default_compute_score +from verl.workers.reward_manager import register +from verl.workers.reward_manager.abstract import AbstractRewardManager + + +async def single_compute_score(evaluation_func, completion, reference, task, task_extra_info, executor, timeout=300.0): + loop = get_event_loop() + try: + # Ensure process_completion is called properly + future = loop.run_in_executor(executor, partial(evaluation_func, task, completion, reference, task_extra_info)) + return await asyncio.wait_for(future, timeout=timeout) + except asyncio.TimeoutError: + print(f"[Timeout] Task timeout: {completion}") + return None # Default value for timed-out rows + except Exception as e: + print(f"[Error] Task failed: {e}, completion: {completion[:80]}") + return None # Default value for failed rows + + +async def parallel_compute_score_async( + evaluation_func, completions, references, tasks, extra_info=None, num_processes=64 +): + if extra_info is None: + extra_info = [None] * len(tasks) + scores = [] + with ProcessPoolExecutor(max_workers=num_processes) as executor: + # to prevent very occasional starvation caused by some anomalous programs ( like infinite loop ), the + # exceptions in async programs will instantly halt the evaluation, and all summoned processes will be killed. + try: + # Create tasks for all rows + tasks_async = [ + single_compute_score(evaluation_func, c, r, t, ei, executor, timeout=300.0) + for c, r, t, ei in zip(completions, references, tasks, extra_info, strict=True) + ] + results = await asyncio.gather(*tasks_async, return_exceptions=False) + except Exception as e: + print(f"[Exception] async gather failed: {e}") + raise + finally: + terminated_count = 0 + for pid, proc in executor._processes.items(): + try: + p = psutil.Process(pid) + p.terminate() + try: + p.wait(timeout=5) + except psutil.TimeoutExpired: + p.kill() + terminated_count += 1 + except Exception: + pass + print(f"[Shutdown] {terminated_count} subprocess(es) terminated.") + + # Process results + for result, completion, reference, task in zip(results, completions, references, tasks, strict=True): + if isinstance(result, Exception) or result is None: + # Handle failed or timed-out tasks + scores.append(0.0) + elif isinstance(result, int | float | bool): + scores.append(float(result)) + else: + scores.append(float(result[0])) + return scores + + +def run_reward_scoring(evaluation_func, completions, references, tasks, extra_info=None, num_processes=64): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + return loop.run_until_complete( + parallel_compute_score_async(evaluation_func, completions, references, tasks, extra_info, num_processes) + ) + finally: + loop.close() + + +@register("prime") +class PrimeRewardManager(AbstractRewardManager): + """ + The Reward Manager used in https://github.com/PRIME-RL/PRIME + """ + + def __init__( + self, + tokenizer: PreTrainedTokenizer, + num_examine: int, + compute_score: Optional[Callable] = None, + reward_fn_key: str = "data_source", + ) -> None: + self.tokenizer = tokenizer + self.num_examine = num_examine # the number of batches of decoded responses to print to the console + self.compute_score = compute_score or default_compute_score + self.reward_fn_key = reward_fn_key + + def verify(self, data): + """ + verify the batch and save as ``acc`` tensor + """ + # batched scoring + prompt_ids = data.batch["prompts"] + + response_ids = data.batch["responses"] + sequences_str = self.tokenizer.batch_decode(response_ids, skip_special_tokens=True) + ground_truth = [data_item.non_tensor_batch["reward_model"]["ground_truth"] for data_item in data] + data_sources = data.non_tensor_batch[self.reward_fn_key] + extra_info = data.non_tensor_batch.get("extra_info", None) + + assert len(sequences_str) == len(ground_truth) == len(data_sources) + try: + scores = run_reward_scoring( + self.compute_score, + completions=sequences_str, + references=ground_truth, + tasks=data_sources, + extra_info=extra_info, + num_processes=64, + ) + except asyncio.TimeoutError: + print("[Timeout] Global reward scoring timed out. Setting all as 0.") + scores = [0.0 for _ in range(len(sequences_str))] + except Exception as e: + print(f"[Error] Unexpected error during scoring. Setting all as 0. {e}") + scores = [0.0 for _ in range(len(sequences_str))] + data.batch["acc"] = torch.tensor(scores, dtype=torch.float32, device=prompt_ids.device) + return scores + + def __call__(self, data: DataProto, return_dict: bool = False) -> torch.Tensor | dict[str, Any]: + """We will expand this function gradually based on the available datasets""" + + # If there is rm score, we directly return rm score. Otherwise, we compute via rm_score_fn + reward_from_rm_scores = self._extract_reward_from_rm_scores(data, return_dict) + if reward_from_rm_scores is not None: + return reward_from_rm_scores + + reward_tensor = torch.zeros_like(data.batch["responses"], dtype=torch.float32) + + already_print_data_sources = {} + + # batched scoring + prompt_ids = data.batch["prompts"] + prompt_length = prompt_ids.shape[-1] + + response_ids = data.batch["responses"] + valid_response_length = data.batch["attention_mask"][:, prompt_length:].sum(dim=-1) + sequences_str = self.tokenizer.batch_decode(response_ids, skip_special_tokens=True) + data_sources = data.non_tensor_batch["data_source"] + + scores = self.verify(data) + + for i in range(len(data)): + data_source = data_sources[i] + reward_tensor[i, valid_response_length[i].item() - 1] = scores[i] + + if data_source not in already_print_data_sources: + already_print_data_sources[data_source] = 0 + + if already_print_data_sources[data_source] < self.num_examine: + already_print_data_sources[data_source] += 1 + print(sequences_str) + + if return_dict: + return {"reward_tensor": reward_tensor} + else: + return reward_tensor diff --git a/verl/verl/workers/reward_manager/registry.py b/verl/verl/workers/reward_manager/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..4e255d8ac8cdc9467f33ba4a63c2f5ff27f44d33 --- /dev/null +++ b/verl/verl/workers/reward_manager/registry.py @@ -0,0 +1,55 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Callable + +from verl.workers.reward_manager.abstract import AbstractRewardManager + +__all__ = ["register", "get_reward_manager_cls"] + +REWARD_MANAGER_REGISTRY: dict[str, type[AbstractRewardManager]] = {} + + +def register(name: str) -> Callable[[type[AbstractRewardManager]], type[AbstractRewardManager]]: + """Decorator to register a reward manager class with a given name. + + Args: + name: `(str)` + The name of the reward manager. + """ + + def decorator(cls: type[AbstractRewardManager]) -> type[AbstractRewardManager]: + if name in REWARD_MANAGER_REGISTRY and REWARD_MANAGER_REGISTRY[name] != cls: + raise ValueError( + f"Reward manager {name} has already been registered: {REWARD_MANAGER_REGISTRY[name]} vs {cls}" + ) + REWARD_MANAGER_REGISTRY[name] = cls + return cls + + return decorator + + +def get_reward_manager_cls(name: str) -> type[AbstractRewardManager]: + """Get the reward manager class with a given name. + + Args: + name: `(str)` + The name of the reward manager. + + Returns: + `(type)`: The reward manager class. + """ + if name not in REWARD_MANAGER_REGISTRY: + raise ValueError(f"Unknown reward manager: {name}") + return REWARD_MANAGER_REGISTRY[name] diff --git a/verl/verl/workers/rollout/__init__.py b/verl/verl/workers/rollout/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f6bd6c28b770fd5996bd23936796ef374ccb8ec1 --- /dev/null +++ b/verl/verl/workers/rollout/__init__.py @@ -0,0 +1,20 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .base import BaseRollout, get_rollout_class +from .hf_rollout import HFRollout +from .naive import NaiveRollout +from .replica import RolloutReplica + +__all__ = ["BaseRollout", "NaiveRollout", "HFRollout", "get_rollout_class", "RolloutReplica"] diff --git a/verl/verl/workers/rollout/__pycache__/__init__.cpython-312.pyc b/verl/verl/workers/rollout/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..40df08b666099c318e81357873c6197942463d79 Binary files /dev/null and b/verl/verl/workers/rollout/__pycache__/__init__.cpython-312.pyc differ diff --git a/verl/verl/workers/rollout/__pycache__/base.cpython-312.pyc b/verl/verl/workers/rollout/__pycache__/base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..30d26e2e2388c274a5c40647ed3870f6b79d334b Binary files /dev/null and b/verl/verl/workers/rollout/__pycache__/base.cpython-312.pyc differ diff --git a/verl/verl/workers/rollout/__pycache__/hf_rollout.cpython-312.pyc b/verl/verl/workers/rollout/__pycache__/hf_rollout.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a72ad741f90d47720989e7a55728fdb60e73d348 Binary files /dev/null and b/verl/verl/workers/rollout/__pycache__/hf_rollout.cpython-312.pyc differ diff --git a/verl/verl/workers/rollout/__pycache__/replica.cpython-312.pyc b/verl/verl/workers/rollout/__pycache__/replica.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6ed69c85c1f6314a16f7e567649e778acbbd6f14 Binary files /dev/null and b/verl/verl/workers/rollout/__pycache__/replica.cpython-312.pyc differ diff --git a/verl/verl/workers/rollout/__pycache__/utils.cpython-312.pyc b/verl/verl/workers/rollout/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..556558e5483ba8c84444f873aa9c9567a0aed731 Binary files /dev/null and b/verl/verl/workers/rollout/__pycache__/utils.cpython-312.pyc differ diff --git a/verl/verl/workers/rollout/base.py b/verl/verl/workers/rollout/base.py new file mode 100644 index 0000000000000000000000000000000000000000..033e7684c1dd5348f90d131d68e8e17790cb40c2 --- /dev/null +++ b/verl/verl/workers/rollout/base.py @@ -0,0 +1,104 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib +from abc import ABC, abstractmethod +from typing import Generator + +import torch +from torch.distributed.device_mesh import DeviceMesh + +from verl import DataProto +from verl.utils.config import omega_conf_to_dataclass +from verl.workers.config import HFModelConfig, RolloutConfig + +__all__ = ["BaseRollout"] + + +class BaseRollout(ABC): + """Base class for rollout.""" + + def __init__( + self, + config: RolloutConfig, + model_config: HFModelConfig, + device_mesh: DeviceMesh, + *args, + **kwargs, + ): + self.config = omega_conf_to_dataclass(config) + self.model_config: HFModelConfig = omega_conf_to_dataclass(model_config) + self.device_mesh = device_mesh + + @abstractmethod + async def resume(self, tags: list[str]): + """Resume rollout weights or kv cache in GPU memory. + + Args: + tags: weights or kv_cache. + """ + pass + + @abstractmethod + async def update_weights( + self, + weights: Generator[tuple[str, torch.Tensor], None, None], + **kwargs, + ): + """Update the weights of the rollout model. + + Args: + weights: A generator that yields the name of the weight tensor and the tensor itself. + """ + pass + + @abstractmethod + async def release(self): + """Release weights and kv cache in GPU memory.""" + pass + + def generate_sequences(self, prompts: DataProto) -> DataProto: + """Batch generate sequences in sync mode. + + Args: + prompts: The input prompts. + + Returns: + The output sequences. + """ + raise NotImplementedError + + +_ROLLOUT_REGISTRY = { + ("vllm", "async"): "verl.workers.rollout.vllm_rollout.ServerAdapter", + ("sglang", "async"): "verl.workers.rollout.sglang_rollout.sglang_rollout.ServerAdapter", + ("trtllm", "async"): "verl.workers.rollout.trtllm_rollout.trtllm_rollout.ServerAdapter", +} + + +def get_rollout_class(rollout_name: str, mode: str = "async") -> type[BaseRollout]: + """Get the rollout class by name. + + Args: + rollout_name: The name of the rollout. + mode: The mode of the rollout, async: server mode. + + Returns: + The rollout class. + """ + assert (rollout_name, mode) in _ROLLOUT_REGISTRY, f"Rollout {rollout_name} with mode {mode} not found" + fqdn = _ROLLOUT_REGISTRY[(rollout_name, mode)] + module_name, class_name = fqdn.rsplit(".", 1) + rollout_module = importlib.import_module(module_name) + return getattr(rollout_module, class_name) diff --git a/verl/verl/workers/rollout/hf_rollout.py b/verl/verl/workers/rollout/hf_rollout.py new file mode 100644 index 0000000000000000000000000000000000000000..b60f275e2051523aee7805ecab3af19d1ebf6dc5 --- /dev/null +++ b/verl/verl/workers/rollout/hf_rollout.py @@ -0,0 +1,177 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Rollout with huggingface models. +TODO: refactor this class. Currently, it will hang when using FSDP HybridShard. We should actually create a single +GPU model. Then, get full state_dict and bind the state_dict to the single GPU model. Then, use the single GPU model +to perform generation. +""" + +import contextlib + +import torch +import torch.distributed +from tensordict import TensorDict +from torch import nn +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +from transformers import GenerationConfig + +from verl import DataProto +from verl.utils.device import get_device_name, get_torch_device +from verl.utils.torch_functional import get_response_mask + +from .base import BaseRollout + +__all__ = ["HFRollout"] + + +class HFRollout(BaseRollout): + def __init__(self, module: nn.Module, config): + super().__init__() + self.config = config + self.module = module + + def generate_sequences(self, prompts: DataProto) -> DataProto: + batch_size = prompts.batch.batch_size[0] + num_chunks = max(batch_size // self.config.get("micro_batch_size", batch_size), 1) + batch_prompts = prompts.chunk(chunks=num_chunks) + output = [self._generate_minibatch(p) for p in batch_prompts] + output = DataProto.concat(output) + return output + + @torch.no_grad() + def _generate_minibatch(self, prompts: DataProto) -> DataProto: + # make sampling args can be overridden by inputs + do_sample = prompts.meta_info.get("do_sample", self.config.do_sample) + is_validate = prompts.meta_info.get("validate", False) + + temperature = prompts.meta_info.get("temperature", self.config.temperature) + response_length = prompts.meta_info.get("response_length", self.config.response_length) + top_p = prompts.meta_info.get("top_p", self.config.get("top_p", 1.0)) + top_k = max(0, prompts.meta_info.get("top_k", self.config.get("top_k", 0))) # to be compatible with vllm + + if not do_sample: + # do_sample==False -> greedy decoding + kwargs = { + "do_sample": False, + "num_beams": 1, + } + elif is_validate: + # do validate and do sample -> use val_kwargs + kwargs = { + "do_sample": True, + "num_beams": 1, + "top_k": max(0, self.config.val_kwargs.top_k), # to be compatible with vllm + "top_p": self.config.val_kwargs.top_p, + "temperature": self.config.val_kwargs.temperature, + "num_return_sequences": 1, # if validate, already repeat in ray_trainer + } + else: + # do_sample -> use rollout config + kwargs = { + "do_sample": True, + "num_beams": 1, + "top_p": top_p, + "top_k": top_k, + "temperature": temperature, + # already repeat in ray_trainer + # https://github.com/verl-project/verl/blob/2fdfbdcba6f2e076f64bc47922d8fe6cf7dc7da5/verl/trainer/ppo/ray_trainer.py#L1117 + "num_return_sequences": 1, + } + + # make config according to generate mode + generation_config = GenerationConfig(**kwargs) + + idx = prompts.batch["input_ids"] # (bs, prompt_length) + prompt_length = idx.size(1) + attention_mask = prompts.batch["attention_mask"] # left-padded attention_mask + position_ids = prompts.batch["position_ids"] + + # used to construct attention_mask + eos_token_id = prompts.meta_info["eos_token_id"] + pad_token_id = prompts.meta_info["pad_token_id"] + + self.module.eval() + param_ctx = contextlib.nullcontext() + + if isinstance(self.module, FSDP): + # recurse need to set to False according to https://github.com/pytorch/pytorch/issues/100069 + param_ctx = FSDP.summon_full_params(self.module, writeback=False, recurse=False) + with param_ctx, torch.autocast(device_type=get_device_name(), dtype=torch.bfloat16): + output = self.module.generate( + input_ids=idx, + attention_mask=attention_mask, + position_ids=position_ids, + do_sample=do_sample, + max_new_tokens=response_length, + eos_token_id=eos_token_id, + pad_token_id=pad_token_id, + generation_config=generation_config, + output_scores=False, # this is potentially very large + return_dict_in_generate=True, + use_cache=True, + ) + + # TODO: filter out the seq with no answers like ds-chat + seq = output.sequences + generated_batch_size = seq.size(0) # bs * num_return_sequences + + # huggingface generate will stop generating when all the batch reaches [EOS]. + # We have to pad to response_length + sequence_length = prompt_length + self.config.response_length + delta_length = sequence_length - seq.shape[1] + + if delta_length > 0: + delta_tokens = torch.ones(size=(generated_batch_size, delta_length), device=seq.device, dtype=seq.dtype) + delta_tokens = pad_token_id * delta_tokens + seq = torch.cat((seq, delta_tokens), dim=1) + assert seq.shape[1] == sequence_length + + # make necessary reputations if num_return_sequences > 1 + num_return_sequences = kwargs.get("num_return_sequences", 1) + if num_return_sequences > 1: + position_ids = position_ids.repeat_interleave(num_return_sequences, dim=0) + attention_mask = attention_mask.repeat_interleave(num_return_sequences, dim=0) + + prompt = seq[:, :prompt_length] # (generated_batch_size, prompt_length) + response = seq[:, prompt_length:] # (generated_batch_size, response_length) + + response_length = response.size(1) + delta_position_id = torch.arange(1, response_length + 1, device=position_ids.device) + delta_position_id = delta_position_id.unsqueeze(0).repeat(generated_batch_size, 1) + + response_position_ids = position_ids[:, -1:] + delta_position_id + position_ids = torch.cat([position_ids, response_position_ids], dim=-1) + + response_attention_mask = get_response_mask( + response_id=response, eos_token=eos_token_id, dtype=attention_mask.dtype + ) + attention_mask = torch.cat((attention_mask, response_attention_mask), dim=-1) + + batch = TensorDict( + { + "prompts": prompt, + "responses": response, + "input_ids": seq, + "attention_mask": attention_mask, + "position_ids": position_ids, + }, + batch_size=generated_batch_size, + ) + + # empty cache before compute old_log_prob + get_torch_device().empty_cache() + + self.module.train() + return DataProto(batch=batch) diff --git a/verl/verl/workers/rollout/llm_server.py b/verl/verl/workers/rollout/llm_server.py new file mode 100644 index 0000000000000000000000000000000000000000..f9e00804ffda557038a79251a8adcdffd3833aba --- /dev/null +++ b/verl/verl/workers/rollout/llm_server.py @@ -0,0 +1,372 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Utility classes for manage and request LLM servers: +- LLMServerManager: manage life-cycle of LLM servers, including launch, tear-down replicas. +- LLMServerClient: proxy client to request LLM servers, used by AgentLoopWorker. +- GlobalRequestLoadBalancer: global load balancer for LLMServerClient. +""" + +import asyncio +import logging +import os +from typing import Any, Optional +from uuid import uuid4 + +import ray +from cachetools import LRUCache +from omegaconf import DictConfig + +from verl.single_controller.ray.base import RayResourcePool, RayWorkerGroup +from verl.utils.ray_utils import auto_await +from verl.utils.rollout_trace import rollout_trace_op +from verl.workers.rollout.replica import RolloutReplica, TokenOutput, get_rollout_replica_class +from verl.workers.rollout.utils import update_prometheus_config + +logger = logging.getLogger(__file__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + +DEFAULT_ROUTING_CACHE_SIZE = 10000 + + +@ray.remote +class GlobalRequestLoadBalancer: + """Global sticky-session + in-flight load balancer shared by all AgentLoopWorkers. + + When a sticky session points to a removed server, the cache entry is + automatically invalidated and a new server is selected. + + Key features: + - **Atomic acquire**: ``acquire_server()`` returns ``(server_id, handle)`` + - **Sticky Session**: Uses LRUCache to map request_id → server_id, ensuring + multi-turn conversations route to the same server. + - **Least-loaded Selection**: When no sticky session exists, selects the + server with the fewest in-flight requests. + - **Dynamic Server Management**: Supports add/remove servers at runtime + for hybrid scaling. + """ + + def __init__(self, servers: dict[str, ray.actor.ActorHandle], max_cache_size: int = DEFAULT_ROUTING_CACHE_SIZE): + if not servers: + raise ValueError("servers must be non-empty") + + self._servers: dict[str, ray.actor.ActorHandle] = dict(servers) + self._inflight_requests: dict[str, int] = {sid: 0 for sid in servers} + self._request_id_to_server: LRUCache = LRUCache(maxsize=max_cache_size) + + def acquire_server(self, request_id: str) -> tuple[str, ray.actor.ActorHandle]: + """Acquire a server for the given request (sticky + least-loaded). + + Returns: + A tuple of ``(server_id, actor_handle)`` in a single atomic call. + """ + # Try sticky session first + if request_id in self._request_id_to_server: + server_id = self._request_id_to_server[request_id] + # Check if server is still in the active pool + if server_id in self._inflight_requests: + self._inflight_requests[server_id] += 1 + return server_id, self._servers[server_id] + # Server was removed, clear stale cache entry and re-select + del self._request_id_to_server[request_id] + + # Select new server (least-loaded among available) + if not self._inflight_requests: + raise RuntimeError("No available servers in load balancer") + + server_id = min(self._inflight_requests, key=self._inflight_requests.get) + self._request_id_to_server[request_id] = server_id + self._inflight_requests[server_id] += 1 + return server_id, self._servers[server_id] + + def release_server(self, server_id: str) -> None: + """Release a server after a request completes.""" + if server_id not in self._inflight_requests: + return + if self._inflight_requests[server_id] > 0: + self._inflight_requests[server_id] -= 1 + + def add_servers(self, servers: dict[str, ray.actor.ActorHandle]) -> None: + """Atomically add multiple servers to the load balancer pool. + + This is more efficient than calling :meth:`add_server` in a loop + because it performs a single bulk update on the internal state. + + Args: + servers: Dict mapping server_id → actor_handle for all servers + to register. + """ + for sid, handle in servers.items(): + self._inflight_requests[sid] = 0 + self._servers[sid] = handle + logger.info(f"[GlobalLoadBalancer] added {len(servers)} servers") + + def remove_servers(self, server_ids: list[str]) -> None: + """Atomically remove multiple servers from the load balancer pool. + + More efficient than calling :meth:`remove_server` in a loop. + + Args: + server_ids: List of server identifiers to remove. + """ + for sid in server_ids: + self._inflight_requests.pop(sid, None) + self._servers.pop(sid, None) + logger.info(f"[GlobalLoadBalancer] removed {len(server_ids)} servers") + + def get_inflight_count(self, server_id: str) -> int: + """Get number of in-flight requests for a server.""" + return self._inflight_requests.get(server_id, 0) + + def get_all_servers(self) -> list[str]: + """Get list of all active server IDs.""" + return list(self._inflight_requests.keys()) + + def get_status(self) -> dict: + """Return current load balancer state for debugging.""" + return { + "servers": dict(self._inflight_requests), + "total_inflight": sum(self._inflight_requests.values()), + "active_servers": len(self._inflight_requests), + "registered_handles": list(self._servers.keys()), + } + + +class LLMServerClient: + """ + A class to manage multiple OpenAI compatible LLM servers. This class provides + - Load balance: least in-flight requests load balancing via global coordination + - Sticky session: send multi-turn chat completions to same server for automatic prefix caching + """ + + def __init__( + self, + config: DictConfig, + load_balancer_handle: ray.actor.ActorHandle, + **kwargs, + ): + """Initialize the LLMServerClient. + + Args: + config (DictConfig): whole config for main entrypoint. + load_balancer_handle (ray.actor.ActorHandle): shared global load balancer actor + that also holds the server-handle registry. + """ + self.config = config + self._load_balancer = load_balancer_handle + + async def _acquire_server(self, request_id: str) -> tuple[str, ray.actor.ActorHandle]: + # Atomic acquire: returns (server_id, handle) in one Ray RPC. + return await self._load_balancer.acquire_server.remote(request_id=request_id) + + def _release_server(self, server_id: str) -> None: + # Fire-and-forget: release is just a counter decrement, no need to await. + # Awaiting here risks blocking the finally clause if the LB actor is unresponsive. + self._load_balancer.release_server.remote(server_id=server_id) + + @rollout_trace_op + async def generate( + self, + request_id, + *, + prompt_ids: list[int], + sampling_params: dict[str, Any], + image_data: Optional[list[Any]] = None, + video_data: Optional[list[Any]] = None, + audio_data: Optional[list[Any]] = None, + mm_processor_kwargs: Optional[dict[str, Any]] = None, + **kwargs: Any, + ) -> TokenOutput: + """Generate tokens from prompt ids. + + Args: + request_id (str): request id for sticky session. + prompt_ids (List[int]): List of prompt token ids. + sampling_params (Dict[str, Any]): Sampling parameters for the chat completion. + + Returns: + TokenOutput | DiffusionOutput: token or diffusion output + """ + server_id, server = await self._acquire_server(request_id) + try: + multimodal_kwargs = {} + if audio_data is not None: + multimodal_kwargs["audio_data"] = audio_data + if mm_processor_kwargs: + multimodal_kwargs["mm_processor_kwargs"] = mm_processor_kwargs + output: TokenOutput = await server.generate.remote( + request_id=uuid4().hex, # use new request_id for each turn + prompt_ids=prompt_ids, + sampling_params=sampling_params, + image_data=image_data, + video_data=video_data, + **multimodal_kwargs, + **kwargs, + ) + return output + finally: + self._release_server(server_id) + + +class LLMServerManager: + """LLMServerManager is responsible for: + - Launch server replicas + - Launch global load balancer + - Elastic launch/tear-down new replicas + + Args: + config (DictConfig): Config for the trainer entrypoint. + worker_group (RayWorkerGroup): Worker group for the server replicas. If not none, init hybrid server, + else init standalone server with a new resource pool. + rollout_resource_pool (RayResourcePool): Resource pool for the server replicas, only needed for TensorRT-LLM. + """ + + def __init__( + self, + config: DictConfig, + worker_group: RayWorkerGroup = None, + rollout_resource_pool: RayResourcePool = None, + ): + self.config = config + self.rollout_config = config.actor_rollout_ref.rollout + self.model_config = config.actor_rollout_ref.model + self.worker_group = worker_group + self.rollout_resource_pool = rollout_resource_pool + + assert worker_group is not None or self.rollout_config.nnodes > 0, "nnodes must be > 0 in standalone mode" + + # for recipe to change + if not hasattr(self, "rollout_replica_class"): + self.rollout_replica_class = get_rollout_replica_class( + self.rollout_config.name, + disaggregation_enabled=self.rollout_config.disaggregation.enabled, + ) + + @classmethod + @auto_await + async def create(cls, *args, **kwargs): + """Create the LLMServerManager.""" + instance = cls(*args, **kwargs) + await instance._initialize_llm_servers() + await instance._init_global_load_balancer() + return instance + + async def _initialize_llm_servers(self, start_rank: int = 0): + """Initialize the LLM server replicas. + + Args: + start_rank: First ``replica_rank`` to assign. Defaults to 0 so that + existing callers are unaffected. Subclasses (e.g. + ``FullyAsyncLLMServerManager``) may pass a non-zero value to avoid + Ray named-actor collisions when hybrid and standalone replicas + coexist. + """ + rollout_world_size = ( + self.rollout_config.tensor_model_parallel_size + * self.rollout_config.data_parallel_size + * self.rollout_config.pipeline_model_parallel_size + ) + # PD inflates per-replica footprint; miss this and init_hybrid slices + # past worker_group → empty workers on replica_rank>=1. + disagg = getattr(self.rollout_config, "disaggregation", None) + if disagg is not None and getattr(disagg, "enabled", False): + prefill_tp = self.rollout_config.tensor_model_parallel_size + # Inline decode_tp default: OmegaConf/Ray serialization drops dataclass methods. + decode_tp = ( + disagg.decode_tensor_model_parallel_size + if disagg.decode_tensor_model_parallel_size is not None + else prefill_tp + ) + rollout_world_size = ( + (prefill_tp * disagg.prefill_replicas + decode_tp * disagg.decode_replicas) + * self.rollout_config.data_parallel_size + * self.rollout_config.pipeline_model_parallel_size + ) + world_size = ( + self.worker_group.world_size + if self.worker_group + else self.rollout_config.n_gpus_per_node * self.rollout_config.nnodes + ) + num_replicas = world_size // rollout_world_size + + self.rollout_replicas = [ + self.rollout_replica_class( + replica_rank=start_rank + replica_rank, + config=self.rollout_config, + model_config=self.model_config, + gpus_per_node=self.rollout_config.n_gpus_per_node, + ) + for replica_rank in range(num_replicas) + ] + + if self.worker_group and self.rollout_config.name != "trtllm": + await asyncio.gather(*[server.init_hybrid(self.worker_group) for server in self.rollout_replicas]) + # TODO: unify trtllm to init_hybrid + elif self.worker_group and self.rollout_config.name == "trtllm": + await asyncio.gather( + *[ + server.init_hybrid_colocated(self.worker_group, self.rollout_resource_pool) + for server in self.rollout_replicas + ] + ) + else: + await asyncio.gather(*[server.init_standalone() for server in self.rollout_replicas]) + + self.server_handles = [server._server_handle for server in self.rollout_replicas] + self.server_addresses = [server._server_address for server in self.rollout_replicas] + print(f"LLMServerManager: {self.server_addresses}") + + # Update Prometheus configuration with server addresses + if self.rollout_config.prometheus.enable: + if self.rollout_config.disable_log_stats: + raise ValueError("PROMETHEUS needs disable_log_stats==False, but it is currently True.") + update_prometheus_config(self.rollout_config.prometheus, self.server_addresses, self.rollout_config.name) + + async def _init_global_load_balancer(self) -> None: + self.global_load_balancer = GlobalRequestLoadBalancer.remote( + servers=dict(zip(self.server_addresses, self.server_handles, strict=True)), + max_cache_size=DEFAULT_ROUTING_CACHE_SIZE, + ) + + def get_client(self, client_cls=LLMServerClient, **kwargs) -> LLMServerClient: + """Get the LLMServerClient to request LLM server replicas. + + Args: + client_cls: The client class to instantiate (default: ``LLMServerClient``). + Pass ``FullyAsyncLLMServerClient`` for abort-resume support. + **kwargs: Forwarded to the client constructor. + """ + return client_cls( + config=self.config, + load_balancer_handle=self.global_load_balancer, + **kwargs, + ) + + def get_addresses(self) -> list[str]: + """Get the OpenAI chat completion API http addresses of the LLM server replicas.""" + return self.server_addresses + + def get_replicas(self) -> list[RolloutReplica]: + """Get the LLM server replicas.""" + return self.rollout_replicas + + @auto_await + async def start_profile(self, **kwargs): + """Start profiling on all rollout replicas.""" + await asyncio.gather(*[replica.start_profile(**kwargs) for replica in self.rollout_replicas]) + + @auto_await + async def stop_profile(self): + """Stop profiling on all rollout replicas.""" + await asyncio.gather(*[replica.stop_profile() for replica in self.rollout_replicas]) diff --git a/verl/verl/workers/rollout/naive/__init__.py b/verl/verl/workers/rollout/naive/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cb6c23bf4327ef199ea9b454f00be88cbaa27967 --- /dev/null +++ b/verl/verl/workers/rollout/naive/__init__.py @@ -0,0 +1,17 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .naive_rollout import NaiveRollout + +__all__ = ["NaiveRollout"] diff --git a/verl/verl/workers/rollout/naive/__pycache__/__init__.cpython-312.pyc b/verl/verl/workers/rollout/naive/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..65a6a9943df505672c675feb509bd0f24b48d7c0 Binary files /dev/null and b/verl/verl/workers/rollout/naive/__pycache__/__init__.cpython-312.pyc differ diff --git a/verl/verl/workers/rollout/naive/__pycache__/naive_rollout.cpython-312.pyc b/verl/verl/workers/rollout/naive/__pycache__/naive_rollout.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1d3a7d7e6459ec610881d223a218139c2f3f6cab Binary files /dev/null and b/verl/verl/workers/rollout/naive/__pycache__/naive_rollout.cpython-312.pyc differ diff --git a/verl/verl/workers/rollout/naive/naive_rollout.py b/verl/verl/workers/rollout/naive/naive_rollout.py new file mode 100644 index 0000000000000000000000000000000000000000..fe56dc4c929b05aa2279b7e1b46e6d9a74e1b175 --- /dev/null +++ b/verl/verl/workers/rollout/naive/naive_rollout.py @@ -0,0 +1,120 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +In single GPU rollout, the sequences are generated directly by sampling from the model. +The output will contain +1. output_ids +2. attention_masks (left padding) +3. eos_masks +4. log_probs +""" + +import torch +import torch.nn.functional as F +from tensordict import TensorDict +from torch import nn + +from verl import DataProto +from verl.utils.torch_functional import logprobs_from_logits + +from ..base import BaseRollout + +__all__ = ["NaiveRollout"] + + +class NaiveRollout(BaseRollout): + def __init__(self, module: nn.Module, config): + """A naive rollout. It requires the module to be compatible with huggingface APIs. That is: + The module should define __call__ to receive input_ids, attention_mask and position_ids. + It outputs a structure that contains logits field. + + Args: + module: module here follows huggingface APIs + config: DictConfig + """ + super().__init__() + self.config = config + self.module = module + + @torch.no_grad() + def generate_sequences(self, prompts: DataProto) -> DataProto: + """Generate sequences""" + idx = prompts.batch["input_ids"] # (bs, prompt_length) + attention_mask = prompts.batch["attention_mask"] # left-padded attention_mask + position_ids = prompts.batch["position_ids"] + + # used to construct attention_mask + eos_token_id = prompts.meta_info["eos_token_id"] + + batch_size = idx.size(0) + prompt_length = idx.size(1) + + self.module.eval() + + prev_attention_mask = torch.ones(size=(batch_size, 1), dtype=attention_mask.dtype, device=attention_mask.device) + + logits_lst = [] + for _ in range(self.config.response_length): + # if the sequence context is growing too long we must crop it at block_size + # idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:] + idx_cond = idx + # forward the model to get the logits for the index in the sequence + # we use huggingface APIs here + output = self.module(input_ids=idx_cond, attention_mask=attention_mask, position_ids=position_ids) + logits = output.logits + # pluck the logits at the final step and scale by desired temperature + logits = logits[:, -1, :] / self.config.temperature # (bs, vocab_size) + # optionally crop the logits to only the top k options + if self.config.top_k is not None: + v, _ = torch.topk(logits, min(self.config.top_k, logits.size(-1))) + logits[logits < v[:, [-1]]] = -float("Inf") + # apply softmax to convert logits to (normalized) probabilities + probs = F.softmax(logits, dim=-1) + # sample from the distribution + if self.config.do_sample: + idx_next = torch.multinomial(probs, num_samples=1) + else: + idx_next = torch.argmax(probs, dim=-1, keepdim=True) + + attention_mask = torch.cat((attention_mask, prev_attention_mask), dim=-1) + + for token_id in eos_token_id: + prev_attention_mask = torch.logical_and(idx_next != token_id, prev_attention_mask.bool()) + prev_attention_mask.to(attention_mask.dtype) + + position_ids = torch.cat((position_ids, position_ids[:, -1:] + 1), dim=-1) + + # append sampled index to the running sequence and continue + idx = torch.cat((idx, idx_next), dim=1) + logits_lst.append(logits) + + logits = torch.stack(logits_lst, dim=1) # (bs, response_length, vocab_size) + prompts = idx[:, :prompt_length] # (bs, prompt_length) + response = idx[:, prompt_length:] # (bs, response_length) + log_probs = logprobs_from_logits(logits=logits, labels=response) + batch = TensorDict( + { + "input_ids": prompts, + "responses": response, + "sequences": idx, + "old_log_probs": log_probs, + "attention_mask": attention_mask, + "position_ids": position_ids, + }, + batch_size=batch_size, + ) + + self.module.train() + + return DataProto(batch=batch) diff --git a/verl/verl/workers/rollout/replica.py b/verl/verl/workers/rollout/replica.py new file mode 100644 index 0000000000000000000000000000000000000000..10e231ec272a102f5c07e619f3845fadb41f832c --- /dev/null +++ b/verl/verl/workers/rollout/replica.py @@ -0,0 +1,402 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import asyncio +import logging +import os +from abc import ABC, abstractmethod +from enum import Enum +from typing import Any, Callable, Optional + +import ray +from omegaconf import DictConfig +from pydantic import BaseModel +from ray.actor import ActorHandle + +from verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup, ResourcePoolManager +from verl.utils.config import omega_conf_to_dataclass +from verl.utils.device import is_torch_npu_available +from verl.workers.config import HFModelConfig, RolloutConfig + +logger = logging.getLogger(__file__) + + +# Max number of concurrent calls to the methods of Rollout, +# excluding calls to generate method. +CONTROL_METHOD_CONCURRENCY = 16 + + +class TokenOutput(BaseModel): + token_ids: list[int] + """response token ids""" + log_probs: Optional[list[float]] = None + """logprobs of response token ids""" + routed_experts: Optional[Any] = None + """routed experts of response token ids""" + stop_reason: Optional[str] = None + """stop reason: 'completed', 'aborted', or None for unknown""" + num_preempted: Optional[int] = None + """number of preempted times for metric calculation""" + extra_fields: dict[str, Any] = {} + """Extra fields for dynamic addition.""" + + +class RolloutMode(Enum): + # Rollout engine and training engine(fsdp/megatron) fused in same process + # Rollout and trainer share GPUs, switch context with weight synchronization. + # Usage scenarios: on-policy training. + HYBRID = "hybrid" + + # Rollout engine colocated with hybrid engine in same ray placement group but in separate process. + # Rollout and hybrid processes share GPUs, switch context without weight synchronization. + # Usage scenarios: GRM (LLM as a judge). + COLOCATED = "colocated" + + # Standalone rollout server with separate GPU resource, disaggregated architecture. + # Usage scenarios: off-policy training. + STANDALONE = "standalone" + + +class RolloutReplica(ABC): + """Rollout replica is an individual server instance, which may be deployed on single or multiple nodes. + It is equivalent to launch server in each node with command line: + + SGLang: + ``` + python -m sglang.launch_server --node-rank 0 --nnode 2 ... + python -m sglang.launch_server --node-rank 1 --nnode 2 ... + ``` + + vLLM: + ``` + vllm serve --data-parallel-size 16 --data-parallel-size-local 8 --data-parallel-start-rank 0 ... + vllm serve --data-parallel-size 16 --data-parallel-size-local 8 --data-parallel-start-rank 8 ... + ``` + + Args: + replica_rank: int, rank of this rollout replica. + config: RolloutConfig, full config. + model_config: DictConfig, model config. + gpus_per_node: int, number of gpus per node. + """ + + def __init__( + self, + replica_rank: int, + config: RolloutConfig, + model_config: DictConfig, + gpus_per_node: int = 8, + is_reward_model: bool = False, + is_teacher_model: bool = False, + name_suffix: str = "", + ) -> None: + self.replica_rank = replica_rank + self.config: RolloutConfig = omega_conf_to_dataclass(config) + self.model_config: HFModelConfig = model_config + + self.world_size = ( + self.config.tensor_model_parallel_size + * self.config.data_parallel_size + * self.config.pipeline_model_parallel_size + ) + self.gpus_per_node = gpus_per_node + self.gpus_per_replica_node = min(gpus_per_node, self.world_size) + assert self.world_size % self.gpus_per_replica_node == 0, ( + f"world_size {self.world_size} must be divisible by gpus_per_node {self.gpus_per_replica_node}" + ) + self.nnodes = self.world_size // self.gpus_per_replica_node + self.is_reward_model = is_reward_model + self.is_teacher_model = is_teacher_model + self.name_suffix = f"_{name_suffix}" if name_suffix else "" + + self.rollout_mode: RolloutMode = None + self.workers: list[ActorHandle] = [] + self.resource_pool: RayResourcePool = None + self.bundle_indices: list[int] = [] + + self.servers: list[ActorHandle] = [] + self._server_address: str = None + self._server_handle: ActorHandle = None + + async def init_hybrid(self, worker_group: RayWorkerGroup): + """Init hybrid rollout server, rollout engine and training engine(fsdp/megatron) fused in same process. + + Args: + worker_group: RayWorkerGroup, fused workers where training engine(fsdp/megatron) have been initialized. + """ + self.rollout_mode = RolloutMode.HYBRID + self.workers = worker_group.workers[ + self.world_size * self.replica_rank : self.world_size * (self.replica_rank + 1) + ] + await self.launch_servers() + + async def init_hybrid_colocated(self, worker_group: RayWorkerGroup, resource_pool: RayResourcePool): + """Init hybrid rollout server, rollout engine and training engine(fsdp/megatron) fused in same process. + + Args: + worker_group: RayWorkerGroup, fused workers where training engine(fsdp/megatron) have been initialized. + resource_pool: RayResourcePool, ray placement group where hybrid engine processes have been launched. + bundle_indices: list[int], bundle indices for this rollout replica. + """ + self.rollout_mode = RolloutMode.HYBRID + self.workers = worker_group.workers[ + self.world_size * self.replica_rank : self.world_size * (self.replica_rank + 1) + ] + self.resource_pool = resource_pool + self.bundle_indices = [self.replica_rank * self.world_size + idx for idx in range(self.world_size)] + await self.launch_servers() + + # TODO(sgm): this should be the default solution, but need to make the RolloutMode more clear. + async def init_colocated(self, resource_pool: RayResourcePool): + """Init colocated rollout server, rollout engine and hybrid engine colocated in same ray placement group + but in separate processes. + + Args: + resource_pool: RayResourcePool, ray placement group where hybrid engine processes have been launched. + """ + self.rollout_mode = RolloutMode.COLOCATED + self.resource_pool = resource_pool + use_gpu = self.rollout_worker_use_gpu() + + if self.is_reward_model: + name_prefix = f"rollout_reward_colocate_{self.replica_rank}{self.name_suffix}" + elif self.is_teacher_model: + name_prefix = f"rollout_teacher_colocate_{self.replica_rank}{self.name_suffix}" + else: + name_prefix = f"rollout_colocate_{self.replica_rank}{self.name_suffix}" + + worker_group = RayWorkerGroup( + resource_pool=self.resource_pool, + ray_cls_with_init=self.get_ray_class_with_init_args(), + bin_pack=False, + name_prefix=name_prefix, + use_gpu=use_gpu, + device_name="cuda" if not is_torch_npu_available(check_device=False) else "npu", + ) + self.workers = worker_group.workers + await self.launch_servers() + + async def init_standalone(self): + """Init standalone rollout server, create new resource pool for this rollout.""" + # create resource pool for this rollout + self.rollout_mode = RolloutMode.STANDALONE + if self.is_reward_model: + resource_pool_name = f"rollout_pool_reward_{self.replica_rank}{self.name_suffix}" + elif self.is_teacher_model: + resource_pool_name = f"rollout_pool_teacher_{self.replica_rank}{self.name_suffix}" + else: + resource_pool_name = f"rollout_pool_{self.replica_rank}{self.name_suffix}" + resource_pool_spec = { + resource_pool_name: [self.gpus_per_replica_node] * self.nnodes, + } + resource_pool_manager = ResourcePoolManager( + resource_pool_spec=resource_pool_spec, + mapping=None, + max_colocate_count=2, + ) + resource_pool_manager.create_resource_pool() + self.resource_pool = resource_pool_manager.resource_pool_dict[resource_pool_name] + + # create worker group for this rollout + if self.is_reward_model: + name_prefix = f"rollout_reward_standalone_{self.replica_rank}{self.name_suffix}" + elif self.is_teacher_model: + name_prefix = f"rollout_teacher_standalone_{self.replica_rank}{self.name_suffix}" + else: + name_prefix = f"rollout_standalone_{self.replica_rank}{self.name_suffix}" + worker_group = RayWorkerGroup( + resource_pool=self.resource_pool, + ray_cls_with_init=self.get_ray_class_with_init_args(), + bin_pack=False, + name_prefix=name_prefix, + use_gpu=True, + device_name="cuda" if not is_torch_npu_available(check_device=False) else "npu", + ) + self.workers = worker_group.workers + await self.launch_servers() + + def get_ray_class_with_init_args(self) -> RayClassWithInitArgs: + """Get rollout worker actor class for colocated and standalone mode.""" + from verl.checkpoint_engine.base import CheckpointEngineWorker + + rollout_worker_actor_cls = ray.remote(CheckpointEngineWorker) + + return RayClassWithInitArgs( + cls=rollout_worker_actor_cls, + rollout_config=self.config, + model_config=self.model_config, + replica_rank=self.replica_rank, + ) + + @abstractmethod + async def launch_servers(self): + """Launch http server in each node.""" + raise NotImplementedError + + @property + def server_address(self) -> str: + """Get rollout server address for OpenAI chat completion.""" + return self._server_address + + @property + def server_handle(self) -> ActorHandle: + """Get rollout server handle for Token-in-token-out generation.""" + return self._server_handle + + @property + def max_concurrency(self) -> int: + # 1000 is Ray's default max_concurrency for async execution. + # Add some margin to account for control method call. + return max(1000, self.config.max_num_seqs + CONTROL_METHOD_CONCURRENCY) + + def rollout_worker_use_gpu(self) -> bool: + return True + + async def wake_up(self): + """Wake up each rollout server.""" + await asyncio.gather(*[server.wake_up.remote() for server in self.servers]) + + async def sleep(self): + """Sleep each rollout server.""" + await asyncio.gather(*[server.sleep.remote() for server in self.servers]) + + async def abort_all_requests(self): + """Partial rollout: abort and save all unfinished requests in each rollout server.""" + await asyncio.gather(*[server.abort_all_requests.remote() for server in self.servers]) + + async def resume_generation(self): + """Resume generation on all servers after abort_all_requests.""" + await asyncio.gather(*[server.resume_generation.remote() for server in self.servers]) + + async def clear_kv_cache(self): + """reset kv cache in each rollout server.""" + await asyncio.gather(*[server.clear_kv_cache.remote() for server in self.servers]) + + async def release_kv_cache(self): + """Release only the kv_cache GPU memory, keeping model weights in place.""" + await asyncio.gather(*[server.release_kv_cache.remote() for server in self.servers]) + + async def resume_kv_cache(self): + """Restore the kv_cache GPU memory after a weight sync.""" + await asyncio.gather(*[server.resume_kv_cache.remote() for server in self.servers]) + + async def start_profile(self, **kwargs): + """Start profiling on the replica.""" + await asyncio.gather(*[server.start_profile.remote(**kwargs) for server in self.servers]) + + async def stop_profile(self): + """Stop profiling on the replica.""" + await asyncio.gather(*[server.stop_profile.remote() for server in self.servers]) + + +class RolloutReplicaRegistry: + """Factory for managing rollout replica implementations.""" + + _registry: dict[str, Callable[[], type[RolloutReplica]]] = {} + + @classmethod + def register(cls, name: str, loader: Callable[[], type[RolloutReplica]]) -> None: + """Register a new rollout replica type.""" + cls._registry[name] = loader + + @classmethod + def get(cls, name: str) -> type[RolloutReplica]: + """Get a rollout replica class by name.""" + if name not in cls._registry: + raise ValueError(f"Unknown rollout mode: {name}. Available: {list(cls._registry.keys())}") + return cls._registry[name]() + + +# Loader functions for built-in types +def _load_vllm(): + from verl.workers.rollout.vllm_rollout.vllm_async_server import vLLMReplica + + return vLLMReplica + + +def _load_sglang(): + os.environ["SGLANG_USE_CPU_ENGINE"] = "1" + + try: + import vllm # noqa: F401 + except ImportError: + import sys + import types + from unittest.mock import Mock + + mock_vllm = types.ModuleType("vllm") + + mock_custom_ops = types.ModuleType("vllm._custom_ops") + mock_custom_ops.scaled_fp8_quant = Mock() + mock_vllm._custom_ops = mock_custom_ops + + mock_model_executor = types.ModuleType("vllm.model_executor") + mock_layers = types.ModuleType("vllm.model_executor.layers") + mock_activation = types.ModuleType("vllm.model_executor.layers.activation") + + class GeluAndMul: # noqa: N801 + pass + + class SiluAndMul: # noqa: N801 + pass + + mock_activation.GeluAndMul = GeluAndMul + mock_activation.SiluAndMul = SiluAndMul + mock_layers.activation = mock_activation + mock_model_executor.layers = mock_layers + mock_vllm.model_executor = mock_model_executor + + sys.modules["vllm"] = mock_vllm + sys.modules["vllm._custom_ops"] = mock_custom_ops + sys.modules["vllm.model_executor"] = mock_model_executor + sys.modules["vllm.model_executor.layers"] = mock_layers + sys.modules["vllm.model_executor.layers.activation"] = mock_activation + + from verl.workers.rollout.sglang_rollout.async_sglang_server import SGLangReplica + + del os.environ["SGLANG_USE_CPU_ENGINE"] + return SGLangReplica + + +def _load_trtllm(): + from verl.workers.rollout.trtllm_rollout.trtllm_async_server import TRTLLMReplica + + return TRTLLMReplica + + +# Register built-in types +RolloutReplicaRegistry.register("vllm", _load_vllm) +RolloutReplicaRegistry.register("sglang", _load_sglang) +RolloutReplicaRegistry.register("trtllm", _load_trtllm) + + +def get_rollout_replica_class(rollout: str, disaggregation_enabled: bool = False) -> type[RolloutReplica]: + """Resolve a replica class by backend name. + + PD-disaggregated SGLang reuses the ``sglang`` backend name; the dispatch + here picks ``SGLangPDReplica`` only when the caller asserts + ``disaggregation_enabled=True`` (sourced from + ``RolloutConfig.disaggregation.enabled``). Validation in + ``RolloutConfig.__post_init__`` blocks the flag for non-SGLang names, so + this function only has to handle the SGLang fork. + """ + if disaggregation_enabled: + if rollout != "sglang": + raise NotImplementedError(f"PD disaggregation is only supported with rollout='sglang'; got {rollout!r}.") + # _load_sglang side-effect: installs vllm mocks needed by SGLangPDReplica's + # transitive imports. Cheap if already installed. + RolloutReplicaRegistry.get("sglang") + from verl.workers.rollout.sglang_rollout.sglang_pd_replica import SGLangPDReplica + + return SGLangPDReplica + return RolloutReplicaRegistry.get(rollout) diff --git a/verl/verl/workers/rollout/schemas.py b/verl/verl/workers/rollout/schemas.py new file mode 100644 index 0000000000000000000000000000000000000000..0b48c94f108b87e29af6ea7a6af068859f763e1d --- /dev/null +++ b/verl/verl/workers/rollout/schemas.py @@ -0,0 +1,713 @@ +# Copyright 2023-2024 SGLang Team +# Copyright 2025 ModelBest Inc. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import difflib +import logging +import os +from enum import Enum +from typing import Any, Optional + +import torch +from pydantic import BaseModel, ConfigDict, model_validator +from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast, ProcessorMixin + +from verl.tools.schemas import OpenAIFunctionToolCall, OpenAIFunctionToolSchema, ToolResponse +from verl.utils.model import compute_position_id_with_mask +from verl.utils.tokenizer import build_multimodal_processor_inputs + +logger = logging.getLogger(__file__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + +BASE_CHAT_HISTORY = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "I am a user."}, +] + + +class FinishReasonTypeEnum(str, Enum): + """The enum for finish reason type.""" + + LENGTH = "length" + STOP = "stop" + TOOL_CALL = "tool_calls" + + @classmethod + def from_str(cls, value: str) -> "FinishReasonTypeEnum": + if value == "stop": + return cls.STOP + elif value == "length": + return cls.LENGTH + elif value == "tool_calls": + return cls.TOOL_CALL + else: + raise ValueError(f"Unsupported finish reason type: {value}") + + +class Message(BaseModel): + role: str + content: str | dict[str, Any] | list[dict[str, Any]] | ToolResponse + tool_calls: Optional[list[OpenAIFunctionToolCall]] = None + + +class AsyncRolloutRequestStateEnum(str, Enum): + """The enum for async rollout request state.""" + + PENDING = "pending" + RUNNING = "running" + COMPLETED = "completed" + FAILED = "failed" + TOOL_CALLING = "tool_calling" + + +class TokenizationSanityCheckModeEnum(str, Enum): + """The enum for tokenization sanity check mode.""" + + DISABLE = "disable" + STRICT = "strict" + IGNORE_STRIPPABLE = "ignore_strippable" + + +class AsyncRolloutRequest(BaseModel): + """The data model for async rollout.""" + + model_config = ConfigDict(arbitrary_types_allowed=True) + + batch_data_id: int = 0 + rollout_offset: int = 0 + request_id: str + state: AsyncRolloutRequestStateEnum + messages: list[Message] + multi_modal_keys: Optional[list[str]] = None + multi_modal_data: Optional[dict[str, Any]] = None + multi_modal_inputs: Optional[dict[str, torch.Tensor]] = None + mm_processor_kwargs: Optional[dict[str, Any]] = None + tool_schemas: Optional[list[OpenAIFunctionToolSchema]] = None + tools_kwargs: dict[str, Any] = {} + input_ids: Optional[torch.Tensor] = None + prompt_ids: Optional[torch.Tensor] = None + response_ids: Optional[torch.Tensor] = None + attention_mask: Optional[torch.Tensor] = None + prompt_attention_mask: Optional[torch.Tensor] = None + response_attention_mask: Optional[torch.Tensor] = None + position_ids: Optional[torch.Tensor] = None + prompt_position_ids: Optional[torch.Tensor] = None + response_position_ids: Optional[torch.Tensor] = None + loss_mask: Optional[torch.Tensor] = None + prompt_loss_mask: Optional[torch.Tensor] = None + response_loss_mask: Optional[torch.Tensor] = None + reward_scores: dict[str, float] + max_prompt_len: int + max_response_len: int = 8192 + max_model_len: int = 32768 + metrics: dict[str, list[Any]] = {} + output_token_ids: torch.Tensor | None = None + rollout_log_probs: torch.Tensor | None = None + + use_inference_chat_template: bool + tokenization_sanity_check_mode: TokenizationSanityCheckModeEnum + generation_prompt_ids: Optional[torch.Tensor] = None + base_conv_wo_gen_prompt_end_pos: int + base_conv_with_gen_prompt_end_pos: int + + @model_validator(mode="before") + @classmethod + def initialize_request(cls, values): + if not (messages := values.get("messages")): + raise ValueError("messages is required for AsyncRolloutRequest initialization") + if not (max_prompt_len := values.get("max_prompt_len")): + raise ValueError("max_prompt_len is required for AsyncRolloutRequest initialization") + if not (processing_class := values.pop("processing_class", None)): + raise ValueError("processing_class is required for AsyncRolloutRequest initialization") + + values["messages"] = [Message.model_validate(msg) for msg in messages] + + # If there is no multi_modal_keys, we assume the multi-modal data is image, video and audio. + if not values.get("multi_modal_keys"): + values["multi_modal_keys"] = ["image", "video", "audio"] + if not values.get("multi_modal_data"): + values["multi_modal_data"] = {key: [] for key in values["multi_modal_keys"]} + else: + # check if all multi_modal_keys are in multi_modal_data + for key in values["multi_modal_keys"]: + if key not in values["multi_modal_data"]: + values["multi_modal_data"][key] = [] + if not values.get("multi_modal_inputs"): + values["multi_modal_inputs"] = {} + if not values.get("mm_processor_kwargs"): + values["mm_processor_kwargs"] = {} + + tools = ( + [tool.model_dump() for tool in tool_schemas] if (tool_schemas := values.get("tool_schemas", [])) else None + ) + + multi_modal_data = values["multi_modal_data"] + mm_processor_kwargs = values["mm_processor_kwargs"] + tokens_without_prompt = cls._handle_apply_chat_template( + processing_class, + messages, + multi_modal_data=multi_modal_data, + mm_processor_kwargs=mm_processor_kwargs, + tools=tools, + add_generation_prompt=False, + tokenize=True, + ) + if ( + values.get("input_ids") is None + or values.get("attention_mask") is None + or values.get("position_ids") is None + ): + tokenization_dict_with_prompt = cls._handle_apply_chat_template( + processing_class, + messages, + multi_modal_data=multi_modal_data, + mm_processor_kwargs=mm_processor_kwargs, + tools=tools, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + ) + + values["input_ids"], values["attention_mask"] = ( + tokenization_dict_with_prompt["input_ids"], + tokenization_dict_with_prompt["attention_mask"], + ) + if values["input_ids"].shape[-1] > max_prompt_len: + # Only log the warning to avoid truncating in the middle of generation prompt. Consider raising an + # error for this case in the future. + # Ensure batch_data_id exists with default value if not provided + if "batch_data_id" not in values: + values["batch_data_id"] = cls.model_fields["batch_data_id"].default + logger.warning( + f"Prompt {values['batch_data_id']} has length {values['input_ids'].shape[-1]} " + f"which is greater than max_prompt_len {max_prompt_len} after applied chat template with tools." + ) + + # Process multi_modal_inputs + multi_modal_inputs = tokenization_dict_with_prompt.copy() + multi_modal_inputs.pop("input_ids", None) + multi_modal_inputs.pop("attention_mask", None) + values["multi_modal_inputs"] = multi_modal_inputs + + values["position_ids"] = values["prompt_position_ids"] = cls._get_position_ids( + processing_class, + values["input_ids"], + values["attention_mask"], + multi_modal_inputs, + mm_processor_kwargs=mm_processor_kwargs, + ) + + values["prompt_ids"], values["prompt_attention_mask"] = values["input_ids"], values["attention_mask"] + values["loss_mask"] = values["prompt_loss_mask"] = torch.zeros_like(values["input_ids"], dtype=torch.bool) + values["generation_prompt_ids"] = values["input_ids"][..., tokens_without_prompt.shape[-1] :] + values["base_conv_wo_gen_prompt_end_pos"] = cls._handle_apply_chat_template( + processing_class, + BASE_CHAT_HISTORY, + multi_modal_data=multi_modal_data, + mm_processor_kwargs=mm_processor_kwargs, + tools=tools, + add_generation_prompt=False, + tokenize=True, + ).shape[-1] + + values["base_conv_with_gen_prompt_end_pos"] = cls._handle_apply_chat_template( + processing_class, + BASE_CHAT_HISTORY, + multi_modal_data=multi_modal_data, + mm_processor_kwargs=mm_processor_kwargs, + tools=tools, + add_generation_prompt=True, + tokenize=True, + ).shape[-1] + + return values + + @staticmethod + def _handle_apply_chat_template( + processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin, + messages: list[Message], + multi_modal_data: dict[str, Any], + mm_processor_kwargs: Optional[dict[str, Any]] = None, + tools: Optional[list[OpenAIFunctionToolSchema]] = None, + add_generation_prompt: bool = False, + tokenize: bool = False, + return_dict: bool = False, + ): + raw_prompt = processing_class.apply_chat_template( + messages, tools=tools, add_generation_prompt=add_generation_prompt, tokenize=False + ) + if not tokenize: + return raw_prompt + + if isinstance(processing_class, PreTrainedTokenizer) or isinstance(processing_class, PreTrainedTokenizerFast): + if any(len(values) > 0 for values in multi_modal_data.values()): + logger.warning( + "There is multi_modal_data but you are not using a processor. Multi-modal data will be ignored." + ) + model_inputs = processing_class(text=[raw_prompt], return_tensors="pt") + elif isinstance(processing_class, ProcessorMixin): + images = images if len(images := multi_modal_data.get("image", [])) > 0 else None + videos = videos if len(videos := multi_modal_data.get("video", [])) > 0 else None + audio = audio if len(audio := multi_modal_data.get("audio", [])) > 0 else None + model_inputs = build_multimodal_processor_inputs( + processing_class, + text=[raw_prompt], + images=images, + videos=videos, + audio=audio, + mm_processor_kwargs=mm_processor_kwargs, + ) + else: + raise ValueError(f"Unsupported processing class type: {type(processing_class)}") + + if hasattr(model_inputs, "convert_to_tensors"): + model_inputs = model_inputs.convert_to_tensors("pt") + model_inputs = dict(model_inputs) + if return_dict: + return model_inputs + else: + return model_inputs["input_ids"] + + @staticmethod + def _get_position_ids( + processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + multi_modal_inputs: Optional[dict[str, torch.Tensor]] = None, + mm_processor_kwargs: Optional[dict[str, Any]] = None, + ) -> torch.Tensor: + # special case for qwen2vl + is_qwen2vl = ( + hasattr(processing_class, "image_processor") + and "Qwen2VLImageProcessor" in processing_class.image_processor.__class__.__name__ + ) + if is_qwen2vl: + from verl.models.transformers.qwen2_vl import get_rope_index + + image_grid_thw = video_grid_thw = second_per_grid_ts = None + if multi_modal_inputs: + image_grid_thw = multi_modal_inputs.get("image_grid_thw") + video_grid_thw = multi_modal_inputs.get("video_grid_thw") + second_per_grid_ts = multi_modal_inputs.get("second_per_grid_ts") + + assert input_ids.dim() == 2 and input_ids.shape[0] == 1, ( + f"input_ids should be 2D with batch size 1, but got shape {input_ids.shape}" + ) + assert attention_mask.dim() == 2 and attention_mask.shape[0] == 1, ( + f"attention_mask should be 2D with batch size 1, but got shape {attention_mask.shape}" + ) + new_position_ids = get_rope_index( + processing_class, + input_ids=input_ids.squeeze(0), + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + attention_mask=attention_mask.squeeze(0), + ) + return new_position_ids # (3, seq_len) + else: + return compute_position_id_with_mask(attention_mask) # (1, seq_len) + + def _update_input_ids( + self, + processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin, + new_input_ids: torch.Tensor, + attention_mask: bool, + loss_mask: bool, + new_multi_modal_inputs: Optional[dict[str, torch.Tensor]] = None, + ) -> None: + """ + Update the input_ids, attention_mask, position_ids, and loss_mask of the request in additive manner. + """ + self.input_ids = torch.cat([self.input_ids, new_input_ids], dim=-1) + attention_mask = torch.ones_like(new_input_ids) * int(attention_mask) + self.attention_mask = torch.cat([self.attention_mask, attention_mask], dim=-1) + loss_mask = torch.ones_like(new_input_ids) * int(loss_mask) + self.loss_mask = torch.cat([self.loss_mask, loss_mask], dim=-1) + + if new_multi_modal_inputs: + self._update_multi_modal_inputs(new_multi_modal_inputs) + + new_position_ids = self._get_position_ids( + processing_class, + new_input_ids, + attention_mask, + new_multi_modal_inputs, + mm_processor_kwargs=self.mm_processor_kwargs, + ) + + last_pos = self.position_ids[..., -1:] + new_position_ids = new_position_ids + (last_pos + 1) + + self.position_ids = torch.cat([self.position_ids, new_position_ids], dim=-1) + + assert ( + self.input_ids.shape[-1] + == self.attention_mask.shape[-1] + == self.position_ids.shape[-1] + == self.loss_mask.shape[-1] + ), f"""Request {self.request_id} has different length of {self.input_ids.shape[-1]=}, + {self.attention_mask.shape[-1]=}, {self.position_ids.shape[-1]=}, {self.loss_mask.shape[-1]=}""" + + def _update_multi_modal_inputs(self, new_multi_modal_inputs: dict[str, torch.Tensor]) -> None: + """ + Update the multi_modal_inputs of the request in additive manner. + """ + for key in new_multi_modal_inputs: + input_tensor = new_multi_modal_inputs[key] + self.multi_modal_inputs[key] = ( + torch.cat([self.multi_modal_inputs[key], input_tensor], dim=0) + if key in self.multi_modal_inputs + else input_tensor + ) + + def get_generation_prompt_ids( + self, processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin + ) -> list[int]: + """ + Get the generation prompt ids for rollout engine. + + Because rollout engine(SGLang) requires the ids to be a list, we need to convert the tensor to a list. + """ + generation_prompt_ids = ( + None + if self.input_ids[..., -self.generation_prompt_ids.shape[-1] :].eq(self.generation_prompt_ids).all() + else self.generation_prompt_ids + ) + if generation_prompt_ids is not None: + self._update_input_ids(processing_class, generation_prompt_ids, attention_mask=True, loss_mask=False) + + if self.use_inference_chat_template: + messages = [msg.model_dump() for msg in self.messages] + tools = [tool.model_dump() for tool in self.tool_schemas] if self.tool_schemas else None + generation_prompt_ids = self._handle_apply_chat_template( + processing_class, + messages, + multi_modal_data=self.multi_modal_data, + mm_processor_kwargs=self.mm_processor_kwargs, + tools=tools, + add_generation_prompt=True, + tokenize=True, + ) + return generation_prompt_ids.squeeze(0).tolist() + else: + return self.input_ids.squeeze(0).tolist() + + def add_user_message( + self, + processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin, + content: str, + ) -> None: + self.messages.append(Message(role="user", content=content)) + messages = [*BASE_CHAT_HISTORY, self.messages[-1]] + tools = [tool.model_dump() for tool in self.tool_schemas] if self.tool_schemas else None + + # We don't need to pass multi_modal_data here because we don't have any multi-modal data from Engine + # Inference, it is pure text. + content_ids = self._handle_apply_chat_template( + processing_class, + messages, + multi_modal_data={}, + mm_processor_kwargs={}, + tools=tools, + add_generation_prompt=False, + tokenize=True, + )[..., self.base_conv_wo_gen_prompt_end_pos :] + self._update_input_ids(processing_class, content_ids, attention_mask=True, loss_mask=False) + + def add_assistant_message( + self, + processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin, + content: str, + content_ids: Optional[torch.Tensor] = None, + tool_calls: Optional[list[OpenAIFunctionToolCall]] = None, + ) -> None: + self.messages.append(Message(role="assistant", content=content, tool_calls=tool_calls)) + if content_ids is None: + messages = [*BASE_CHAT_HISTORY, self.messages[-1]] + tools = [tool.model_dump() for tool in self.tool_schemas] if self.tool_schemas else None + + # We don't need to pass multi_modal_data here because we don't have any multi-modal data from Engine + # Inference, it is pure text. + content_ids = self._handle_apply_chat_template( + processing_class, + messages, + multi_modal_data={}, + mm_processor_kwargs={}, + tools=tools, + add_generation_prompt=False, + tokenize=True, + )[..., self.base_conv_with_gen_prompt_end_pos :] + self._update_input_ids(processing_class, content_ids, attention_mask=True, loss_mask=True) + + def add_tool_response_messages( + self, + processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin, + contents: list[ToolResponse], + ) -> None: + if not contents or all(content.is_empty() for content in contents): + return + # We also handle the case when tool returns image + # We require the processing of the image and video to be done at tool.execute() level + delta_multi_modal_data = {key: [] for key in self.multi_modal_keys} + for content in contents: + if content.is_text_only(): + self.messages.append(Message(role="tool", content=content.text)) + else: + content_list = [] + # When we update multi_model_keys, we also need to update this logic + if content.image: + content_list.extend([{"type": "image"} for _ in content.image]) + delta_multi_modal_data["image"].extend(content.image) + if content.video: + content_list.extend([{"type": "video"} for _ in content.video]) + delta_multi_modal_data["video"].extend(content.video) + if content.text: + content_list.append({"type": "text", "text": content.text}) + self.messages.append(Message(role="tool", content=content_list)) + + messages = [*BASE_CHAT_HISTORY, *self.messages[-len(contents) :]] + tools = [tool.model_dump() for tool in self.tool_schemas] if self.tool_schemas else None + + for key in self.multi_modal_keys: + if len(delta_multi_modal_data[key]) > 0: + self.multi_modal_data[key].extend(delta_multi_modal_data[key]) + + # We just passed the new multi-modal data to the chat template to update the input_ids. + content_info = self._handle_apply_chat_template( + processing_class, + messages, + multi_modal_data=delta_multi_modal_data, + mm_processor_kwargs=self.mm_processor_kwargs, + tools=tools, + add_generation_prompt=False, + tokenize=True, + return_dict=True, + ) + content_ids = content_info["input_ids"][..., self.base_conv_wo_gen_prompt_end_pos :] + + # process multi_modal_inputs + multi_modal_inputs = content_info.copy() + multi_modal_inputs.pop("input_ids", None) + multi_modal_inputs.pop("attention_mask", None) + + # chat templates include generation prompt tokens (e.g., "assistant\n") + # So when tool response is added, we need to explicitly remove these tokens. + self._remove_generation_prompt_ids_if_present() + + self._update_input_ids( + processing_class, + content_ids, + attention_mask=True, + loss_mask=False, + new_multi_modal_inputs=multi_modal_inputs, + ) + + def update_metrics(self, metrics: Any, tool_id: str) -> None: + """ + metrics: should be a dict of tools_name -> Any + """ + if self.metrics.get(tool_id) is None: + self.metrics[tool_id] = [] + self.metrics[tool_id].append(metrics) + + def _get_prompt_diffs( + self, + processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin, + full_prompt_ids: torch.Tensor, + current_prompt_ids: torch.Tensor, + diff_surrounding_chars: int = 10, + ) -> list[dict[str, Any]]: + """Get differences between full prompt and current prompt with surrounding context. + + This function helps debug tokenization mismatches by showing the differences between + full prompt and current prompt with surrounding context. Instead of just showing + the exact diff, it includes additional tokens before and after to help locate + the issue in the chat template. + + For example, if the actual diff is a newline change from "\n\n" to "\n", with + diff_surrounding_chars the output might look like: + + full_prompt_chunk: "<|im_start|>assistant\n\nI think..." + current_prompt_chunk: "<|im_start|>assistant\nI think..." + + This context makes it much easier to identify where in the chat template the + mismatch occurs. + + Args: + processing_class: The processing class to use for decoding the token IDs + full_prompt_ids: Token IDs from applying chat template to all messages at once + current_prompt_ids: Token IDs from incremental chat template application + diff_surrounding_chars: Number of surrounding characters to include for context (default: 10) + + Returns: + List of dicts containing the differing chunks with context and their indices + """ + full_prompt_ids = full_prompt_ids.squeeze(0) + current_prompt_ids = current_prompt_ids.squeeze(0) + full_prompt = processing_class.decode(full_prompt_ids, skip_special_tokens=False) + current_prompt = processing_class.decode(current_prompt_ids, skip_special_tokens=False) + s = difflib.SequenceMatcher(None, full_prompt, current_prompt, autojunk=False) + diffs = [] + for tag, i1, i2, j1, j2 in s.get_opcodes(): + if tag == "equal": + continue + + # Get the surrounding context for better readability + start_i = max(0, i1 - diff_surrounding_chars) + end_i = min(len(full_prompt), i2 + diff_surrounding_chars) + start_j = max(0, j1 - diff_surrounding_chars) + end_j = min(len(current_prompt), j2 + diff_surrounding_chars) + + diffs.append( + { + "full_prompt_chunk": full_prompt[start_i:end_i], + "current_prompt_chunk": current_prompt[start_j:end_j], + "indices": (start_i, end_i, start_j, end_j), + } + ) + return diffs + + def _remove_generation_prompt_ids_if_present(self) -> None: + """ + Remove generation prompt IDs from input tensors if they are present at the end. + """ + if self.input_ids[..., -self.generation_prompt_ids.shape[-1] :].eq(self.generation_prompt_ids).all(): + self.input_ids = self.input_ids[..., : -self.generation_prompt_ids.shape[-1]] + self.attention_mask = self.attention_mask[..., : -self.generation_prompt_ids.shape[-1]] + self.position_ids = self.position_ids[..., : -self.generation_prompt_ids.shape[-1]] + self.loss_mask = self.loss_mask[..., : -self.generation_prompt_ids.shape[-1]] + + def finalize( + self, + processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin, + reward_scores: dict[str, list[float]], + finish_reason_type: FinishReasonTypeEnum = FinishReasonTypeEnum.STOP, + ) -> None: + self.state = AsyncRolloutRequestStateEnum.COMPLETED + self.reward_scores = reward_scores + + # In case we failed to generate the assistant message and the generation prompt ids were already added to + # input_ids, remove them from the end of input_ids + self._remove_generation_prompt_ids_if_present() + + self.response_ids = self.input_ids[..., self.prompt_ids.shape[-1] :] + + if self.tokenization_sanity_check_mode != TokenizationSanityCheckModeEnum.DISABLE: + # When there is a diff, we log the diffs with diff_surrounding_chars context + diff_surrounding_chars = 10 + + messages = [msg.model_dump() for msg in self.messages] + tools = [tool.model_dump() for tool in self.tool_schemas] if self.tool_schemas else None + full_prompt_info = self._handle_apply_chat_template( + processing_class, + messages, + multi_modal_data=self.multi_modal_data, + mm_processor_kwargs=self.mm_processor_kwargs, + tools=tools, + add_generation_prompt=False, + tokenize=True, + return_dict=True, + ) + full_prompt_ids = full_prompt_info["input_ids"] + + # We must use dict(full_prompt_info) to convert BatchFeature values to a new dict + # because np.array() only keeps the keys for BatchFeature. + full_prompt_multi_modal_inputs = full_prompt_info.copy() + full_prompt_multi_modal_inputs.pop("input_ids", None) + full_prompt_multi_modal_inputs.pop("attention_mask", None) + + for multi_modal_inputs_key in self.multi_modal_inputs: + if multi_modal_inputs_key in full_prompt_multi_modal_inputs: + if ( + not self.multi_modal_inputs[multi_modal_inputs_key] + .eq(full_prompt_multi_modal_inputs[multi_modal_inputs_key]) + .all() + ): + logger.warning( + f"Multi-modal data {multi_modal_inputs_key} is not consistent. " + f"This may lead to unexpected behavior during training. " + f"Please review your multi_modal_inputs logic." + ) + else: + logger.warning( + f"Multi-modal inputs key {multi_modal_inputs_key} is not found in the multi_modal_inputs. " + f"This may lead to unexpected behavior during training." + f"Please review your multi_modal_inputs logic." + ) + + if diffs := self._get_prompt_diffs( + processing_class, full_prompt_ids, self.input_ids, diff_surrounding_chars=diff_surrounding_chars + ): + log_warning = False + if self.tokenization_sanity_check_mode == TokenizationSanityCheckModeEnum.STRICT: + log_warning = True + elif self.tokenization_sanity_check_mode == TokenizationSanityCheckModeEnum.IGNORE_STRIPPABLE: + non_strippable_diffs_exist = any( + d["full_prompt_chunk"].strip() or d["current_prompt_chunk"].strip() for d in diffs + ) + if non_strippable_diffs_exist: + log_warning = True + + if log_warning: + mode_str = f" ({self.tokenization_sanity_check_mode.value})" + logger.warning( + f"Inconsistent training and inference tokenization detected{mode_str}. This may lead to " + f"unexpected behavior during training. Please review your chat template to determine if this " + f"is intentional. For more information, refer to the multiturn README.md." + ) + logger.warning( + f"Showing {diff_surrounding_chars} characters before and after the diffs for context and " + f"better readability." + ) + diff_details_list = [] + for d in diffs: + i1, i2, j1, j2 = d["indices"] + diff_details_list.append( + f"idx {i1}:{i2} -> {j1}:{j2} | full_prompt_chunk: {repr(d['full_prompt_chunk'])} | " + f"current_prompt_chunk: {repr(d['current_prompt_chunk'])}" + ) + diff_details = "\n".join(diff_details_list) + logger.warning(f"Found differences:\n{diff_details}") + + if finish_reason_type == FinishReasonTypeEnum.STOP: + pass + elif finish_reason_type == FinishReasonTypeEnum.LENGTH: + pass + else: + raise ValueError(f"Unsupported finalize finish reason type: {finish_reason_type}") + self.truncate_output_ids(processing_class) + + assert ( + self.input_ids.shape[-1] + == self.attention_mask.shape[-1] + == self.position_ids.shape[-1] + == self.loss_mask.shape[-1] + ), f"""Request {self.request_id} has different length of {self.input_ids.shape[-1]=}, + {self.attention_mask.shape[-1]=}, {self.position_ids.shape[-1]=}, {self.loss_mask.shape[-1]=}""" + + def truncate_output_ids( + self, processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin + ) -> None: + self.input_ids = self.input_ids[..., : self.max_model_len] + self.attention_mask = self.attention_mask[..., : self.max_model_len] + self.position_ids = self.position_ids[..., : self.max_model_len] + self.loss_mask = self.loss_mask[..., : self.max_model_len] + self.response_ids = self.input_ids[..., self.prompt_ids.shape[-1] :][..., : self.max_response_len] + self.response_attention_mask = self.attention_mask[..., self.prompt_attention_mask.shape[-1] :][ + ..., : self.max_response_len + ] + self.response_position_ids = self.position_ids[..., self.prompt_position_ids.shape[-1] :][ + ..., : self.max_response_len + ] + self.response_loss_mask = self.loss_mask[..., self.prompt_loss_mask.shape[-1] :][..., : self.max_response_len] diff --git a/verl/verl/workers/rollout/sglang_rollout/__init__.py b/verl/verl/workers/rollout/sglang_rollout/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..337ec74bce1922b8af69a9879fcbf0f2b382c555 --- /dev/null +++ b/verl/verl/workers/rollout/sglang_rollout/__init__.py @@ -0,0 +1,12 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and diff --git a/verl/verl/workers/rollout/sglang_rollout/async_sglang_server.py b/verl/verl/workers/rollout/sglang_rollout/async_sglang_server.py new file mode 100644 index 0000000000000000000000000000000000000000..3731f30bc3e01d1cbb7e29cc587d0c4a913d9324 --- /dev/null +++ b/verl/verl/workers/rollout/sglang_rollout/async_sglang_server.py @@ -0,0 +1,822 @@ +# Copyright 2023-2024 SGLang Team +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import asyncio +import dataclasses +import json +import logging +import os +import secrets +from pathlib import Path +from typing import Any, Optional + +import ray +import sglang +import sglang.srt.entrypoints.engine +import torch +from packaging import version +from ray.actor import ActorHandle +from sglang.srt.entrypoints.http_server import ( + ServerArgs, + _GlobalState, + app, + set_global_state, +) +from sglang.srt.managers.io_struct import ( + ContinueGenerationReqInput, + GenerateReqInput, + PauseGenerationReqInput, + ReleaseMemoryOccupationReqInput, + ResumeMemoryOccupationReqInput, +) +from sglang.srt.managers.tokenizer_manager import ServerStatus + +from verl.utils.config import omega_conf_to_dataclass +from verl.utils.device import get_visible_devices_keyword +from verl.utils.net_utils import get_free_port, is_valid_ipv6_address +from verl.utils.profiler import DistProfiler, build_sglang_profiler_args +from verl.workers.config import HFModelConfig, RolloutConfig +from verl.workers.rollout.replica import RolloutMode, RolloutReplica, TokenOutput +from verl.workers.rollout.sglang_rollout.sglang_rollout import _set_envs_and_config +from verl.workers.rollout.sglang_rollout.utils import SGLANG_LORA_NAME +from verl.workers.rollout.utils import get_max_position_embeddings, run_uvicorn + +logger = logging.getLogger(__file__) +logger.setLevel(logging.INFO) + +visible_devices_keyword = get_visible_devices_keyword() + + +def _extract_prompt_logprobs_sglang( + meta_info: dict, + num_prompt_logprobs: int, + sequence_length: int, + result_dict: dict[str, list], +) -> None: + """Shape SGLang input-logprobs into the vLLM ``extract_prompt_logprobs`` contract. + Populates ``result_dict`` with two ``[sequence_length, max(num_prompt_logprobs, 1)]`` + lists — ``prompt_ids`` and ``prompt_logprobs`` — so the distillation teacher + consumer in ``teacher_manager.AsyncTeacherLLMServerManager`` can treat vLLM and + SGLang teachers interchangeably. + SGLang returns input logprobs with length ``S == len(input_ids)`` whose first + entry has ``logprob=None`` (no predicting context). That matches the vLLM + convention, so we skip entry 0 and append a trailing dummy row to keep the + total length equal to the consumer's ``len(sequence_ids)`` assertion. + """ + input_token_logprobs = meta_info.get("input_token_logprobs") or [] + if num_prompt_logprobs > 0: + input_top_logprobs = meta_info.get("input_top_logprobs") or [] + prompt_ids_ls: list[list[int]] = [] + prompt_logprobs_ls: list[list[float]] = [] + # Entry 0 has logprob=None (no predicting context); skip it, matching vLLM. + for position in range(1, len(input_token_logprobs)): + if num_prompt_logprobs == 0: + logprob, token_id, _ = input_token_logprobs[position] + prompt_ids_ls.append([int(token_id)]) + prompt_logprobs_ls.append([float(logprob)]) + else: + top_entries = input_top_logprobs[position] + # SGLang returns ranked best-first; we preserve that ordering so rank + # 0 is the top-1 token, matching the vLLM extractor's rank-1 slot. + ids = [int(tok_id) for _, tok_id, _ in top_entries] + logprobs = [float(logprob) for logprob, _, _ in top_entries] + assert len(ids) == num_prompt_logprobs, ( + f"SGLang returned {len(ids)} top logprobs at position {position}, expected {num_prompt_logprobs}." + ) + prompt_ids_ls.append(ids) + prompt_logprobs_ls.append(logprobs) + # Trailing dummy row so total length == len(sequence_ids), matching vLLM. + pad_width = max(num_prompt_logprobs, 1) + prompt_ids_ls.append([0] * pad_width) + prompt_logprobs_ls.append([0.0] * pad_width) + assert len(prompt_ids_ls) == sequence_length, ( + f"SGLang prompt_logprobs length ({len(prompt_ids_ls)}) does not match " + f"sequence length ({sequence_length}); check logprob_start_len=0 invariant." + ) + result_dict["prompt_ids"] = prompt_ids_ls + result_dict["prompt_logprobs"] = prompt_logprobs_ls + + +class SGLangHttpServer: + """SGLang http server in single node, this is equivalent to launch server with command line: + ``` + python -m sglang.launch_server --node-rank 0 --nnode 1 ... + ``` + + Args: + config (DictConfig): full config. + rollout_mode (RolloutMode): rollout mode. + replica_rank (int): replica rank, a replica may contain multiple nodes. + node_rank (int): node rank. + nnodes (int): number of nodes. + cuda_visible_devices (str): cuda visible devices. + """ + + def __init__( + self, + config: RolloutConfig, + model_config: HFModelConfig, + rollout_mode: RolloutMode, + workers: list[ActorHandle], + replica_rank: int, + node_rank: int, + nnodes: int, + cuda_visible_devices: str, + base_gpu_id: int, + disaggregation_role: str = "null", + disaggregation_bootstrap_port: Optional[int] = None, + ): + print( + f"SGLang http server: {rollout_mode=}, {replica_rank=}, {node_rank=}, " + f"{nnodes=}, {cuda_visible_devices=}, role={disaggregation_role}" + ) + os.environ[visible_devices_keyword] = cuda_visible_devices + + assert disaggregation_role in ("null", "prefill", "decode"), ( + f"disaggregation_role must be 'null'|'prefill'|'decode', got {disaggregation_role!r}" + ) + self._disaggregation_role = disaggregation_role + self._disaggregation_bootstrap_port = disaggregation_bootstrap_port + + self.config: RolloutConfig = omega_conf_to_dataclass(config) + self.model_config: HFModelConfig = omega_conf_to_dataclass(model_config, dataclass_type=HFModelConfig) + max_position_embeddings = get_max_position_embeddings(self.model_config.hf_config) + if self.config.max_model_len is None: + self.config.max_model_len = max_position_embeddings + else: + if self.config.max_model_len > max_position_embeddings: + raise ValueError( + f"max_model_len ({self.config.max_model_len}) should be less than or equal to " + f"max_position_embeddings ({max_position_embeddings})" + ) + self.rollout_mode = rollout_mode + self.workers = workers + + self.replica_rank = replica_rank + self.node_rank = node_rank + self.nnodes = nnodes + self.base_gpu_id = base_gpu_id + # model weights version, set by ServerAdapter when update weights. + self.global_steps = None + + # PD peer linkage populated post-launch by SGLangPDReplica.set_pd_peer. + self._pd_decode_peers: list[ActorHandle] = [] + self._pd_bootstrap_host: Optional[str] = None + + if self.rollout_mode != RolloutMode.HYBRID and self.config.load_format == "dummy": + logger.warning(f"rollout mode is {self.rollout_mode}, load_format is dummy, set to auto") + self.config.load_format = "auto" + + # used for http server + self._server_address = ray.util.get_node_ip_address().strip("[]") + self._server_port = None + + # used for controlling sglang server profiler + profiler_config = self.config.profiler + tool_config = None + if profiler_config is not None: + if profiler_config.tool in ["torch", "npu"]: + tool_config = omega_conf_to_dataclass((profiler_config.tool_config or {}).get(profiler_config.tool)) + else: + logger.warning(f"agent loop only support torch and npu profiler, got {profiler_config.tool}") + profiler_config = None + self.profiler_controller = DistProfiler(self.replica_rank, config=profiler_config, tool_config=tool_config) + + # For multi-node, we need dist_init_addr so nodes can coordinate NCCL init. + # For single-node, let SGLang handle port selection internally via nccl_port, + # which also avoids port conflicts. + self._master_address = None + self._master_port = None + self._master_sock = None + if self.nnodes > 1 and self.node_rank == 0: + self._master_address = self._server_address + self._master_port, self._master_sock = get_free_port(self._server_address, with_alive_sock=True) + logger.info( + f"SGLangHttpServer, replica_rank: {self.replica_rank}, " + f"master address: {self._master_address}, port: {self._master_port}" + ) + + def get_master_address(self): + """Get master address and port for init NCCL process group.""" + return self._master_address, self._master_port + + def get_server_address(self): + """Get http server address and port.""" + assert self._server_port is not None, "http server is not launched, port is None" + return self._server_address, self._server_port + + async def set_pd_peer(self, decode_peers: list, bootstrap_host: str): + assert isinstance(decode_peers, list) and decode_peers + self._pd_decode_peers = list(decode_peers) + self._pd_bootstrap_host = bootstrap_host + + def _prepend_cu12_lib_to_ld_library_path(self) -> None: + """Ray runtime_env.pip installs cu12 into a transient venv, not the usual + site-packages. NIXL's UCX plugin dlopens libcudart.so.12 from + LD_LIBRARY_PATH; wrong path ⇒ scheduler subprocess dies with SIGABRT.""" + try: + import nvidia.cuda_runtime as cu12_mod + except ImportError as e: + logger.warning( + f"nvidia.cuda_runtime not importable: {e}. " + f"NIXL may fail with 'libcudart.so.12: cannot open shared object'." + ) + return + cu12_lib = str(Path(cu12_mod.__file__).parent / "lib") + if not os.path.isdir(cu12_lib): + return + existing = os.environ.get("LD_LIBRARY_PATH", "") + if cu12_lib in existing.split(":"): + return + os.environ["LD_LIBRARY_PATH"] = f"{cu12_lib}:{existing}" if existing else cu12_lib + logger.info(f"Prepended {cu12_lib} to LD_LIBRARY_PATH for NIXL/UCX dlopen.") + + async def launch_server(self, master_address: str = None, master_port: int = None): + if self._disaggregation_role != "null": + self._prepend_cu12_lib_to_ld_library_path() + + if self.nnodes > 1: + if self.node_rank != 0: + assert master_address and master_port, "non-master node should provide master address and port" + self._master_address = master_address + self._master_port = master_port + + engine_kwargs = self.config.get("engine_kwargs", {}).get("sglang", {}) or {} + attention_backend = engine_kwargs.pop("attention_backend", None) + quantization = self.config.get("quantization", None) + if quantization is not None: + if quantization == "fp8": + assert version.parse(sglang.__version__) >= version.parse("0.5.5"), ( + "sglang>=0.5.5 is required for FP8 quantization" + ) + FP8_BLOCK_QUANT_KWARGS = { + "activation_scheme": "dynamic", + "fmt": "e4m3", + "quant_method": "fp8", + "weight_block_size": [128, 128], + } + fp8_block_quant_kwargs = dict(FP8_BLOCK_QUANT_KWARGS) + else: + raise ValueError(f"Currently only support fp8 quantization, got: {quantization}") + infer_tp = self.config.tensor_model_parallel_size * self.config.data_parallel_size + args = { + "model_path": self.model_config.local_path, + "dtype": self.config.dtype, + "mem_fraction_static": self.config.gpu_memory_utilization, + "disable_cuda_graph": self.config.enforce_eager, + "enable_memory_saver": True, + "base_gpu_id": self.base_gpu_id, + "gpu_id_step": 1, + "tp_size": infer_tp, + "dp_size": self.config.data_parallel_size, + "ep_size": self.config.expert_parallel_size, + "node_rank": self.node_rank, + "load_format": self.config.load_format, + "nnodes": self.nnodes, + "trust_remote_code": self.model_config.trust_remote_code, + "max_running_requests": self.config.get("max_num_seqs", None), + "log_level": "error", + "mm_attention_backend": "fa3", + "attention_backend": attention_backend if attention_backend is not None else "fa3", + "skip_tokenizer_init": self.config.skip_tokenizer_init, + "skip_server_warmup": True, + "quantization": quantization, + "json_model_override_args": json.dumps({"quantization_config": fp8_block_quant_kwargs}) + if quantization == "fp8" + else json.dumps({}), + **engine_kwargs, + } + + # update lora-related args + if self.model_config.lora_rank > 0: + args.update( + { + "enable_lora": True, + "max_lora_rank": self.model_config.lora_rank, + "lora_target_modules": self.model_config.target_modules, + } + ) + # Only set dist_init_addr for multi-node; for single-node, let SGLang + # handle port selection internally via nccl_port to avoid conflicts. + if self.nnodes > 1: + dist_init_addr = ( + f"[{self._master_address}]:{self._master_port}" + if is_valid_ipv6_address(self._master_address) + else f"{self._master_address}:{self._master_port}" + ) + args["dist_init_addr"] = dist_init_addr + + if self.config.prometheus.enable: + if self.config.prometheus.served_model_name: + # Extract model name from path if it's a full path + served_model_name = self.config.prometheus.served_model_name + if "/" in served_model_name: + # If it's a full path, extract the last part as model name + served_model_name = served_model_name.split("/")[-1] + args["served_model_name"] = served_model_name + + # start sglang metrics + args["enable_metrics"] = True + + # enable_weights_cpu_backup is supported in sglang>=0.5.3 + if "enable_weights_cpu_backup" in [f.name for f in dataclasses.fields(ServerArgs)]: + enable_weights_cpu_backup = ( + True if self.rollout_mode == RolloutMode.COLOCATED or self.model_config.lora_rank > 0 else False + ) + args["enable_weights_cpu_backup"] = enable_weights_cpu_backup + + if self._disaggregation_role != "null": + disagg = self.config.disaggregation + args["disaggregation_mode"] = self._disaggregation_role + args["disaggregation_transfer_backend"] = disagg.transfer_backend + # Bind HTTP + bootstrap to the routable node IP; default 127.0.0.1 + # makes decode-to-prefill bootstrap connection fail across nodes. + args["host"] = self._server_address + if self._disaggregation_bootstrap_port is not None: + args["disaggregation_bootstrap_port"] = self._disaggregation_bootstrap_port + if disagg.decode_tensor_model_parallel_size is not None: + args["disaggregation_decode_tp"] = disagg.decode_tensor_model_parallel_size + if disagg.ib_device is not None: + args["disaggregation_ib_device"] = disagg.ib_device + + if self.config.enable_rollout_routing_replay: + args.update({"enable_return_routed_experts": True}) + + # mtp + if self.config.mtp is not None and self.config.mtp.enable and self.config.mtp.enable_rollout: + # Enable weights CPU backup for sglang >= 0.5.6 + if sglang.__version__ < "0.5.6": + raise ValueError(f"sglang version {sglang.__version__} is not supported for MTP rollout") + + args["speculative_algorithm"] = self.config.mtp.speculative_algorithm + args["speculative_num_steps"] = self.config.mtp.speculative_num_steps + args["speculative_eagle_topk"] = self.config.mtp.speculative_eagle_topk + args["speculative_num_draft_tokens"] = self.config.mtp.speculative_num_draft_tokens + + args["enable_weights_cpu_backup"] = True + args["enable_draft_weights_cpu_backup"] = True + + # NOTE: We can't directly call SGLang's launch_server since it's not an async function. + # https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/entrypoints/http_server.py + sglang.srt.entrypoints.engine._set_envs_and_config = _set_envs_and_config + os.environ["SGLANG_BLOCK_NONZERO_RANK_CHILDREN"] = "0" + server_args = ServerArgs(**args) + # For SGLang main branch or version >= 0.5.10 + # The latest main branch of SGLang has wrapped the _launch_subprocesses function inside the Engine class + if version.parse(sglang.__version__) >= version.parse("0.5.10"): + from sglang.srt.entrypoints.http_server import Engine + + self.tokenizer_manager, self.template_manager, self.scheduler_info, *_ = Engine._launch_subprocesses( + server_args=server_args, + init_tokenizer_manager_func=sglang.srt.entrypoints.engine.init_tokenizer_manager, + run_scheduler_process_func=sglang.srt.entrypoints.engine.run_scheduler_process, + run_detokenizer_process_func=sglang.srt.entrypoints.engine.run_detokenizer_process, + ) + elif version.parse(sglang.__version__) >= version.parse("0.5.7"): + from sglang.srt.entrypoints.http_server import _launch_subprocesses + + self.tokenizer_manager, self.template_manager, self.scheduler_info, *_ = _launch_subprocesses( + server_args=server_args, + init_tokenizer_manager_func=sglang.srt.entrypoints.engine.init_tokenizer_manager, + run_scheduler_process_func=sglang.srt.entrypoints.engine.run_scheduler_process, + run_detokenizer_process_func=sglang.srt.entrypoints.engine.run_detokenizer_process, + ) + else: + from sglang.srt.entrypoints.http_server import _launch_subprocesses + + self.tokenizer_manager, self.template_manager, self.scheduler_info, *_ = _launch_subprocesses( + server_args=server_args + ) + + # In multi-node cases, non-zero rank nodes should not launch http server. + if self.node_rank > 0: + return + + set_global_state( + _GlobalState( + tokenizer_manager=self.tokenizer_manager, + template_manager=self.template_manager, + scheduler_info=self.scheduler_info, + ) + ) + app.is_single_tokenizer_mode = True + + # Set warmup_thread_{kw}args to avoid AttributeError in lifespan function + app.server_args = server_args + app.warmup_thread_kwargs = {"server_args": server_args} + app.warmup_thread_args = (server_args, None, None) + + # Manually add Prometheus middleware before starting server + # This ensures /metrics endpoint is available immediately + if server_args.enable_metrics: + from sglang.srt.utils.common import add_prometheus_middleware + + add_prometheus_middleware(app) + + self._server_port, self._server_task = await run_uvicorn(app, server_args, self._server_address) + self.tokenizer_manager.server_status = ServerStatus.Up + + async def wake_up(self): + if self.node_rank != 0: + return + + if self.rollout_mode == RolloutMode.HYBRID: + # In hybrid mode, rollout is wake up in `update_weights` + raise ValueError(f"wake_up not support rollout_mode {self.rollout_mode}") + elif self.rollout_mode == RolloutMode.COLOCATED: + # Directly call engine to wake up without sync weights. + obj = ResumeMemoryOccupationReqInput(tags=["kv_cache", "weights"]) + await self.tokenizer_manager.resume_memory_occupation(obj, None) + await self.tokenizer_manager.flush_cache() + elif self.rollout_mode == RolloutMode.STANDALONE: + # In standalone mode, resume kv_cache if free_cache_engine is enabled + obj = ResumeMemoryOccupationReqInput(tags=["kv_cache"]) + await self.tokenizer_manager.resume_memory_occupation(obj, None) + await self.tokenizer_manager.flush_cache() + + @property + def lora_as_adapter(self) -> bool: + return ( + self.model_config.lora_rank > 0 or self.model_config.lora.get("rank", 0) > 0 + ) and not self.model_config.lora.get("merge", False) + + async def sleep(self): + if self.node_rank != 0 or not self.config.free_cache_engine: + return + + # When using LoRA as adapter (merge=False), only release kv_cache — + # keep base weights in GPU so we only need to sync adapter deltas. + # Mirrors the vLLM sleep() pattern in vllm_async_server.py. + if self.lora_as_adapter: + tags = ["kv_cache"] + else: + tags = ["kv_cache", "weights"] + + if self.rollout_mode == RolloutMode.HYBRID: + obj = ReleaseMemoryOccupationReqInput(tags=tags) + await self.tokenizer_manager.release_memory_occupation(obj, None) + elif self.rollout_mode == RolloutMode.COLOCATED: + obj = ReleaseMemoryOccupationReqInput(tags=tags) + await self.tokenizer_manager.release_memory_occupation(obj, None) + elif self.rollout_mode == RolloutMode.STANDALONE: + # In standalone mode, resume kv_cache if free_cache_engine is enabled + obj = ReleaseMemoryOccupationReqInput(tags=["kv_cache"]) + await self.tokenizer_manager.release_memory_occupation(obj, None) + + async def clear_kv_cache(self): + if self.node_rank == 0: + await self.tokenizer_manager.flush_cache() + + async def release_kv_cache(self): + """Release only kv_cache GPU memory, keeping model weights intact.""" + if self.node_rank != 0 or not self.config.free_cache_engine: + return + obj = ReleaseMemoryOccupationReqInput(tags=["kv_cache"]) + await self.tokenizer_manager.release_memory_occupation(obj, None) + + async def resume_kv_cache(self): + """Restore kv_cache GPU memory after a weight sync. Counterpart to release_kv_cache().""" + if self.node_rank != 0: + return + obj = ResumeMemoryOccupationReqInput(tags=["kv_cache"]) + await self.tokenizer_manager.resume_memory_occupation(obj, None) + await self.tokenizer_manager.flush_cache() + + async def generate( + self, + prompt_ids: torch.Tensor, + sampling_params: dict[str, Any], + request_id: str, + image_data: Optional[list[Any]] = None, + video_data: Optional[list[Any]] = None, + bootstrap_host: Optional[str] = None, + bootstrap_port: Optional[int] = None, + bootstrap_room: Optional[int] = None, + ) -> TokenOutput: + # PD top-level dispatch: prefill mints a bootstrap_room and fans out + # paired local-prefill + remote-decode calls; decode returns the tokens + # (prefill only materialises KV and pushes via NIXL). Random peer + # choice avoids systematic skew from heavy-tailed RL prompt lengths. + if self._disaggregation_role == "prefill" and self._pd_decode_peers and bootstrap_room is None: + room = secrets.randbits(63) + decode_peer = self._pd_decode_peers[secrets.randbelow(len(self._pd_decode_peers))] + prefill_coro = self.generate( + prompt_ids, + dict(sampling_params), + f"{request_id}_P", + image_data=image_data, + video_data=video_data, + bootstrap_host=self._pd_bootstrap_host, + bootstrap_port=self._disaggregation_bootstrap_port, + bootstrap_room=room, + ) + decode_coro = decode_peer.generate.remote( + prompt_ids, + dict(sampling_params), + f"{request_id}_D", + image_data=image_data, + video_data=video_data, + bootstrap_host=self._pd_bootstrap_host, + bootstrap_port=self._disaggregation_bootstrap_port, + bootstrap_room=room, + ) + _, decode_output = await asyncio.gather(prefill_coro, decode_coro) + return decode_output + + # TODO(@wuxibin): switch to `/generate` http endpoint once multi-modal support ready. + max_possible_tokens = self.config.max_model_len - len(prompt_ids) - 1 + + if max_possible_tokens < 0: + raise ValueError( + f"Prompt length ({len(prompt_ids)}) exceeds the model's maximum context length " + f"({self.config.max_model_len})." + ) + + if "max_new_tokens" in sampling_params: + max_new_tokens = sampling_params.pop("max_new_tokens") + elif "max_tokens" in sampling_params: + # support vllm-style 'max_tokens' param + max_new_tokens = sampling_params.pop("max_tokens") + else: + # Cap max_tokens by response_length to ensure tensor alignment, + # and by remaining budget to prevent OOM in multi-turn rollouts. + max_new_tokens = min( + self.config.response_length, self.config.prompt_length + self.config.response_length - len(prompt_ids) + ) + + # Clamp max_new_tokens to the valid range [0, max_possible_tokens] + max_new_tokens = max(0, min(max_new_tokens, max_possible_tokens)) + + assert max_new_tokens <= max_possible_tokens, ( + f"max_new_tokens {max_new_tokens} exceeds available context space {max_possible_tokens}" + ) + sampling_params["max_new_tokens"] = max_new_tokens + return_logprob = sampling_params.pop("logprobs", False) + + # vLLM-style "prompt_logprobs=K" from the distillation teacher: request + # input-token logprobs for every position (top-K when K>0, sampled-token + # logprob only when K==0). Translate to SGLang's per-request logprob API. + prompt_logprobs = sampling_params.pop("prompt_logprobs", None) + if prompt_logprobs is not None: + return_logprob = True + + request = { + "rid": request_id, + "input_ids": prompt_ids, + "sampling_params": sampling_params, + "return_logprob": return_logprob, + "image_data": image_data, + # TODO: support video input for sglang + # video_data=video_data, + } + + if prompt_logprobs is not None: + request["logprob_start_len"] = 0 + if prompt_logprobs > 0: + request["top_logprobs_num"] = prompt_logprobs + + if self.config.enable_rollout_routing_replay: + request.update({"return_routed_experts": True}) + + # SGLang's scheduler rejects disagg-mode requests without bootstrap_room. + if bootstrap_room is not None: + request["bootstrap_host"] = bootstrap_host + request["bootstrap_port"] = bootstrap_port + request["bootstrap_room"] = bootstrap_room + + generate_request = GenerateReqInput(**request) + + # Add lora request + if self.model_config.lora_rank > 0: + generate_request.lora_path = SGLANG_LORA_NAME + + output = await self.tokenizer_manager.generate_request(generate_request, None).__anext__() + meta_info = output.get("meta_info", {}) + finish_reason = meta_info.get("finish_reason") + finish_reason = finish_reason["type"] if finish_reason else None + if return_logprob: + token_ids = list(output.get("output_ids", [])) + output_token_logprobs = meta_info.get("output_token_logprobs") or [] + if output_token_logprobs and len(output_token_logprobs) == len(token_ids): + log_probs = [float(log_prob) for log_prob, _, _ in output_token_logprobs] + else: + # SGLang may return mismatched lengths (e.g. max_new_tokens=0 + # produces a phantom logprob entry with empty output_ids), or + # an abort may leave an empty logprob payload. + if len(output_token_logprobs) != len(token_ids): + logger.error( + f"output_token_logprobs length ({len(output_token_logprobs)}) != " + f"output_ids length ({len(token_ids)}) for request {request_id}" + ) + token_ids = [] + log_probs = [] + else: + token_ids = output["output_ids"] + log_probs = None + + routed_experts = None + if self.config.enable_rollout_routing_replay: + if self.config.skip_tokenizer_init: + routed_experts = output.get("meta_info", {}).get("routed_experts", None) + else: + from sglang.srt.layers.moe.routed_experts_capturer import extract_routed_experts_from_meta_info + + hf_config = self.model_config.hf_config + if not hasattr(hf_config, "num_hidden_layers") or not hasattr(hf_config, "num_experts_per_tok"): + raise AttributeError( + "enable_rollout_routing_replay is set, but hf_config is missing " + "'num_hidden_layers' or 'num_experts_per_tok'. This feature requires an MoE model " + "configuration that defines these attributes." + ) + routed_experts = extract_routed_experts_from_meta_info(output).reshape( + -1, hf_config.num_hidden_layers, hf_config.num_experts_per_tok + ) + + extra_fields = {"global_steps": self.global_steps} + if prompt_logprobs is not None: + _extract_prompt_logprobs_sglang( + meta_info=meta_info, + num_prompt_logprobs=prompt_logprobs, + sequence_length=len(prompt_ids), + result_dict=extra_fields, + ) + + return TokenOutput( + token_ids=token_ids, + log_probs=log_probs, + routed_experts=routed_experts, + stop_reason=finish_reason, + extra_fields=extra_fields, + ) + + async def set_global_steps(self, global_steps: int): + """Set the global steps of the model weights.""" + self.global_steps = global_steps + + async def abort_all_requests(self): + if self.node_rank != 0: + return + await self.tokenizer_manager.pause_generation(PauseGenerationReqInput(mode="abort")) + + async def resume_generation(self): + if self.node_rank != 0: + return + await self.tokenizer_manager.continue_generation(ContinueGenerationReqInput()) + + async def start_profile(self, **kwargs): + if ( + self.profiler_controller.check_enable() + and self.profiler_controller.check_this_rank() + and self.profiler_controller.is_discrete_mode() + ): + profile_args = build_sglang_profiler_args( + self.profiler_controller.config, self.profiler_controller.tool_config, self.replica_rank + ) + tokenizer_manager = getattr(self, "tokenizer_manager", None) + if tokenizer_manager is None: + return + await tokenizer_manager.start_profile(**profile_args) + + async def stop_profile(self): + if ( + self.profiler_controller.check_enable() + and self.profiler_controller.check_this_rank() + and self.profiler_controller.is_discrete_mode() + ): + tokenizer_manager = getattr(self, "tokenizer_manager", None) + if tokenizer_manager is None: + return + await tokenizer_manager.stop_profile() + + +class SGLangReplica(RolloutReplica): + def __init__( + self, + replica_rank: int, + config: RolloutConfig, + model_config: HFModelConfig, + gpus_per_node: int = 8, + is_reward_model: bool = False, + is_teacher_model: bool = False, + name_suffix: str = "", + ): + super().__init__( + replica_rank, config, model_config, gpus_per_node, is_reward_model, is_teacher_model, name_suffix + ) + self.server_class = ray.remote(SGLangHttpServer) + + async def launch_servers(self): + """Launch http server in each node.""" + assert len(self.workers) == self.world_size, ( + f"worker number {len(self.workers)} not equal to world size {self.world_size}" + ) + + # get (node_id, CUDA_VISIBLE_DEVICES) of all workers + worker_infos = await asyncio.gather( + *[ + worker.__ray_call__.remote( + lambda self: (ray.get_runtime_context().get_node_id(), os.environ[visible_devices_keyword]) + ) + for worker in self.workers + ] + ) + worker_cuda_visible_devices = [worker_info[1] for worker_info in worker_infos] + worker_node_ids = [worker_info[0] for worker_info in worker_infos] + base_gpu_id = 0 + infer_tp = self.config.tensor_model_parallel_size * self.config.data_parallel_size + replica_world_size = infer_tp * self.config.pipeline_model_parallel_size + if os.environ.get(f"RAY_EXPERIMENTAL_NOSET_{visible_devices_keyword}", None): + logger.warning(f"RAY_EXPERIMENTAL_NOSET_{visible_devices_keyword} is set True!") + base_gpu_id = (0 + self.replica_rank * replica_world_size) % self.gpus_per_node + # create server actor in each node with node affinity and cuda visible devices + for node_rank in range(self.nnodes): + workers = self.workers[ + node_rank * self.gpus_per_replica_node : (node_rank + 1) * self.gpus_per_replica_node + ] + node_cuda_visible_devices_set = worker_cuda_visible_devices[ + node_rank * self.gpus_per_replica_node : (node_rank + 1) * self.gpus_per_replica_node + ] + node_cuda_visible_devices = ",".join( + map( + str, + sorted( + set( + int(device) + for worker_devices_set in node_cuda_visible_devices_set + for device in worker_devices_set.split(",") + if device.strip() + ) + ), + ) + ) + + node_id = worker_node_ids[node_rank * self.gpus_per_replica_node] + if self.is_reward_model: + name = f"sglang_server_reward_{self.replica_rank}_{node_rank}{self.name_suffix}" + elif self.is_teacher_model: + name = f"sglang_server_teacher_{self.replica_rank}_{node_rank}{self.name_suffix}" + else: + name = f"sglang_server_{self.replica_rank}_{node_rank}{self.name_suffix}" + server = self.server_class.options( + scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy( + node_id=node_id, + soft=False, + ), + runtime_env={"env_vars": {f"RAY_EXPERIMENTAL_NOSET_{visible_devices_keyword}": "1"}}, + name=name, + max_concurrency=self.max_concurrency, + ).remote( + config=self.config, + model_config=self.model_config, + rollout_mode=self.rollout_mode, + workers=workers, + replica_rank=self.replica_rank, + node_rank=node_rank, + nnodes=self.nnodes, + cuda_visible_devices=node_cuda_visible_devices, + base_gpu_id=base_gpu_id, + ) + self.servers.append(server) + + # launch http server in each node + master_address, master_port = None, None + if self.nnodes > 1: + master_address, master_port = await self.servers[0].get_master_address.remote() + await asyncio.gather( + *[ + server.launch_server.remote(master_address=master_address, master_port=master_port) + for server in self.servers + ] + ) + + # get http server address from first server + server_address, server_port = await self.servers[0].get_server_address.remote() + self._server_handle = self.servers[0] + self._server_address = ( + f"[{server_address}]:{server_port}" + if is_valid_ipv6_address(server_address) + else f"{server_address}:{server_port}" + ) + + async def abort_all_requests(self): + """Abort all ongoing generation requests on the primary server. + + SGLang control RPCs are only served by the node-rank 0 server for a + multi-node replica, so avoid broadcasting this call to every server. + """ + await self.servers[0].abort_all_requests.remote() + + async def resume_generation(self): + """Resume generation on the primary server after abort_all_requests.""" + await self.servers[0].resume_generation.remote() diff --git a/verl/verl/workers/rollout/sglang_rollout/http_server_engine.py b/verl/verl/workers/rollout/sglang_rollout/http_server_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..63034857ce7074e83f57fee0024f5a8b8be996ac --- /dev/null +++ b/verl/verl/workers/rollout/sglang_rollout/http_server_engine.py @@ -0,0 +1,973 @@ +# Copyright 2025 z.ai +# Copyright 2023-2024 SGLang Team +# Copyright 2025 ModelBest Inc. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# This file is adapted from multiple sources: +# 1. THUDM/slime project +# Original source: https://github.com/THUDM/slime/blob/main/slime/backends/sglang_utils/http_server_engine.py +# Copyright 2025 z.ai +# Licensed under the Apache License, Version 2.0 +# 2. SGLang project +# Original source: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/entrypoints/http_server_engine.py +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 +# +# Modifications made by z.ai and ModelBest Inc. include but are not limited to: +# - Enhanced error handling and retry logic +# - Added async support with connection pooling +# - Extended functionality for distributed weight updates +# - Improved logging and monitoring capabilities +# - Additional configuration options and optimizations + +"""HTTP Server Engine Adapter for SGLang. + +This module provides HTTP-based adapters for SGLang engines, allowing communication +with SGLang servers through HTTP requests instead of direct engine calls. + +Classes: + HttpServerAdapter: Synchronous HTTP adapter for SGLang engines + AsyncHttpServerAdapter: Asynchronous HTTP adapter for SGLang engines + +Functions: + launch_server_process: Launch and initialize an SGLang HTTP server process +""" + +import asyncio +import logging +import multiprocessing +import os +import time +from contextlib import asynccontextmanager +from typing import Any, Callable, Optional + +import aiohttp +import requests +from sglang.srt.entrypoints.EngineBase import EngineBase +from sglang.srt.entrypoints.http_server import launch_server +from sglang.srt.managers.io_struct import UpdateWeightsFromTensorReqInput +from sglang.srt.server_args import ServerArgs +from sglang.srt.utils import kill_process_tree + +# Configure logger +logger = logging.getLogger(__name__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + +# Default configuration constants +DEFAULT_TIMEOUT = 60.0 +DEFAULT_MAX_ATTEMPTS = 3 +DEFAULT_RETRY_DELAY = 2.0 +DEFAULT_MAX_CONNECTIONS = 2000 +DEFAULT_MAX_WAIT_TIME = 300.0 + + +def _read_response(response: requests.Response): + if response.status_code == 204 or not response.content: + return {} + try: + return response.json() + except ValueError: + return { + "content_type": response.headers.get("Content-Type", ""), + "text": response.text, + } + + +async def _read_async_response(resp: aiohttp.ClientResponse) -> dict[str, Any]: + if resp.status == 204 or (resp.content_length == 0): + return {} + + try: + return await resp.json(content_type=None) + except Exception: + try: + text = await resp.text() + except Exception: + return {} + return { + "content_type": (resp.headers.get("Content-Type") or ""), + "text": text, + } + + +def launch_server_process( + server_args: ServerArgs, + timeout: float = DEFAULT_TIMEOUT, + max_wait_time=DEFAULT_MAX_WAIT_TIME, + first_rank_in_node=False, +) -> multiprocessing.Process: + """Launch an SGLang HTTP server process and wait for it to be ready. + + This function starts a new process running an SGLang HTTP server, then waits + for the server to become ready by polling its health endpoints. It ensures + the server is fully operational before returning. + + Args: + server_args (ServerArgs): Server configuration arguments including host, port, and other settings + timeout (float, optional): Timeout for individual HTTP requests during health checks. + Defaults to DEFAULT_TIMEOUT. + + Returns: + multiprocessing.Process: The launched multiprocessing.Process instance + + Raises: + RuntimeError: If the server process terminates unexpectedly during startup or cache flush + TimeoutError: If server fails to become ready within reasonable time (300 seconds) + requests.RequestException: If health check requests fail repeatedly + + Note: + This function will return immediately for non-master nodes (node_rank != 0), + but the process will still be started and returned. + This is for consistency; except for the process obtained by node_rank = 0, + other processes have no actual effect. + """ + p = multiprocessing.Process(target=launch_server, args=(server_args,)) + if server_args.node_rank != 0 or not first_rank_in_node: + logger.info(f"Server process started with PID {p.pid} for node rank {server_args.node_rank}", flush=True) + return p + + p.start() + + base_url = server_args.url() + headers = { + "Content-Type": "application/json; charset=utf-8", + "Authorization": f"Bearer {server_args.api_key}", + } + + # Health check with overall timeout + start_time = time.time() + + with requests.Session() as session: + while time.time() - start_time < max_wait_time: + if not p.is_alive(): + raise RuntimeError("Server process terminated unexpectedly during startup") + + try: + if server_args.is_embedding: + response = session.get(f"{base_url}/health", headers=headers, timeout=timeout) + else: + response = session.get(f"{base_url}/health_generate", headers=headers, timeout=timeout) + if response.status_code == 200: + break + except requests.RequestException as e: + logger.debug(f"Health check failed: {e}") + + time.sleep(2) + else: + p.terminate() + logger.error(f"Server in {base_url} failed to become healthy within timeout period") + raise TimeoutError("Server failed to become healthy within timeout period") + + # Ensure cache is ready + while time.time() - start_time < max_wait_time: + if not p.is_alive(): + raise RuntimeError("Server process terminated unexpectedly during cache flush") + + try: + response = session.get(f"{base_url}/flush_cache", headers=headers, timeout=timeout) + if response.status_code == 200: + break + except requests.RequestException as e: + logger.debug(f"Cache flush check failed: {e}") + + time.sleep(2) + else: + p.terminate() + raise TimeoutError("Server cache flush failed within timeout period") + + return p + + +class HttpServerAdapter(EngineBase): + """HTTP-based adapter for SGLang engines. + + This adapter allows interaction with SGLang engines through HTTP requests + instead of direct engine calls. It launches an HTTP server process and + provides methods to communicate with it via REST API calls. + + You can use this class to launch a server from a HttpServerAdapter instance. + We recommend using this class only when you need to use http server. + Otherwise, you can use Engine directly. + + Attributes: + router_ip (Optional[str]): IP address of the router for worker registration + router_port (Optional[int]): Port of the router for worker registration + server_args (ServerArgs): Server configuration arguments + node_rank (int): Rank of this node in distributed setup + process (multiprocessing.Process): The launched server process + timeout (float): HTTP request timeout in seconds + max_attempts (int): Maximum number of attempts for requests + retry_delay (float): Base delay between retries in seconds + """ + + def __init__( + self, + router_ip: Optional[str] = None, + router_port: Optional[int] = None, + timeout: float = DEFAULT_TIMEOUT, + max_attempts: int = DEFAULT_MAX_ATTEMPTS, + retry_delay: float = DEFAULT_RETRY_DELAY, + first_rank_in_node: bool = False, + max_start_wait_time: float = DEFAULT_MAX_WAIT_TIME, + launch_server: bool = True, + **kwargs: Any, + ) -> None: + """Initialize the HTTP server engine adapter. + + Args: + router_ip (Optional[str], optional): IP address of router for worker registration. + Defaults to None. + router_port (Optional[int], optional): Port of router for worker registration. + Defaults to None. + timeout (float, optional): HTTP request timeout in seconds. + Defaults to DEFAULT_TIMEOUT. + max_attempts (int, optional): Maximum number of retry attempts for failed requests. + Defaults to DEFAULT_MAX_ATTEMPTS. + retry_delay (float, optional): Base delay between retries in seconds. + Defaults to DEFAULT_RETRY_DELAY. + launch_server (bool, optional): Whether to launch the server process. + Defaults to True. + **kwargs (Any): Additional arguments passed to ServerArgs + + Note: + TODO: @ChangyiYang Enable SGLang router for this http server engine + If both router_ip and router_port are provided and this is the master node + (node_rank == 0), the adapter will automatically register with the router. + """ + self.router_ip: Optional[str] = router_ip + self.router_port: Optional[int] = router_port + self.timeout: float = timeout + self.max_attempts: int = max_attempts + self.retry_delay: float = retry_delay + self.server_args: ServerArgs = ServerArgs(**kwargs) + self.node_rank: int = self.server_args.node_rank + self.max_start_wait_time: float = max_start_wait_time + + logger.info( + f"Launch HttpServerAdapter at: {self.server_args.host}:{self.server_args.port} with {first_rank_in_node}" + ) + if launch_server: + self.process: multiprocessing.Process = launch_server_process( + self.server_args, self.timeout, self.max_start_wait_time, first_rank_in_node + ) + + if self.node_rank == 0 and self.router_ip and self.router_port: + self._register_with_router() + + def _register_with_router(self) -> None: + """Register worker with router with error handling. + + This method attempts to register the current worker with a router service. + If registration fails, it logs an error but does not raise an exception, + allowing the server to continue operating without router integration. + + Raises: + Does not raise exceptions - all errors are logged and handled gracefully. + """ + try: + url = f"http://{self.router_ip}:{self.router_port}/add_worker" + params = {"url": f"http://{self.server_args.host}:{self.server_args.port}"} + response = requests.post(url, params=params, timeout=self.timeout) + response.raise_for_status() + logger.info("Successfully registered with router") + except Exception as e: + logger.error(f"Failed to register with router: {e}") + # Don't raise here - server can still work without router + + def _make_request( + self, + endpoint: str, + payload: Optional[dict[str, Any]] = None, + method: str = "POST", + timeout: float = DEFAULT_TIMEOUT, + only_master: bool = True, + ) -> dict[str, Any]: + """Make a HTTP request with retry logic and consistent error handling. + + Args: + endpoint (str): The API endpoint to call (without leading slash) + payload (Optional[Dict[str, Any]], optional): The JSON payload to send. + Defaults to empty dict if None. + method (str, optional): HTTP method to use. Defaults to "POST". + + Returns: + Dict[str, Any]: The JSON response from the server + + Raises: + requests.HTTPError: If the HTTP request fails with a client/server error + RuntimeError: If all retry attempts are exhausted + + Note: + - For non-master nodes (node_rank != 0), returns empty dict immediately + - Uses exponential backoff for retries + - Logs warnings for timeout and connection errors, errors for HTTP errors + """ + if only_master and self.node_rank != 0: + return {} + + url = f"http://{self.server_args.host}:{self.server_args.port}/{endpoint}" + + for attempt in range(self.max_attempts): + try: + if method.upper() == "GET": + response = requests.get(url, timeout=self.timeout) + else: + response = requests.post(url, json=payload or {}, timeout=self.timeout) + + response.raise_for_status() + return _read_response(response) + + except requests.exceptions.Timeout: + logger.warning(f"Request to {endpoint} timed out (attempt {attempt + 1})") + except requests.exceptions.ConnectionError: + logger.warning(f"Connection error for {endpoint} (attempt {attempt + 1})") + except requests.exceptions.HTTPError as e: + logger.error(f"HTTP error for {endpoint}: {e}") + raise + except Exception as e: + logger.error(f"Unexpected error for {endpoint}: {e}") + if attempt == self.max_attempts - 1: + raise + + if attempt < self.max_attempts - 1: + time.sleep(self.retry_delay * (2**attempt)) + + raise RuntimeError(f"Failed to complete request to {endpoint} after {self.max_attempts} attempts") + + def update_weights_from_tensor(self, req: UpdateWeightsFromTensorReqInput) -> dict[str, Any]: + """Update model weights from tensor data. + + The HTTP server will only post meta data, and the real weights will be + copied directly from GPUs. + + Args: + serialized_named_tensors (List[str]): List of serialized tensor data + load_format (Optional[str], optional): Format specification for loading weights. + Defaults to None. + flush_cache (bool, optional): Whether to flush cache after updating weights. + Defaults to False. + + Returns: + Dict[str, Any]: Server response containing update status + + Note: + The model should be on GPUs rather than CPU for this functionality to work properly. + If you encounter issues, ensure your model is loaded on GPU devices rather than CPU. + """ + import base64 + + named_tensors = req.serialized_named_tensors + load_format = req.load_format + flush_cache = req.flush_cache + + if named_tensors: + serialized_named_tensors = [ + base64.b64encode(named_tensor).decode("utf-8") for named_tensor in named_tensors + ] + else: + serialized_named_tensors = [] + + return self._make_request( + "update_weights_from_tensor", + { + "serialized_named_tensors": serialized_named_tensors, + "load_format": load_format, + "flush_cache": flush_cache, + }, + ) + + def shutdown(self) -> None: + """Shutdown the HTTP server and clean up resources. + + This method performs the following cleanup operations: + 1. Unregisters the worker from the router (if configured) + 2. Terminates the server process tree + + All operations are performed with error handling to ensure graceful shutdown + even if individual steps fail. + + Note: + This method should be called when the adapter is no longer needed + to ensure proper cleanup of resources and processes. + """ + # Unregister from router + if self.router_ip and self.router_port: + try: + url = f"http://{self.router_ip}:{self.router_port}/remove_worker" + params = {"url": f"http://{self.server_args.host}:{self.server_args.port}"} + requests.post(url, params=params, timeout=5.0) # Short timeout for shutdown + logger.info("Successfully unregistered from router") + except Exception as e: + logger.warning(f"Failed to unregister from router: {e}") + + # Kill server process + if hasattr(self, "process") and self.process is not None: + try: + kill_process_tree(self.process.pid) + logger.info("Server process terminated") + except Exception as e: + logger.error(f"Failed to terminate server process: {e}") + + def generate( + self, + prompt: Optional[str] = None, + sampling_params: Optional[dict[str, Any]] = None, + input_ids: Optional[list[int]] = None, + image_data: Optional[Any] = None, + return_logprob: bool = False, + logprob_start_len: Optional[int] = None, + top_logprobs_num: Optional[int] = None, + token_ids_logprob: Optional[list[int]] = None, + lora_path: Optional[str] = None, + custom_logit_processor: Optional[Callable] = None, + ) -> dict[str, Any]: + """Generate text using the SGLang server. + + Args: + prompt (Optional[str], optional): Text prompt for generation. Defaults to None. + sampling_params (Optional[Dict[str, Any]], optional): Parameters controlling + text generation sampling. Defaults to None. + input_ids (Optional[List[int]], optional): Alternative to prompt, direct token IDs input. + Defaults to None. + image_data (Optional[Any], optional): Image data for multimodal generation. + Defaults to None. + return_logprob (bool, optional): Whether to return log probabilities. + Defaults to False. + logprob_start_len (Optional[int], optional): Starting length for log probability calculation. + Defaults to None. + top_logprobs_num (Optional[int], optional): Number of top log probabilities to return. + Defaults to None. + token_ids_logprob (Optional[List[int]], optional): Specific token IDs for + log probability calculation. Defaults to None. + lora_path (Optional[str], optional): Path to LoRA adapter weights. Defaults to None. + custom_logit_processor (Optional[Callable], optional): Custom logit processing function. + Defaults to None. + + Returns: + Dict[str, Any]: Generated text and associated metadata from the server + + Note: + Either prompt or input_ids should be provided, but not both. + The response format depends on the server configuration and parameters. + """ + payload = { + "text": prompt, + "sampling_params": sampling_params, + "input_ids": input_ids, + "image_data": image_data, + "return_logprob": return_logprob, + "logprob_start_len": logprob_start_len, + "top_logprobs_num": top_logprobs_num, + "token_ids_logprob": token_ids_logprob, + "lora_path": lora_path, + "custom_logit_processor": custom_logit_processor, + } + # Filter out None values + payload = {k: v for k, v in payload.items() if v is not None} + + return self._make_request("generate", payload, only_master=False) + + def reward_score( + self, + prompt: Optional[str] = None, + input_ids: Optional[list[int]] = None, + image_data: Optional[Any] = None, + lora_path: Optional[str] = None, + ) -> dict[str, Any]: + assert self.server_args.is_embedding, "Score is only supported for embedding models" + payload = { + "text": prompt, + "input_ids": input_ids, + "image_data": image_data, + "lora_path": lora_path, + } + # Filter out None values + payload = {k: v for k, v in payload.items() if v is not None} + + return self._make_request("classify", payload, only_master=False) + + def flush_cache(self) -> dict[str, Any]: + """Flush the cache of the server. + + This method repeatedly attempts to flush the server cache until successful. + The flush operation will not return status 200 when there are pending requests. + + Returns: + Dict[str, Any]: Server response indicating cache flush status. + For non-master nodes, returns empty dict. + + Note: + Uses retry logic with limited attempts (max_attempts * 2) to avoid infinite loops. + Each retry includes a delay to allow pending requests to complete. + """ + if self.node_rank != 0: + return {} + + # Use retry logic with limited attempts to avoid infinite loops + for attempt in range(self.max_attempts * 2): # Allow more retries for cache flush + try: + response = requests.get( + f"http://{self.server_args.host}:{self.server_args.port}/flush_cache", timeout=self.timeout + ) + if response.status_code == 200: + return _read_response(response) + except Exception as e: + logger.warning(f"Error flushing cache (attempt {attempt + 1}): {e}") + + time.sleep(self.retry_delay) + + logger.error("Failed to flush cache after maximum attempts") + return {} + + def release_memory_occupation(self, tags: Optional[list[str]] = None) -> dict[str, Any]: + """Release GPU memory occupation temporarily. + + Args: + tags (Optional[List[str]], optional): List of tags to specify which memory to release. + If None, releases all memory. Defaults to None. ["weights", "kv_cache"] + + Returns: + Dict[str, Any]: Server response indicating memory release status + """ + return self._make_request("release_memory_occupation", {"tags": tags}) + + def resume_memory_occupation(self, tags: Optional[list[str]] = None) -> dict[str, Any]: + """Resume GPU memory occupation. + + Args: + tags (Optional[List[str]], optional): List of tags to specify which memory to resume. + If None, resumes all memory. Defaults to None. ["weights", "kv_cache"] + + Returns: + Dict[str, Any]: Server response indicating memory resume status + """ + return self._make_request("resume_memory_occupation", {"tags": tags}) + + def abort_request(self, rid: str = "", abort_all: bool = False) -> dict[str, Any]: + """Abort a request. + + Args: + rid (str): The ID of the request to abort + abort_all (bool, optional): Whether to abort all requests. Defaults to False. + + Returns: + Dict[str, Any]: Server response indicating abort status + """ + return self._make_request("abort_request", {"rid": rid, "abort_all": abort_all}) + + +class AsyncHttpServerAdapter(HttpServerAdapter): + """Asynchronous HTTP-based adapter for SGLang engines. + + This class inherits from HttpServerAdapter and adds async capabilities + for non-blocking HTTP requests to the SGLang server. It provides the same + functionality as the synchronous version but with async/await support. + + The async adapter is useful when you need to make multiple concurrent requests + or integrate with async frameworks. It uses aiohttp for efficient async HTTP + communication and maintains connection pooling for better performance. + + Attributes: + max_connections (int): Maximum number of connections in the connection pool + """ + + def __init__( + self, + router_ip: Optional[str] = None, + router_port: Optional[int] = None, + timeout: float = DEFAULT_TIMEOUT, + max_attempts: int = DEFAULT_MAX_ATTEMPTS, + retry_delay: float = DEFAULT_RETRY_DELAY, + max_connections: int = DEFAULT_MAX_CONNECTIONS, + first_rank_in_node: bool = False, + launch_server: bool = True, + **kwargs: Any, + ) -> None: + """Initialize the async HTTP server engine adapter. + + Args: + router_ip (Optional[str], optional): IP address of router for worker registration. + Defaults to None. + router_port (Optional[int], optional): Port of router for worker registration. + Defaults to None. + timeout (float, optional): HTTP request timeout in seconds. + Defaults to DEFAULT_TIMEOUT. + max_attempts (int, optional): Maximum number of retry attempts for failed requests. + Defaults to DEFAULT_MAX_ATTEMPTS. + retry_delay (float, optional): Base delay between retries in seconds. + Defaults to DEFAULT_RETRY_DELAY. + max_connections (int, optional): Maximum number of connections in the connection pool. + Defaults to DEFAULT_MAX_CONNECTIONS. + launch_server (bool, optional): Whether to launch the server process. + Defaults to True. + **kwargs (Any): Additional arguments passed to ServerArgs + """ + super().__init__( + router_ip, + router_port, + timeout, + max_attempts, + retry_delay, + first_rank_in_node, + launch_server=launch_server, + **kwargs, + ) + self.max_connections: int = max_connections + + @asynccontextmanager + async def _get_session(self) -> aiohttp.ClientSession: + """Context manager for safe session access with proper connection pooling. + + Yields: + aiohttp.ClientSession: Session instance for making HTTP requests + + Note: + This method creates a new session for each request to avoid resource competition + while still maintaining proper connection pooling through the shared connector. + """ + # Create a new session for each request to avoid resource competition + connector = aiohttp.TCPConnector( + limit=self.max_connections, + limit_per_host=self.max_connections // 4, + ttl_dns_cache=300, + use_dns_cache=True, + ) + timeout = aiohttp.ClientTimeout(total=self.timeout) + session = aiohttp.ClientSession(connector=connector, timeout=timeout) + + try: + yield session + finally: + # Always close the session to free up resources + if not session.closed: + await session.close() + + async def _make_async_request( + self, + endpoint: str, + payload: Optional[dict[str, Any]] = None, + method: str = "POST", + timeout: float = DEFAULT_TIMEOUT, + only_master: bool = True, + ) -> dict[str, Any]: + """Make an async HTTP request with retry logic and consistent error handling. + + Args: + endpoint (str): The API endpoint to call (without leading slash) + payload (Optional[Dict[str, Any]], optional): The JSON payload to send. + Defaults to empty dict if None. + method (str, optional): HTTP method to use. Defaults to "POST". + + Returns: + Dict[str, Any]: The JSON response from the server + + Raises: + aiohttp.ClientResponseError: If the HTTP request fails with a client/server error + RuntimeError: If all retry attempts are exhausted + + Note: + - For non-master nodes (node_rank != 0), returns empty dict immediately + - Uses exponential backoff for retries + - Logs warnings for timeout and connection errors, errors for HTTP errors + """ + if only_master and self.node_rank != 0: + return {} + + url = f"http://{self.server_args.host}:{self.server_args.port}/{endpoint}" + + for attempt in range(self.max_attempts): + try: + async with self._get_session() as session: + if method.upper() == "GET": + async with session.get(url, timeout=timeout) as response: + response.raise_for_status() + return await _read_async_response(response) + else: + async with session.post(url, json=payload or {}, timeout=timeout) as response: + response.raise_for_status() + return await _read_async_response(response) + + except asyncio.TimeoutError: + logger.warning(f"Async request to {endpoint} timed out (attempt {attempt + 1})") + except aiohttp.ClientConnectorError: + logger.warning(f"Connection error for {endpoint} (attempt {attempt + 1})") + except aiohttp.ClientResponseError as e: + logger.error(f"HTTP error for {endpoint}: {e}") + raise + except Exception as e: + logger.error(f"Unexpected error for {endpoint}: {e}") + if attempt == self.max_attempts - 1: + raise + + if attempt < self.max_attempts - 1: + await asyncio.sleep(self.retry_delay * (2**attempt)) + + raise RuntimeError(f"Failed to complete async request to {endpoint} after {self.max_attempts} attempts") + + async def release_memory_occupation(self, tags: Optional[list[str]] = None) -> dict[str, Any]: + """Release GPU memory occupation temporarily (async version). + + Args: + tags (Optional[List[str]], optional): List of tags to specify which memory to release. + If None, releases all memory. Defaults to None. ["weights", "kv_cache"] + + Returns: + Dict[str, Any]: Server response indicating memory release status + """ + return await self._make_async_request("release_memory_occupation", {"tags": tags}) + + async def resume_memory_occupation(self, tags: Optional[list[str]] = None) -> dict[str, Any]: + """Resume GPU memory occupation (async version). + + Similar to AsyncEngine, this method handles first-time weight reloading + by calling release_memory_occupation if needed. + + Args: + tags (Optional[List[str]], optional): List of tags to specify which memory to resume. + If None, resumes all memory. Defaults to None. ["weights", "kv_cache"] + + Returns: + Dict[str, Any]: Server response indicating memory resume status + """ + return await self._make_async_request("resume_memory_occupation", {"tags": tags}) + + async def update_weights_from_tensor( + self, + req: UpdateWeightsFromTensorReqInput, + ) -> dict[str, Any]: + """Update model weights from tensor data asynchronously. + + Args: + serialized_named_tensors (List[str]): List of serialized tensor data + load_format (Optional[str], optional): Format specification for loading weights. + Defaults to None. + flush_cache (bool, optional): Whether to flush cache after updating weights. + Defaults to True. + + Returns: + Dict[str, Any]: Server response containing update status + """ + import base64 + + named_tensors = req.serialized_named_tensors + load_format = req.load_format + flush_cache = req.flush_cache + + serialized_named_tensors = [base64.b64encode(named_tensor).decode("utf-8") for named_tensor in named_tensors] + return await self._make_async_request( + "update_weights_from_tensor", + { + "serialized_named_tensors": serialized_named_tensors, + "load_format": load_format, + "flush_cache": flush_cache, + }, + ) + + async def load_lora_adapter_from_tensor(self, req): + return await self._make_async_request( + "load_lora_adapter_from_tensors", + { + "lora_name": req.lora_name, + "config_dict": req.config_dict, + "serialized_tensors": req.serialized_tensors, + }, + ) + + async def unload_lora_adapter(self, lora_name: str): + return await self._make_async_request( + "unload_lora_adapter", + { + "lora_name": lora_name, + }, + ) + + async def available_models(self): + return await self._make_async_request(endpoint="v1/models", method="GET") + + async def flush_cache(self) -> dict[str, Any]: + """Flush the cache of the server asynchronously. + + Similar to the sync version, this method retries until the cache + is successfully flushed. It uses async sleep between retries. + + Returns: + Dict[str, Any]: Server response indicating cache flush status. + For non-master nodes, returns empty dict. + + Note: + Uses retry logic with limited attempts (max_attempts * 4) to avoid infinite loops. + Each retry includes an async delay to allow pending requests to complete. + """ + if self.node_rank != 0: + return {} + + # Use retry logic with limited attempts to avoid infinite loops + for attempt in range(self.max_attempts * 4): # Allow more retries for cache flush + try: + async with self._get_session() as session: + url = f"http://{self.server_args.host}:{self.server_args.port}/flush_cache" + async with session.get(url) as response: + if response.status == 200: + return await _read_async_response(response) + except Exception as e: + logger.warning(f"Error flushing cache (attempt {attempt + 1}): {e}") + + await asyncio.sleep(self.retry_delay) + + logger.error("Failed to flush cache after maximum attempts") + return {} + + async def generate( + self, + prompt: Optional[str] = None, + sampling_params: Optional[dict[str, Any]] = None, + input_ids: Optional[list[int]] = None, + image_data: Optional[Any] = None, + return_logprob: bool = False, + logprob_start_len: Optional[int] = None, + top_logprobs_num: Optional[int] = None, + token_ids_logprob: Optional[list[int]] = None, + lora_path: Optional[str] = None, + custom_logit_processor: Optional[Callable] = None, + ) -> dict[str, Any]: + """Generate text using the SGLang server asynchronously.""" + logger.info("generate() started") + + payload = { + "text": prompt, + "sampling_params": sampling_params, + "input_ids": input_ids, + "image_data": image_data, + "return_logprob": return_logprob, + "logprob_start_len": logprob_start_len, + "top_logprobs_num": top_logprobs_num, + "token_ids_logprob": token_ids_logprob, + "lora_path": lora_path, + "custom_logit_processor": custom_logit_processor, + } + + # Filter out None values + payload = {k: v for k, v in payload.items() if v is not None} + + # Send request + response = await self._make_async_request("generate", payload, timeout=self.timeout, only_master=False) + + return response + + async def async_generate( + self, + prompt: Optional[str] = None, + sampling_params: Optional[dict[str, Any]] = None, + input_ids: Optional[list[int]] = None, + image_data: Optional[Any] = None, + return_logprob: bool = False, + logprob_start_len: Optional[int] = None, + top_logprobs_num: Optional[int] = None, + token_ids_logprob: Optional[list[int]] = None, + lora_path: Optional[str] = None, + custom_logit_processor: Optional[Callable] = None, + ) -> dict[str, Any]: + """Async generate method that mirrors AsyncEngine.async_generate interface. + + This method provides compatibility with AsyncEngine's async_generate method + by forwarding the call to the generate method. It ensures API consistency + between direct engine usage and HTTP-based engine usage. + + Args: + prompt (Optional[str], optional): Text prompt for generation. Defaults to None. + sampling_params (Optional[Dict[str, Any]], optional): Parameters controlling + text generation sampling. Defaults to None. + input_ids (Optional[List[int]], optional): Alternative to prompt, direct token IDs input. + Defaults to None. + image_data (Optional[Any], optional): Image data for multimodal generation. + Defaults to None. + return_logprob (bool, optional): Whether to return log probabilities. + Defaults to False. + logprob_start_len (Optional[int], optional): Starting length for log probability calculation. + Defaults to None. + top_logprobs_num (Optional[int], optional): Number of top log probabilities to return. + Defaults to None. + token_ids_logprob (Optional[List[int]], optional): Specific token IDs for + log probability calculation. Defaults to None. + lora_path (Optional[str], optional): Path to LoRA adapter weights. Defaults to None. + custom_logit_processor (Optional[Callable], optional): Custom logit processing function. + Defaults to None. + + Returns: + Dict[str, Any]: Generated text and associated metadata from the server + + Note: + This method is provided for API compatibility with AsyncEngine. + It forwards all calls to the generate method. + """ + return await self.generate( + prompt=prompt, + sampling_params=sampling_params, + input_ids=input_ids, + image_data=image_data, + return_logprob=return_logprob, + logprob_start_len=logprob_start_len, + top_logprobs_num=top_logprobs_num, + token_ids_logprob=token_ids_logprob, + lora_path=lora_path, + custom_logit_processor=custom_logit_processor, + ) + + async def reward_score( + self, + prompt: Optional[str] = None, + input_ids: Optional[list[int]] = None, + image_data: Optional[Any] = None, + lora_path: Optional[str] = None, + ) -> dict[str, Any]: + logger.info("reward_score() started") + payload = { + "text": prompt, + "input_ids": input_ids, + "image_data": image_data, + "lora_path": lora_path, + } + # Filter out None values + payload = {k: v for k, v in payload.items() if v is not None} + + # Send request + response = await self._make_async_request("classify", payload, timeout=self.timeout, only_master=False) + + return response + + async def async_reward_score( + self, + prompt: Optional[str] = None, + input_ids: Optional[list[int]] = None, + image_data: Optional[Any] = None, + lora_path: Optional[str] = None, + ) -> dict[str, Any]: + return await self.reward_score( + prompt=prompt, + input_ids=input_ids, + image_data=image_data, + lora_path=lora_path, + ) + + async def abort_request(self, rid: str = "", abort_all: bool = False) -> dict[str, Any]: + """Abort a request asynchronously. + + Args: + rid (str): The ID of the request to abort + abort_all (bool, optional): Whether to abort all requests. Defaults to False. + + Returns: + Dict[str, Any]: Server response indicating abort status + """ + return await self._make_async_request("abort_request", {"rid": rid, "abort_all": abort_all}) diff --git a/verl/verl/workers/rollout/sglang_rollout/sglang_pd_replica.py b/verl/verl/workers/rollout/sglang_rollout/sglang_pd_replica.py new file mode 100644 index 0000000000000000000000000000000000000000..cd649c3c97497f494835e8b50265ee2221319571 --- /dev/null +++ b/verl/verl/workers/rollout/sglang_rollout/sglang_pd_replica.py @@ -0,0 +1,229 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""SGLang PD-disaggregated replica: 1 prefill + N decode servers per replica, +asymmetric TP supported. MVP: prefill_replicas=1, whole replica on one node.""" + +import asyncio +import logging +import os +from dataclasses import replace as _dc_replace +from typing import Optional + +import ray +from omegaconf import DictConfig +from ray.actor import ActorHandle + +from verl.utils.device import is_torch_npu_available +from verl.utils.net_utils import get_free_port, is_valid_ipv6_address +from verl.workers.config import RolloutConfig +from verl.workers.rollout.sglang_rollout.async_sglang_server import ( + SGLangReplica, + visible_devices_keyword, +) + +logger = logging.getLogger(__file__) +logger.setLevel(logging.INFO) + + +class SGLangPDReplica(SGLangReplica): + """Replica that runs SGLang in prefill-decode disaggregated mode.""" + + def __init__( + self, + replica_rank: int, + config: RolloutConfig, + model_config: DictConfig, + gpus_per_node: int = 8, + is_reward_model: bool = False, + is_teacher_model: bool = False, + ): + super().__init__( + replica_rank, + config, + model_config, + gpus_per_node, + is_reward_model, + is_teacher_model, + ) + disagg = self.config.disaggregation + assert disagg.enabled, "SGLangPDReplica requires rollout.disaggregation.enabled=True" + + if disagg.prefill_replicas != 1: + raise NotImplementedError(f"prefill_replicas=1 only (got {disagg.prefill_replicas})") + self._n_prefill = disagg.prefill_replicas + self._n_decode = disagg.decode_replicas + + self._prefill_tp = self.config.tensor_model_parallel_size + # Inline decode_tp default: OmegaConf/Ray serialization drops dataclass methods. + self._decode_tp = ( + disagg.decode_tensor_model_parallel_size + if disagg.decode_tensor_model_parallel_size is not None + else self._prefill_tp + ) + + pd_world_size = self._prefill_tp + self._n_decode * self._decode_tp + if pd_world_size > gpus_per_node: + raise NotImplementedError( + f"PD replica needs {pd_world_size} GPUs but gpus_per_node={gpus_per_node}; " + f"use more replicas to span nodes" + ) + + if self.config.data_parallel_size != 1: + raise NotImplementedError(f"data_parallel_size=1 only (got {self.config.data_parallel_size})") + self.world_size = pd_world_size + self.gpus_per_replica_node = min(self.gpus_per_node, self.world_size) + assert self.world_size % self.gpus_per_replica_node == 0 + self.nnodes = self.world_size // self.gpus_per_replica_node + + self._prefill_servers: list[ActorHandle] = [] + self._decode_servers: list[ActorHandle] = [] + self._prefill_server_address: Optional[str] = None + self._decode_server_addresses: list[str] = [] + self._bootstrap_port: Optional[int] = None + + async def launch_servers(self): + assert len(self.workers) == self.world_size + assert not is_torch_npu_available(check_device=False), "PD on NPU not validated" + + worker_infos = await asyncio.gather( + *[ + worker.__ray_call__.remote( + lambda self: ( + ray.get_runtime_context().get_node_id(), + os.environ[visible_devices_keyword], + ) + ) + for worker in self.workers + ] + ) + + # Hold the bootstrap socket open until prefill binds it; closing earlier + # opens a TOCTOU window where another process can grab the port. + bootstrap_port = self.config.disaggregation.bootstrap_port + self._bootstrap_sock = None + if bootstrap_port is None: + prefill_host_ip = ray.util.get_node_ip_address().strip("[]") + bootstrap_port, self._bootstrap_sock = get_free_port(prefill_host_ip, with_alive_sock=True) + self._bootstrap_port = bootstrap_port + + prefill_end = self._prefill_tp + prefill_workers = self.workers[0:prefill_end] + prefill_node_id = worker_infos[0][0] + prefill_devs = self._collect_cuda_devices(worker_infos[0:prefill_end]) + + if self._bootstrap_sock is not None: + self._bootstrap_sock.close() + self._bootstrap_sock = None + + [prefill_server] = await self._launch_one( + role="prefill", + workers=prefill_workers, + node_id=prefill_node_id, + cuda_visible_devices=prefill_devs, + bootstrap_port=self._bootstrap_port, + tp=self._prefill_tp, + actor_name=f"sglang_server_{self.replica_rank}_0", + ) + self._prefill_servers = [prefill_server] + + prefill_address, prefill_port = await prefill_server.get_server_address.remote() + + def _fmt(addr, port): + return f"[{addr}]:{port}" if is_valid_ipv6_address(addr) else f"{addr}:{port}" + + self._prefill_server_address = _fmt(prefill_address, prefill_port) + + self._decode_servers = [] + self._decode_server_addresses = [] + for i in range(self._n_decode): + start = self._prefill_tp + i * self._decode_tp + end = start + self._decode_tp + workers_i = self.workers[start:end] + node_id_i = worker_infos[start][0] + devs_i = self._collect_cuda_devices(worker_infos[start:end]) + + [decode_server] = await self._launch_one( + role="decode", + workers=workers_i, + node_id=node_id_i, + cuda_visible_devices=devs_i, + bootstrap_port=self._bootstrap_port, + tp=self._decode_tp, + actor_name=f"sglang_server_decode_{self.replica_rank}_{i}", + ) + self._decode_servers.append(decode_server) + + d_addr, d_port = await decode_server.get_server_address.remote() + self._decode_server_addresses.append(_fmt(d_addr, d_port)) + + self._server_address = self._prefill_server_address + self._server_handle = prefill_server + self.servers = list(self._prefill_servers) + list(self._decode_servers) + + await prefill_server.set_pd_peer.remote(list(self._decode_servers), prefill_address) + + logger.info( + f"SGLangPDReplica rank={self.replica_rank} launched: " + f"prefill={self._prefill_server_address}, " + f"decodes=[{', '.join(self._decode_server_addresses)}], " + f"bootstrap_port={self._bootstrap_port}" + ) + + @staticmethod + def _collect_cuda_devices(worker_infos) -> str: + devs = set() + for _, dev_str in worker_infos: + for d in dev_str.split(","): + if d.strip(): + devs.add(int(d)) + return ",".join(str(d) for d in sorted(devs)) + + async def _launch_one( + self, + role: str, + workers: list[ActorHandle], + node_id: str, + cuda_visible_devices: str, + bootstrap_port: int, + tp: int, + actor_name: str, + ) -> list[ActorHandle]: + base_gpu_id = 0 + if os.environ.get(f"RAY_EXPERIMENTAL_NOSET_{visible_devices_keyword}", None): + base_gpu_id = (0 + self.replica_rank * self.world_size) % self.gpus_per_node + + pool_config = _dc_replace(self.config, tensor_model_parallel_size=tp) + + server = self.server_class.options( + scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy( + node_id=node_id, soft=False + ), + runtime_env={"env_vars": {f"RAY_EXPERIMENTAL_NOSET_{visible_devices_keyword}": "1"}}, + name=actor_name, + max_concurrency=self.max_concurrency, + ).remote( + config=pool_config, + model_config=self.model_config, + rollout_mode=self.rollout_mode, + workers=workers, + replica_rank=self.replica_rank, + node_rank=0, + nnodes=1, + cuda_visible_devices=cuda_visible_devices, + base_gpu_id=base_gpu_id, + disaggregation_role=role, + disaggregation_bootstrap_port=bootstrap_port, + ) + await server.launch_server.remote(master_address=None, master_port=None) + return [server] diff --git a/verl/verl/workers/rollout/sglang_rollout/sglang_rollout.py b/verl/verl/workers/rollout/sglang_rollout/sglang_rollout.py new file mode 100644 index 0000000000000000000000000000000000000000..b6b959f499daee9196d826811da29e3fec162071 --- /dev/null +++ b/verl/verl/workers/rollout/sglang_rollout/sglang_rollout.py @@ -0,0 +1,372 @@ +# Copyright 2023-2024 SGLang Team +# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import logging +import multiprocessing as mp +import os +from dataclasses import asdict +from typing import Generator + +import ray +import sglang.srt.entrypoints.engine +import torch +from peft import LoraConfig +from sglang.srt.server_args import ServerArgs +from sglang.srt.utils import ( + MultiprocessingSerializer, + assert_pkg_version, + is_cuda, + set_prometheus_multiproc_dir, + set_ulimit, +) +from sglang.srt.weight_sync.utils import _preprocess_tensor_for_update_weights +from sglang.srt.weight_sync.utils import update_weights as sgl_update_weights +from torch.distributed.device_mesh import DeviceMesh, init_device_mesh + +from verl.utils.net_utils import is_valid_ipv6_address +from verl.workers.config import HFModelConfig, RolloutConfig +from verl.workers.rollout.base import BaseRollout +from verl.workers.rollout.sglang_rollout.http_server_engine import AsyncHttpServerAdapter +from verl.workers.rollout.sglang_rollout.utils import ( + SGLANG_LORA_NAME, + get_named_tensor_buckets, +) + +logger = logging.getLogger(__file__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + + +# patch to avoid issue https://github.com/sgl-project/sglang/issues/6723 +def _set_envs_and_config(server_args: ServerArgs): + # Set global environments + os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" + os.environ["NCCL_CUMEM_ENABLE"] = "0" + os.environ["NCCL_NVLS_ENABLE"] = str(int(server_args.enable_nccl_nvls)) + os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1" + os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "4" + os.environ["CUDA_MODULE_LOADING"] = "AUTO" + # Enable faulthandler in subprocesses + os.environ["PYTHONFAULTHANDLER"] = "1" + + # Set prometheus env vars + if server_args.enable_metrics: + set_prometheus_multiproc_dir() + + # Set ulimit + set_ulimit() + + # Check flashinfer version + if server_args.attention_backend == "flashinfer": + assert_pkg_version( + "flashinfer_python", + "0.2.5", + "Please uninstall the old version and reinstall the latest version by following the instructions at https://docs.flashinfer.ai/installation.html.", + ) + if is_cuda(): + assert_pkg_version( + "sgl-kernel", + "0.1.1", + "Please reinstall the latest version with `pip install sgl-kernel --force-reinstall`", + ) + + # Set mp start method + mp.set_start_method("spawn", force=True) + + +sglang.srt.entrypoints.engine._set_envs_and_config = _set_envs_and_config + + +# because chatCompletion is an async method, it makes the whole ray actor be an async actor +# which can not call loop.run_until_complete. So we need to make the engine to be an async class +class ServerAdapter(BaseRollout): + """SGLang server adapter used in native http server mode, serve as http client to request SGLang server + to resume/release/update weights and kv_cache. + + - hybrid mode: reside in each hybrid worker to sync weights between training engine and SGLang server. + - standalone/colocated mode: just a dummy placeholder to occupy the GPU to prevent ray scheduling new GPU actor. + """ + + def __init__( + self, + config: RolloutConfig, + model_config: HFModelConfig, + device_mesh: DeviceMesh, + replica_rank: int = -1, + ): + super().__init__(config, model_config, device_mesh) + if self.config.get("quantization", None) == "fp8": + import sglang + from packaging import version + + assert version.parse(sglang.__version__) >= version.parse("0.5.5"), ( + "sglang>=0.5.5 is required for FP8 quantization" + ) + FP8_BLOCK_QUANT_KWARGS = { + "activation_scheme": "dynamic", + "fmt": "e4m3", + "quant_method": "fp8", + "weight_block_size": [128, 128], + } + fp8_block_quant_kwargs = dict(FP8_BLOCK_QUANT_KWARGS) + self.model_config.hf_config.quantization_config = fp8_block_quant_kwargs + self._engine: AsyncHttpServerAdapter = None + + rank = int(os.environ["RANK"]) + local_world_size = int(os.environ["RAY_LOCAL_WORLD_SIZE"]) + # PD asymmetric layout inflates per-replica footprint; must match + # agent_loop.py:_initialize_llm_servers or trainer-to-replica mapping breaks. + disagg = getattr(self.config, "disaggregation", None) + prefill_tp = self.config.tensor_model_parallel_size + if disagg is not None and getattr(disagg, "enabled", False): + # Inline decode_tp default: OmegaConf/Ray serialization drops dataclass methods. + decode_tp = ( + disagg.decode_tensor_model_parallel_size + if disagg.decode_tensor_model_parallel_size is not None + else prefill_tp + ) + rollout_world_size = ( + (prefill_tp * disagg.prefill_replicas + decode_tp * disagg.decode_replicas) + * self.config.data_parallel_size + * self.config.pipeline_model_parallel_size + ) + else: + rollout_world_size = prefill_tp * self.config.data_parallel_size * self.config.pipeline_model_parallel_size + if replica_rank == -1: + self.replica_rank = rank // rollout_world_size + else: + self.replica_rank = replica_rank + self.rollout_rank = rank % rollout_world_size + self.node_rank = self.rollout_rank // local_world_size + self.local_rank = self.rollout_rank % local_world_size + + # Map each trainer rank to its co-located SGLang server so weight-update + # IPC handles stay on the GPU where they were created. Offset math + # assumes prefill_replicas == 1 (enforced by SGLangPDReplica); if that + # ever lifts, update both this block and SGLangPDReplica.launch_servers. + self._pd_role = None + self._pd_server_index = None + self._pd_tp_local_rank = None + if disagg is not None and getattr(disagg, "enabled", False): + decode_tp = ( + disagg.decode_tensor_model_parallel_size + if disagg.decode_tensor_model_parallel_size is not None + else prefill_tp + ) + # Modulo by single-group footprint so if DP>1 is ever enabled, + # each DP group's ranks resolve to the same role offsets. + footprint = prefill_tp + disagg.decode_replicas * decode_tp + local = self.rollout_rank % footprint + if local < prefill_tp: + self._pd_role = "prefill" + self._pd_server_index = 0 + self._pd_tp_local_rank = local + else: + off = local - prefill_tp + self._pd_role = "decode" + self._pd_server_index = off // decode_tp + self._pd_tp_local_rank = off % decode_tp + self._has_server = (disagg is None or not getattr(disagg, "enabled", False)) or (self._pd_role is not None) + + # sleep_level controls what gets released during sleep/release: + # 2 (default) = release weights + kv_cache (full sleep, merge path) + # 1 = release kv_cache only (keep base weights, adapter path) + # Set by engine_workers.update_weights() when lora.merge=False. + self.sleep_level = 2 + + async def _init_server_adapter(self): + if self._engine is not None: + return + + if not self._has_server: + return + + # device_mesh is needed to gather cuda ipc handle to update weights. + if self.device_mesh is None: + assert torch.distributed.is_initialized(), "torch distributed must be initialized" + infer_tp = self.config.tensor_model_parallel_size * self.config.data_parallel_size + infer_pp = self.config.pipeline_model_parallel_size + infer_world_size = infer_tp * infer_pp + dp = torch.distributed.get_world_size() // infer_world_size + self.device_mesh = init_device_mesh( + "cpu", mesh_shape=(dp, infer_tp, infer_pp), mesh_dim_names=["dp", "infer_tp", "infer_pp"] + ) + + # Only the role's TP-rank-0 builds an adapter; others participate in + # FSDP collectives but skip HTTP dispatch. + if self._pd_role is not None: + if self._pd_tp_local_rank != 0: + return + else: + if self.device_mesh["infer_tp"].get_local_rank() != 0: + return + + if self._pd_role == "prefill": + actor_name = f"sglang_server_{self.replica_rank}_0" + timeout_kwargs = {} + elif self._pd_role == "decode": + actor_name = f"sglang_server_decode_{self.replica_rank}_{self._pd_server_index}" + # Decode init on long-prompt workloads can stall past the default + # (60s × 12); shorter timeout + fewer attempts avoids trainer lockup. + timeout_kwargs = {"timeout": 10.0, "max_attempts": 2} + else: + actor_name = f"sglang_server_{self.replica_rank}_{self.node_rank}" + timeout_kwargs = {} + + self.server_actor = ray.get_actor(actor_name) + server_address, server_port = await self.server_actor.get_server_address.remote() + host = f"[{server_address}]" if is_valid_ipv6_address(server_address) else server_address + logger.info( + f"ServerAdapter {self._pd_role or 'colocated'}: " + f"replica_rank={self.replica_rank}, rollout_rank={self.rollout_rank}, " + f"server={host}:{server_port}, actor={actor_name}" + ) + + self._engine = AsyncHttpServerAdapter( + model_path=self.model_config.local_path, + host=host, + port=server_port, + launch_server=False, + trust_remote_code=self.model_config.trust_remote_code, + **timeout_kwargs, + ) + + def _is_server_tp_leader(self) -> bool: + """True if this rank is TP-rank-0 of its server's group. + + In PD, the role's TP (prefill_tp or decode_tp) may differ from the + config-level TP that device_mesh was built with, so use + _pd_tp_local_rank when PD is active. + """ + if self._pd_role is not None: + return self._pd_tp_local_rank == 0 + return self.device_mesh["infer_tp"].get_local_rank() == 0 + + async def resume(self, tags: list[str]): + """Resume rollout weights or kv cache in GPU memory. + + Args: + tag: weights or kv_cache. + """ + await self._init_server_adapter() + if self._engine is None: + return + if self._is_server_tp_leader() and self.config.free_cache_engine: + await self._engine.resume_memory_occupation(tags=tags) + + async def release(self): + """Release weights and kv cache in GPU memory. + + When sleep_level=1 (LoRA adapter mode), only releases kv_cache + to keep base weights alive across training iterations. + When sleep_level=2 (default/merge mode), releases everything. + """ + await self._init_server_adapter() + if self._engine is None: + return + if self._is_server_tp_leader() and self.config.free_cache_engine: + if self.sleep_level == 1: + tags = ["kv_cache"] + else: + tags = ["kv_cache", "weights"] + await self._engine.release_memory_occupation(tags=tags) + + async def update_weights( + self, weights: Generator[tuple[str, torch.Tensor], None, None], global_steps: int = None, **kwargs + ): + """ + Update model weights using tensor buckets, similar to THUDM/slime's implementation. + + Notes: + - For the best performance of `rebuild_cuda_tensor`, it is recommended to: + 1. Enable `RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES`. + 2. Manually set `CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7` + when using Tensor Parallelism (TP >= 8). + - See reference implementations in SLIME: + - Main logic: https://github.com/THUDM/slime/blob/fb7605cc5fb09af0f9369d37f7192f12bddee577/slime/ray/ppo_actor.py#L452 + - runtime envs: https://github.com/THUDM/slime/blob/fb7605cc5fb09af0f9369d37f7192f12bddee577/slime/ray/ppo_actor.py#L39 + """ + await self._init_server_adapter() + # All ranks MUST iterate the weights generator below — DTensor.full_tensor() + # all_gather's across the FSDP group and skipping deadlocks the others. + # Only HTTP dispatch is gated on self._engine. + + peft_config, base_sync_done = kwargs.get("peft_config", None), kwargs.get("base_sync_done", False) + if peft_config and base_sync_done: + if self.device_mesh["infer_tp"].get_local_rank() == 0: + # unload lora + models_result = await self._engine.available_models() + exists = any(item["id"] == SGLANG_LORA_NAME for item in models_result["data"]) + if exists: + await self._engine.unload_lora_adapter(SGLANG_LORA_NAME) + + # load lora by tensor + serialize_peft_config, serialize_named_tensors = self.wrap_lora_params(peft_config, weights) + from sglang.srt.managers.io_struct import LoadLoRAAdapterFromTensorsReqInput + + req = LoadLoRAAdapterFromTensorsReqInput( + lora_name=SGLANG_LORA_NAME, + config_dict=serialize_peft_config, + serialized_tensors=serialize_named_tensors, + ) + # send http request + await self._engine.load_lora_adapter_from_tensor(req) + else: + update_weights_bucket_bytes = int(self.config.checkpoint_engine.update_weights_bucket_megabytes) << 20 + if self.config.get("quantization", None) == "fp8": + from verl.utils.sglang.sglang_fp8_utils import SGLangFP8QuantizerHelper + + logger.info("Convert bf16 weights to fp8 format before loading") + fp8_quantizer_helper = SGLangFP8QuantizerHelper(self.model_config.hf_config.quantization_config) + weights = fp8_quantizer_helper.quant_weights_by_name( + weights, + dtype=self.model_config.hf_config.dtype, + ) + else: + weights = weights + + async for params_batch in get_named_tensor_buckets(weights, update_weights_bucket_bytes): + await sgl_update_weights( + engine=self._engine, + params_batch=params_batch, + device_mesh_key="infer_tp", + device_mesh=self.device_mesh, + ) + + if self._engine is not None and self._is_server_tp_leader(): + await self._engine.flush_cache() + if global_steps is not None: + await self.server_actor.set_global_steps.remote(global_steps) + + def wrap_lora_params(self, peft_config: LoraConfig, weights: Generator[tuple[str, torch.Tensor]]): + # peft config + peft_config_json = asdict(peft_config) + peft_config_json["task_type"] = peft_config_json["task_type"].value + peft_config_json["peft_type"] = peft_config_json["peft_type"].value + peft_config_json["target_modules"] = list(peft_config_json["target_modules"]) + + # lora weights + processed_weights: dict[str, torch.Tensor] = { + name: _preprocess_tensor_for_update_weights(tensor.detach()) for name, tensor in weights + } + + infer_tp_size = self.device_mesh["infer_tp"].mesh.size()[0] + serialized_named_tensors = [] + for i in range(infer_tp_size): + serialized_tensors = MultiprocessingSerializer.serialize(processed_weights, output_str=True) + serialized_named_tensors.append(serialized_tensors) + + return peft_config_json, serialized_named_tensors diff --git a/verl/verl/workers/rollout/sglang_rollout/utils.py b/verl/verl/workers/rollout/sglang_rollout/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a54a9ed8a7a4d64dcea86f26155bbc81efae1b32 --- /dev/null +++ b/verl/verl/workers/rollout/sglang_rollout/utils.py @@ -0,0 +1,111 @@ +# Copyright 2023-2024 SGLang Team +# Copyright 2025 ModelBest Inc. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pickle +from typing import Any, Iterator, Optional + +import numpy as np +import torch +import torch.distributed as dist + +from verl.utils.device import get_device_name +from verl.workers.rollout.utils import ensure_async_iterator + +SGLANG_LORA_NAME = "verl_actor_lora_name" + + +def broadcast_pyobj( + data: list[Any], + rank: int, + dist_group: Optional[torch.distributed.ProcessGroup] = None, + src: int = 0, + force_cpu_device: bool = False, +): + """from https://github.com/sgl-project/sglang/blob/844e2f227ab0cce6ef818a719170ce37b9eb1e1b/python/sglang/srt/utils.py#L905 + + Broadcast inputs from src rank to all other ranks with torch.dist backend. + The `rank` here refer to the source rank on global process group (regardless + of dist_group argument). + """ + device = torch.device(get_device_name() if not force_cpu_device else "cpu") + + if rank == src: + if len(data) == 0: + tensor_size = torch.tensor([0], dtype=torch.long, device=device) + dist.broadcast(tensor_size, src=src, group=dist_group) + else: + serialized_data = pickle.dumps(data) + size = len(serialized_data) + + tensor_data = torch.ByteTensor(np.frombuffer(serialized_data, dtype=np.uint8)).to(device) + tensor_size = torch.tensor([size], dtype=torch.long, device=device) + + dist.broadcast(tensor_size, src=src, group=dist_group) + dist.broadcast(tensor_data, src=src, group=dist_group) + return data + else: + tensor_size = torch.tensor([0], dtype=torch.long, device=device) + dist.broadcast(tensor_size, src=src, group=dist_group) + size = tensor_size.item() + + if size == 0: + return [] + + tensor_data = torch.empty(size, dtype=torch.uint8, device=device) + dist.broadcast(tensor_data, src=src, group=dist_group) + + serialized_data = bytes(tensor_data.cpu().numpy()) + data = pickle.loads(serialized_data) + return data + + +async def get_named_tensor_buckets( + iterable: Iterator[tuple[str, torch.Tensor]], bucket_bytes: int +) -> Iterator[list[tuple[str, torch.Tensor]]]: + """ + Group tensors into buckets based on a specified size in megabytes. + + Args: + iterable: An iterator of tuples containing tensor names and tensors. + bucket_bytes: The maximum size of each bucket in bytes. + + Yields: + Lists of tuples, where each tuple contains a tensor name and its corresponding tensor. + + Example: + >>> tensors = [('tensor1', torch.randn(1000, 1000)), ('tensor2', torch.randn(2000, 2000))] + >>> for bucket in get_named_tensor_buckets(tensors, bucket_size_mb=10): + ... print(bucket) + [('tensor1', tensor(...)), ('tensor2', tensor(...))] + + """ + if bucket_bytes <= 0: + raise ValueError(f"bucket_bytes must be greater than 0, got {bucket_bytes}") + + current_bucket = [] + current_size = 0 + async for name, tensor in ensure_async_iterator(iterable): + tensor_size = tensor.element_size() * tensor.numel() + if current_size + tensor_size > bucket_bytes: + if current_bucket: + yield current_bucket + current_bucket = [(name, tensor.clone())] + current_size = tensor_size + else: + current_bucket.append((name, tensor.clone())) + current_size += tensor_size + + if current_bucket: + yield current_bucket diff --git a/verl/verl/workers/rollout/tokenizer.py b/verl/verl/workers/rollout/tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..1e1212e50dce4785767cdd52c3dcc6288d08fa02 --- /dev/null +++ b/verl/verl/workers/rollout/tokenizer.py @@ -0,0 +1,163 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +The base tokenizer class, required for any hybrid engine based rollout or inference with vLLM. +""" + +from abc import ABC, abstractmethod + +import numpy as np +import torch + +__all__ = ["HybridEngineBaseTokenizer"] + + +class HybridEngineBaseTokenizer(ABC): + """the tokenizer property and function name should align with HF's to meet vllm requirement""" + + @property + @abstractmethod + def vocab_size(self): + """ + `int`: Size of the base vocabulary (without the added tokens). + """ + pass + + @property + @abstractmethod + def pad_token_id(self): + """ + `Optional[int]`: Id of the padding token in the vocabulary. Returns `None` if the token has not been set. + """ + pass + + @property + @abstractmethod + def eos_token_id(self): + """ + `Optional[int]`: Id of the end of sentence token in the vocabulary. Returns `None` if the token has not been + set. + """ + pass + + @property + @abstractmethod + def all_special_ids(self) -> list[int]: + """ + `List[int]`: List the ids of the special tokens(`''`, `''`, etc.) mapped to class attributes. + """ + pass + + @property + @abstractmethod + def all_special_tokens(self) -> list[str]: + """ + `List[str]`: A list of the unique special tokens (`''`, `''`, ..., etc.). + + Convert tokens of `tokenizers.AddedToken` type to string. + """ + pass + + @abstractmethod + def encode(self, text): + """ + Converts a string to a sequence of ids (integer), using the tokenizer and vocabulary. + + Args: + text (`str`, `List[str]` or `List[int]`): + The first sequence to be encoded. This can be a string, a list of strings (tokenized string using the + `tokenize` method) or a list of integers. + + text_pair (`str`, `List[str]` or `List[int]`, *optional*): + Optional second sequence to be encoded. This can be a string, a list of strings (tokenized string using + the `tokenize` method) or a list of integers. + """ + pass + + @abstractmethod + def decode( + self, + token_ids: int | list[int] | np.ndarray | torch.Tensor, + skip_special_tokens: bool = False, + clean_up_tokenization_spaces: bool = None, + **kwargs, + ) -> str: + """ + Converts a sequence of ids in a string, using the tokenizer and vocabulary with options to remove special + tokens and clean up tokenization spaces. + + Similar to doing `self.convert_tokens_to_string(self.convert_ids_to_tokens(token_ids))`. + + Args: + token_ids (`Union[int, List[int], np.ndarray, torch.Tensor]`): + List of tokenized input ids. Can be obtained using the `__call__` method. + skip_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to remove special tokens in the decoding. + clean_up_tokenization_spaces (`bool`, *optional*): + Whether or not to clean up the tokenization spaces. If `None`, will default to + `self.clean_up_tokenization_spaces`. + kwargs (additional keyword arguments, *optional*): + Will be passed to the underlying model specific decode method. + + Returns: + `str`: The decoded sentence. + """ + pass + + @abstractmethod + def convert_ids_to_tokens(self, ids: int | list[int], skip_special_tokens: bool = False) -> str | list[str]: + """ + Converts a single index or a sequence of indices in a token or a sequence of tokens, using the vocabulary and + added tokens. + + Args: + ids (`int` or `List[int]`): + The token id (or token ids) to convert to tokens. + skip_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to remove special tokens in the decoding. + + Returns: + `str` or `List[str]`: The decoded token(s). + """ + pass + + @abstractmethod + def get_added_vocab(self) -> dict[str, int]: + """ + Returns the added tokens in the vocabulary as a dictionary of token to index. Results might be different from + the fast call because for now we always add the tokens even if they are already in the vocabulary. This is + something we should change. + + Returns: + `Dict[str, int]`: The added tokens. + """ + pass + + @abstractmethod + def convert_tokens_to_string(self, tokens: list[str]) -> str: + """ + Converts a sequence of tokens in a single string. The most simple way to do it is `" ".join(tokens)` but we + often want to remove sub-word tokenization artifacts at the same time. + + Args: + tokens (`List[str]`): The token to join in a string. + + Returns: + `str`: The joined tokens. + """ + pass + + @property + def is_fast(self): + return False diff --git a/verl/verl/workers/rollout/trtllm_rollout/__init__.py b/verl/verl/workers/rollout/trtllm_rollout/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d828409b82e82a2f9aae138e1150608cf891a667 --- /dev/null +++ b/verl/verl/workers/rollout/trtllm_rollout/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/verl/verl/workers/rollout/trtllm_rollout/trtllm_async_rollout.md b/verl/verl/workers/rollout/trtllm_rollout/trtllm_async_rollout.md new file mode 100644 index 0000000000000000000000000000000000000000..ebb1d21516e37b43a84e625cdb62bbd6bdeb6d1c --- /dev/null +++ b/verl/verl/workers/rollout/trtllm_rollout/trtllm_async_rollout.md @@ -0,0 +1,291 @@ +# Running VeRL with TensorRT-LLM Rollout + +We provide initial support for [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM) as an asynchronous rollout engine in VERL's reinforcement learning pipeline. It covers key features such as distributed inference with Ray-based orchestration, dynamic weight updates via IPC (Inter-Process Communication), and efficient GPU memory management for GRPO training. + +TRT-LLM rollout uses hybrid engine colocate mode, where training and inference workers are colocated on the same GPUs. Memory is managed via `resume()`/`release()` APIs to enable GPU sharing between training and inference workloads. + +While the current design factors in multi-node use cases, more extensive multi-node testing and functionality will be delivered in the near future. Current focus is on FSDP and Megatron backend support for Qwen model variants. + +--- + +## 1. Quick Start + + +```bash +# GRPO with FSDP training engine and TP1 +>> INFER_BACKEND=trtllm ROLLOUT_TP=1 bash examples/grpo_trainer/run_qwen3_8b_fsdp.sh +``` + +Note that using the TRT-LLM rollout requires setting the following environment variables before launching the Ray cluster, as included in the above script. + +```bash +# Clean all SLURM/MPI/PMIx env to avoid pmix mismatch error. +for v in $(env | awk -F= '/^(PMI|PMIX|MPI|OMPI|SLURM)_/{print $1}'); do + unset "$v" +done +``` + +## 2. Architecture Design + +### 2.1 High-Level Component Diagram + +```mermaid +%%{init: {'theme':'base', 'themeVariables': { 'fontSize':'18px', 'edgeLabelBackground':'#eeeeee'}}}%% +flowchart TB + space1[" "] + style space1 fill:none,stroke:none + + subgraph VERL["VERL Training Pipeline"] + subgraph Workers["Training Workers"] + Actor["Actor Worker"] + Critic["Critic Worker"] + RefModel["Ref Model Worker"] + end + + Actor -->|Weight Updates
IPC| Rollout["TensorRT-LLM Rollout"] + + subgraph RayCluster["Rollout Workers
(Ray Cluster)"] + space2[" "] + style space2 fill:none,stroke:none + + subgraph AsyncRollout["ServerAdapter
(per DP rank)"] + DPLeader["• DP Leader coordination"] + IPCMgmt["• IPC handle management"] + HTTPAdapter["• HTTP adapter for server communication"] + end + + AsyncRollout -->|HTTP/REST API| HTTPServer + + subgraph HTTPServer["TRTLLMHttpServer
(Ray Actor per Replica)"] + OpenAI["• OpenAI Server wrapper"] + EngMgmt["• AsyncLLM engine management"] + MemMgmt["• Memory management (resume/release)"] + end + + HTTPServer --> AsyncLLM + + subgraph AsyncLLM["TensorRT-LLM
AsyncLLM Engine"] + GPUWorkers["• GPU workers (Tensor Parallel)"] + KVCache["• KV Cache management"] + CUDAGraph["• CUDA Graph optimization"] + end + end + end + + space1 ~~~ VERL + + style VERL fill:#e1f5ff + style RayCluster fill:#fff4e6 + style AsyncRollout fill:#f3e5f5 + style HTTPServer fill:#e8f5e9 + style AsyncLLM fill:#fce4ec +``` + +### 2.2 Agent Loop Architecture + +TRT-LLM rollout follows the same Agent Loop architecture described in the [VERL documentation](https://verl.readthedocs.io/en/latest/advance/agent_loop.html). + +With TensorRT-LLM rollout, the AsyncLLM engine runs in the same process as the TRTLLMHttpServer (Ray actor). The engine spawns Ray workers as ModelRunner through Ray's native orchestration with placement groups. + +AsyncLLM engine communicates with Ray workers through TensorRT-LLM's internal communication layer. When the server receives a request, it directly calls the AsyncLLM engine to generate response_ids. The Ray workers are separate processes from FSDP/Megatron-LM workers but are co-located on the same GPUs in hybrid engine mode. + +The diagram below illustrates TRT-LLM's implementation in hybrid engine mode (Ray Workers and FSDP workers share GPUs): + +```mermaid +flowchart TB + generate[generate] + + generate --> Server + + Server[TRTLLMHttpServer
AsyncLLM Engine] + + Server --> Workers + + subgraph Workers["TRT-LLM group (TP4)"] + direction LR + subgraph W0[ ] + RW0[Ray Worker-0] + F0[FSDP-0] + end + subgraph W1[ ] + RW1[Ray Worker-1] + F1[FSDP-1] + end + subgraph W2[ ] + RW2[Ray Worker-2] + F2[FSDP-2] + end + subgraph W3[ ] + RW3[Ray Worker-3] + F3[FSDP-3] + end + end + + style Server fill:#ffb6c1 + style RW0 fill:#ffffe0 + style RW1 fill:#ffffe0 + style RW2 fill:#ffffe0 + style RW3 fill:#ffffe0 + style F0 fill:#ffb6c1 + style F1 fill:#ffb6c1 + style F2 fill:#ffb6c1 + style F3 fill:#ffb6c1 + style W0 fill:#d3d3d3 + style W1 fill:#d3d3d3 + style W2 fill:#d3d3d3 + style W3 fill:#d3d3d3 + style Workers fill:#f5f5f5 +``` + + +### 2.3 Ray Placement Group Architecture + +1. **Placement APIs & GPU Assignment**: TRT-LLM rollout leverages TRT-LLM's Ray-based APIs (`placement_groups`, `placement_bundle_indices`, `per_worker_gpu_share`) to control GPU placement. Each replica (corresponding to one `TRTLLMHttpServer`) is assigned GPU bundles from placement groups based on its replica rank and TP size. + +2. **Server Placement**: `TRTLLMHttpServer` is pinned to the same node as its first bundle using `NodeAffinitySchedulingStrategy`, ensuring efficient communication between the HTTP server and its Ray workers. + +3. **GPU Sharing**: In hybrid engine mode, training and inference workers share GPUs. Memory is managed via `resume()`/`release()` APIs. The resource pool uses `max_colocate_count=3` internally to support colocation of ActorRollout, RewardModel, and Critic workers. + +4. **Multi-Node Design**: The placement group slicing algorithm supports spanning multiple placement groups for multi-node deployments. **Note**: Formal multi-node testing and functionality will be delivered in subsequent MRs. + +The following diagram shows an example of TP=4 and DP=2. Replica 0 takes bundles 0-3 and Replica 1 takes bundles 4-7 from the same placement group, with each replica managing TP workers across its assigned bundles: + +```mermaid +flowchart TB + subgraph RayCluster["Ray Cluster Resource Pool"] + subgraph PG0["Placement Group 0 (Node 0)"] + B0_0["Bundle 0: GPU 0"] + B0_1["Bundle 1: GPU 1"] + B0_2["Bundle 2: GPU 2"] + B0_3["Bundle 3: GPU 3"] + B0_4["Bundle 4: GPU 4"] + B0_5["Bundle 5: GPU 5"] + B0_6["Bundle 6: GPU 6"] + B0_7["Bundle 7: GPU 7"] + end + + subgraph PG1["Placement Group 1 (Node 1)"] + B1_0["Bundle 0: GPU 0"] + B1_1["Bundle 1: GPU 1"] + B1_2["Bundle 2: GPU 2"] + B1_3["Bundle 3: GPU 3"] + B1_4["Bundle 4: GPU 4"] + B1_5["Bundle 5: GPU 5"] + B1_6["Bundle 6: GPU 6"] + B1_7["Bundle 7: GPU 7"] + end + + PG0 --> Assignment + PG1 --> Assignment + + Assignment["Assigned to TRTLLMReplica"] + + Assignment --> Replica0 + Assignment --> Replica1 + + Replica0["Replica 0
(bundles 0-3 from PG0)
TP=4, DP=2"] + Replica1["Replica 1
(bundles 4-7 from PG0)
TP=4, DP=2"] + end + + style PG0 fill:#e3f2fd + style PG1 fill:#e3f2fd + style Replica0 fill:#c8e6c9 + style Replica1 fill:#c8e6c9 +``` + +--- + +## 3. Core Components + +### 3.1 `TRTLLMHttpServer` + +**Purpose**: Ray actor that wraps TensorRT-LLM's AsyncLLM engine and exposes an OpenAI-compatible HTTP API. + +**Key Responsibilities**: +- Initialize and manage AsyncLLM engine with placement group constraints +- Wrap AsyncLLM with OpenAIServer to expose HTTP endpoints +- Handle HTTP server lifecycle (launch, shutdown) +- Process generation requests with sampling parameters +- Coordinate memory management (wake_up/sleep) for GPU sharing with training workers + + +### 3.2 `TRTLLMReplica` + +**Purpose**: Manages the mapping between replicas and Ray placement groups, orchestrating server deployment. + +**Key Responsibilities**: +- Calculate placement group and bundle index assignments per replica +- Pin TRTLLMHttpServer to specific nodes using NodeAffinitySchedulingStrategy +- Launch and coordinate HTTP servers across distributed nodes +- Validate placement group configurations + + +### 3.3 `ServerAdapter` + +**Purpose**: Rollout worker that handles weight updates, memory management, and generation via HTTP adapter. + +Each DP rank has one leader (the first TP rank within that DP group), and that leader coordinates weight updates to the corresponding TRTLLMHttpServer replica. + +**Key Responsibilities**: +- Act as DP leader for weight synchronization across exclude_dp mesh +- Convert PyTorch tensors to IPC handles for zero-copy weight updates +- Stream weight updates in chunks to avoid memory exhaustion +- Coordinate resume/release operations for memory management +- Initialize HTTP adapter for server communication + + +### 3.4 `AsyncTRTLLMHttpAdapter` + +**Purpose**: HTTP client for communicating with TRTLLMHttpServer. + +**Key Features**: +- Async request handling with retry logic +- Connection pooling for high throughput +- Exponential backoff on failures +- Timeout management + +--- + +## 4. Data Flow Diagrams + +### 4.1 Generation Request Flow + +```mermaid +sequenceDiagram + participant Client as Client/Actor + participant Rollout as ServerAdapter + participant Adapter as AsyncHttpAdapter + participant Server as TRTLLMHttpServer + participant AsyncLLM as AsyncLLM Engine + + Client->>Rollout: generate(prompts) + + rect rgb(240, 248, 255) + Note over Rollout: Init adapter if needed + end + + Rollout->>Adapter: POST /v1/completions
{prompt_ids, sampling_params} + + rect rgb(255, 250, 240) + Note over Adapter: Retry loop with backoff + end + + Adapter->>Server: HTTP POST + + rect rgb(245, 255, 245) + Note over Server: Parse request
Validate params + end + + Server->>AsyncLLM: generate_async() + + rect rgb(255, 245, 245) + Note over AsyncLLM: Schedule to execution queue + Note over AsyncLLM: Run inference (TP workers)
- Forward pass
- Sample tokens
- Update KV cache + end + + AsyncLLM-->>Server: Output (token_ids, log_probs) + + Server-->>Adapter: JSON response + Adapter-->>Rollout: TokenOutput + Rollout-->>Client: Results +``` diff --git a/verl/verl/workers/rollout/trtllm_rollout/trtllm_async_server.py b/verl/verl/workers/rollout/trtllm_rollout/trtllm_async_server.py new file mode 100644 index 0000000000000000000000000000000000000000..c0617616e47346773a13596727df746a37b6218e --- /dev/null +++ b/verl/verl/workers/rollout/trtllm_rollout/trtllm_async_server.py @@ -0,0 +1,622 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import asyncio +import logging +import os +from typing import Any, Optional + +import ray +import torch +from omegaconf import DictConfig +from ray.actor import ActorHandle +from ray.util import placement_group_table +from ray.util.placement_group import PlacementGroup + +from verl.single_controller.ray import SubRayResourcePool +from verl.utils.config import omega_conf_to_dataclass +from verl.utils.net_utils import is_valid_ipv6_address +from verl.utils.profiler import DistProfiler +from verl.workers.config import HFModelConfig, RolloutConfig +from verl.workers.rollout.replica import RolloutMode, RolloutReplica, TokenOutput +from verl.workers.rollout.utils import get_max_position_embeddings, qwen2_5_vl_dedup_image_tokens, run_uvicorn + +logger = logging.getLogger(__file__) +logger.setLevel(logging.INFO) + + +def _resolve_chat_stop_tokens(model_config) -> tuple[int, list[int]]: + """Return (end_id, stop_token_ids) for TorchSampler. + + Both TRTLLM's samplers stops only on end_id. For chat-format prompts the model + naturally ends each assistant turn with a chat-end token (e.g. <|im_end|> + for Qwen, <|eot_id|> for Llama-3) that is *different* from the base-model + eos_token_id. If end_id is set to the base eos the sampler ignores the + chat-end token and the model loops into a second turn, inflating response + lengths until max_tokens is hit. + + For models without a distinct chat-end token the return values are + identical to the current default (end_id = hf_config.eos_token_id). + """ + eos_token_id = model_config.hf_config.eos_token_id + all_stop_ids: list[int] = list(eos_token_id) if isinstance(eos_token_id, list) else [eos_token_id] + + if model_config.generation_config is not None: + gen_eos = model_config.generation_config.eos_token_id + if gen_eos is not None: + for t in gen_eos if isinstance(gen_eos, list) else [gen_eos]: + if t not in all_stop_ids: + all_stop_ids.append(t) + + chat_end_id = None + if model_config.tokenizer is not None: + _chat_stop_strings = ["<|im_end|>", "<|eot_id|>", "<|end_of_turn|>"] + _added_vocab = model_config.tokenizer.get_added_vocab() + for stop_str in _chat_stop_strings: + if stop_str in _added_vocab: + tid = _added_vocab[stop_str] + if tid not in all_stop_ids: + all_stop_ids.append(tid) + if chat_end_id is None: + chat_end_id = tid + + primary_end_id = chat_end_id if chat_end_id is not None else eos_token_id + logger.warning(f"TRT-LLM stop token IDs: {all_stop_ids}, end_id: {primary_end_id}") + return primary_end_id, all_stop_ids + + +@ray.remote +class TRTLLMHttpServer: + """TensorRT LLM HTTP server in single node. + + Args: + config (DictConfig): full config. + model_config (HFModelConfig): model config. + is_reward_model (bool): whether this is a reward model. + rollout_mode (RolloutMode): rollout mode. + workers (list[ActorHandle]): list of rollout workers. + replica_rank (int): replica rank, a replica may contain multiple nodes. + max_colocate_count (int): max colocate count. + pgs (list[PlacementGroup]): placement groups. + bundle_indices (list[list[int]]): bundle indices. + """ + + def __init__( + self, + config: RolloutConfig, + model_config: HFModelConfig, + is_reward_model: bool, + rollout_mode: RolloutMode, + workers: list[ActorHandle], + replica_rank: int, + max_colocate_count: int, + pgs: list[PlacementGroup] = None, + bundle_indices: list[list[int]] = None, + ): + os.environ["TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL"] = "1" + assert torch.cuda.is_available(), "TRTLLM http server should run on GPU node" + + self.config: RolloutConfig = omega_conf_to_dataclass(config) + self.model_config: HFModelConfig = omega_conf_to_dataclass(model_config, dataclass_type=HFModelConfig) + self.is_reward_model = is_reward_model + max_position_embeddings = get_max_position_embeddings(self.model_config.hf_config) + if self.config.max_model_len is None: + self.config.max_model_len = max_position_embeddings + else: + if self.config.max_model_len > max_position_embeddings: + raise ValueError( + f"max_model_len ({self.config.max_model_len}) should be less than or equal to " + f"max_position_embeddings ({max_position_embeddings})" + ) + self.rollout_mode = rollout_mode + self.workers = workers + self.replica_rank = replica_rank + self.max_colocate_count = max_colocate_count + self.pgs = pgs + self.bundle_indices = bundle_indices + # model weights version, set by ServerAdapter when update weights. + self.global_steps = None + # Set when generation is allowed; cleared during weight sync to block new requests. + self._generation_allowed = asyncio.Event() + self._generation_allowed.set() + + self.profiler_controller = self._init_profiler_controller() + + if self.rollout_mode != RolloutMode.HYBRID and self.config.load_format == "dummy": + logger.warning(f"rollout mode is {self.rollout_mode}, load_format is dummy, set to auto") + self.config.load_format = "auto" + + self.is_vlm_model = ( + self.model_config.hf_config is not None and hasattr(self.model_config.hf_config, "vision_config") + ) or hasattr(self.model_config, "vision_config") + + # used for http server + self._server_address = ray.util.get_node_ip_address().strip("[]") + self._server_port = None + + logger.info(f"TRTLLMHttpServer, replica_rank: {self.replica_rank}") + + _end_id, _stop_ids = _resolve_chat_stop_tokens(self.model_config) + + logger.info(f"TRT-LLM resolved end_id={_end_id}, stop_ids={_stop_ids}") + + self._use_torch_sampler = bool(int(os.environ.get("TLLM_USE_TORCHSAMPLER", "0"))) + + if self._use_torch_sampler: + self.sampling_args = { + "detokenize": True, + "end_id": _end_id, + "stop_token_ids": _stop_ids, + "pad_id": self.model_config.hf_config.pad_token_id, + "include_stop_str_in_output": True, + } + else: + self.sampling_args = { + "detokenize": False, + "end_id": -1, + "pad_id": self.model_config.hf_config.pad_token_id, + "stop_token_ids": _stop_ids, + "include_stop_str_in_output": True, + } + logger.info(f"use_torch_sampler={self._use_torch_sampler}, sampling_args={self.sampling_args}") + + def get_server_address(self): + """Get http server address and port.""" + assert self._server_port is not None, "http server is not launched, port is None" + return self._server_address, self._server_port + + async def launch_server(self): + from tensorrt_llm import AsyncLLM + from tensorrt_llm.llmapi import CapacitySchedulerPolicy, CudaGraphConfig, KvCacheConfig, SchedulerConfig + + try: + from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType, SleepConfig + except ImportError: + ExecutorMemoryType = None + SleepConfig = None + from tensorrt_llm.serve import OpenAIServer + + assert self.config.pipeline_model_parallel_size == 1, "pipeline_model_parallel_size > 1 is not supported yet" + + engine_kwargs = self.config.get("engine_kwargs", {}).get("trtllm", {}) or {} + # Pop kv_cache_config from engine_kwargs to merge into KvCacheConfig constructor, + # otherwise **engine_kwargs unpacking in llm_kwargs would overwrite the entire + # KvCacheConfig object, losing free_gpu_memory_fraction and enable_block_reuse. + kv_cache_overrides = engine_kwargs.pop("kv_cache_config", {}) + kv_cache_config = KvCacheConfig( + enable_block_reuse=self.config.enable_prefix_caching, + free_gpu_memory_fraction=self.config.gpu_memory_utilization, + **kv_cache_overrides, + ) + + per_worker_gpu_share = 1.0 / self.max_colocate_count + + quantization = self.config.quantization + if quantization is not None: + if quantization == "fp8": + FP8_BLOCK_QUANT_KWARGS = { + "activation_scheme": "dynamic", + "fmt": "e4m3", + "quant_method": "fp8", + "weight_block_size": [128, 128], + } + engine_kwargs["model_kwargs"] = {"quantization_config": FP8_BLOCK_QUANT_KWARGS} + if self.config.load_format != "dummy": + raise ValueError("FP8 quantization is only supported for dummy load format") + else: + raise ValueError(f"Currently only support fp8 quantization, got: {quantization}") + + llm_kwargs = { + "model": self.model_config.local_path, + "backend": "pytorch", + "dtype": self.config.dtype, + "enable_chunked_prefill": self.config.enable_chunked_prefill, + "skip_tokenizer_init": self.config.skip_tokenizer_init, + "orchestrator_type": "ray", + "kv_cache_config": kv_cache_config, + "max_seq_len": self.config.max_model_len, + "max_batch_size": self.config.max_num_seqs, + "max_num_tokens": self.config.max_num_batched_tokens, + "tensor_parallel_size": self.config.tensor_model_parallel_size, + "pipeline_parallel_size": self.config.pipeline_model_parallel_size, + "moe_expert_parallel_size": self.config.expert_parallel_size, + "moe_tensor_parallel_size": self.config.moe_tensor_parallel_size, + "load_format": self.config.load_format, + "trust_remote_code": self.model_config.trust_remote_code, + "placement_groups": self.pgs, + "placement_bundle_indices": self.bundle_indices, + "per_worker_gpu_share": per_worker_gpu_share, + "sleep_config": SleepConfig( + restore_modes={ + ExecutorMemoryType.MODEL_WEIGHTS_MAIN: "NONE", + ExecutorMemoryType.KV_CACHE: "NONE", + } + ) + if self.config.enable_sleep_mode and SleepConfig is not None + else None, + "allreduce_strategy": "NCCL", + "sampler_type": "TorchSampler" if self._use_torch_sampler else "TRTLLMSampler", + **engine_kwargs, + } + + self_defined_extension = { + "ray_worker_extension_cls": "verl.workers.rollout.trtllm_rollout.trtllm_worker_extension.WorkerExtension", + } + if self.is_vlm_model: + llm_kwargs.update(self_defined_extension) + else: + # TODO: once TRT-LLM WorkerExtension includes wait_for_engine_idle, + # replace with "tensorrt_llm.llmapi.rlhf_utils.WorkerExtension" directly. + llm_kwargs.update( + { + "ray_worker_extension_cls": ( + "verl.workers.rollout.trtllm_rollout.trtllm_worker_extension.RlhfWorkerExtension" + ), + } + ) + + if self.is_reward_model: + llm_kwargs.update( + { + "cuda_graph_config": None, + "disable_overlap_scheduler": True, + } + ) + else: + llm_kwargs.update( + { + "cuda_graph_config": CudaGraphConfig( + enable_padding=True, + batch_sizes=self.config.cudagraph_capture_sizes, + max_batch_size=0 if self.config.cudagraph_capture_sizes else self.config.max_num_seqs, + ), + "scheduler_config": SchedulerConfig( + capacity_scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + ), + } + ) + + self.llm = await AsyncLLM(**llm_kwargs) + import inspect + + init_params = inspect.signature(OpenAIServer.__init__).parameters + if "generator" in init_params: + trtllm_server = OpenAIServer( + generator=self.llm, + model=self.model_config.local_path, + tool_parser=None, + server_role=None, + metadata_server_cfg=None, + ) + else: + trtllm_server = OpenAIServer( + llm=self.llm, + model=self.model_config.local_path, + tool_parser=None, + server_role=None, + metadata_server_cfg=None, + ) + + app = trtllm_server.app + self._server_port, self._server_task = await run_uvicorn(app, None, self._server_address) + + async def generate( + self, + prompt_ids: str | list[int], + sampling_params: dict[str, Any], + request_id: str, + image_data: Optional[list[Any]] = None, + video_data: Optional[list[Any]] = None, + audio_data: Optional[list[Any]] = None, + mm_processor_kwargs: Optional[dict[str, Any]] = None, + ) -> TokenOutput: + from tensorrt_llm.llmapi import SamplingParams + + max_tokens = min( + self.config.response_length, + self.config.prompt_length + self.config.response_length - len(prompt_ids), + ) + max_tokens = max(0, min(max_tokens, self.config.max_model_len - len(prompt_ids))) + sampling_params["max_tokens"] = max_tokens + # TorchSampler: logprobs=0 means sampled-token logprob; TRTLLMSampler: logprobs=1 + _want_logprobs = sampling_params.pop("logprobs", False) + if self._use_torch_sampler: + sampling_params["logprobs"] = 0 if _want_logprobs else None + else: + sampling_params["logprobs"] = 1 if _want_logprobs else None + if sampling_params["top_k"] == -1: + sampling_params["top_k"] = 0 + sampling_params.update(self.sampling_args) + + trt_llm_sampling_params = SamplingParams(**sampling_params) + if audio_data is not None: + raise NotImplementedError("TRT-LLM rollout does not support audio inputs yet.") + + await self._generation_allowed.wait() + if self.is_vlm_model and (image_data or video_data): + deduped_ids = qwen2_5_vl_dedup_image_tokens(prompt_ids, self.model_config.processor) + org_prompt = self.llm.tokenizer.decode(deduped_ids) + input_dict = { + "prompt": org_prompt, + "multi_modal_data": {}, + "mm_processor_kwargs": dict(mm_processor_kwargs or {}), + } + if image_data: + input_dict["multi_modal_data"]["image"] = image_data + if video_data: + input_dict["multi_modal_data"]["video"] = video_data + + outputs = await self.llm.generate_async( + inputs=input_dict, + sampling_params=trt_llm_sampling_params, + ) + else: + outputs = await self.llm.generate_async( + inputs=prompt_ids, + sampling_params=trt_llm_sampling_params, + ) + token_ids = outputs.outputs[0].token_ids + log_probs = None + if outputs.outputs[0].logprobs is not None: + # When logprobs=1, TRT-LLM returns only the sampled token's logprob at each position. + # Extract log_probs before checking finish_reason so cancelled (partial) requests also + # return log_probs for their already-generated tokens. + log_probs = [list(d.values())[0].logprob for d in outputs.outputs[0].logprobs] + if outputs.outputs[0].finish_reason == "cancelled": + return TokenOutput( + token_ids=token_ids, + log_probs=log_probs, + stop_reason="aborted", + extra_fields={"global_steps": self.global_steps}, + ) + return TokenOutput(token_ids=token_ids, log_probs=log_probs, extra_fields={"global_steps": self.global_steps}) + + async def set_global_steps(self, global_steps: int): + """Set the global steps of the model weights.""" + self.global_steps = global_steps + + async def abort_all_requests(self): + """Abort all in-flight requests and block new ones. Call resume_generation() to unblock.""" + self._generation_allowed.clear() + await self.llm.pause_generation() + # TODO: remove once TRT-LLM is upgraded to a version where pause_generation() + # drains internally (https://github.com/NVIDIA/TensorRT-LLM/pull/13784). + await self.llm.collective_rpc("wait_for_engine_idle") + if self.config.enable_prefix_caching: + await self.llm.collective_rpc("reset_prefix_cache") + + async def resume_generation(self): + """Unblock new generation requests after abort_all_requests().""" + await self.llm.resume_generation() + self._generation_allowed.set() + + async def clear_kv_cache(self): + """Invalidate prefix cache entries after weight update.""" + await self.llm.collective_rpc("reset_prefix_cache") + + async def release_kv_cache(self): + """Release only kv_cache GPU memory, keeping model weights intact. + + This is used during weight sync to free GPU memory for new weights. + """ + if not self.config.free_cache_engine: + return + await self.llm.release(tags=["kv_cache"]) + + async def resume_kv_cache(self): + """Restore kv_cache GPU memory after a weight sync. Counterpart to release_kv_cache().""" + await self.llm.resume(tags=["kv_cache"]) + + async def wake_up(self): + from verl.workers.rollout.trtllm_rollout.trtllm_rollout import ServerAdapter + + if self.rollout_mode == RolloutMode.HYBRID: + # In hybrid mode, rollout is wake up in `update_weights` + raise ValueError(f"wake_up not support rollout_mode {self.rollout_mode}") + if self.rollout_mode == RolloutMode.COLOCATED: + await self.llm.resume(tags=ServerAdapter.get_full_tags()) + elif self.rollout_mode == RolloutMode.STANDALONE: + logger.info("skip wake_up in standalone mode") + + async def sleep(self): + from verl.workers.rollout.trtllm_rollout.trtllm_rollout import ServerAdapter + + if not self.config.free_cache_engine: + return + + if self.rollout_mode == RolloutMode.HYBRID: + await self.llm.release(tags=ServerAdapter.get_full_tags()) + elif self.rollout_mode == RolloutMode.COLOCATED: + await self.llm.release(tags=ServerAdapter.get_full_tags()) + elif self.rollout_mode == RolloutMode.STANDALONE: + logger.info("skip sleep in standalone mode") + + async def report_device_ids(self) -> list[str]: + """Report GPU device UUIDs from TRT-LLM workers.""" + return await self.llm.collective_rpc( + "report_device_id", + unique_reply_rank=0, + ) + + async def start_profile(self, **kwargs): + if self.profiler_controller.check_enable() and self.profiler_controller.check_this_rank(): + await self.llm.collective_rpc("start_profile") + + async def stop_profile(self): + if self.profiler_controller.check_enable() and self.profiler_controller.check_this_rank(): + await self.llm.collective_rpc("stop_profile") + + def _init_profiler_controller(self) -> DistProfiler: + profiler_config = self.config.profiler + tool_config = None + if profiler_config is not None: + if profiler_config.tool in ["torch", "npu"]: + tool_config = omega_conf_to_dataclass((profiler_config.tool_config or {}).get(profiler_config.tool)) + elif profiler_config.tool == "nsys": + # nsys config lives in global_tool_config, not tool_config + from verl.utils.profiler.config import NsightToolConfig + + raw = (profiler_config.global_tool_config or {}).get("nsys") + tool_config = omega_conf_to_dataclass(raw) if raw is not None else NsightToolConfig() + elif profiler_config.tool is not None: + logger.warning(f"trtllm rollout: unsupported profiler tool '{profiler_config.tool}', disabling") + profiler_config = None + return DistProfiler(self.replica_rank, config=profiler_config, tool_config=tool_config) + + +class TRTLLMReplica(RolloutReplica): + def __init__( + self, + replica_rank: int, + config: RolloutConfig, + model_config: DictConfig, + gpus_per_node: int = 8, + is_reward_model: bool = False, + is_teacher_model: bool = False, + name_suffix: str = "", + ) -> None: + if is_teacher_model: + raise NotImplementedError("TRTLLMReplica doesn't support teacher model yet.") + super().__init__( + replica_rank, config, model_config, gpus_per_node, is_reward_model, is_teacher_model, name_suffix + ) + self.node_ip = ray.util.get_node_ip_address().strip("[]") + + def rollout_worker_use_gpu(self) -> bool: + return False + + def get_pgs_and_bundle_indices(self) -> tuple[list[PlacementGroup], list[list[int]]]: + """Get placement groups and bundle indices for the replica.""" + + start_pg_index = 0 + local_bundle_index = 0 + + # For SubRayResourcePool, the replica is assigned sub pool specific for this replica. + if isinstance(self.resource_pool, SubRayResourcePool): + assert self.resource_pool.subgroup_world_size == self.world_size, ( + "Subgroup world size must be equal to world size" + ) + local_bundle_index = self.resource_pool.start_bundle_index + # For RayResourcePool, the replica is assigned to entire resource pool. + # We need to find start pg index and local bundle index based on replica rank. + else: + # In standalone mode, init_standalone() creates a per-replica RayResourcePool + # that contains only world_size bundles for this replica. Start at bundle 0. + # In colocated/hybrid mode, the shared pool spans all replicas, so offset by rank. + if self.rollout_mode == RolloutMode.STANDALONE: + local_bundle_index = 0 + else: + local_bundle_index = self.world_size * self.replica_rank + + while ( + start_pg_index < len(self.resource_pool.pgs) + and local_bundle_index >= self.resource_pool.pgs[start_pg_index].bundle_count + ): + local_bundle_index -= self.resource_pool.pgs[start_pg_index].bundle_count + start_pg_index += 1 + assert ( + start_pg_index < len(self.resource_pool.pgs) + and local_bundle_index < self.resource_pool.pgs[start_pg_index].bundle_count + ), "Start pg index or local bundle index out of range" + + # Global Bundle View for Replica x 2 & TP=4: + # ┌───────────────────┬───────────────────┐ + # │ Placement Group 0 │ Placement Group 1 │ + # ├────┬────┬────┬────┼────┬────┬────┬────┤ + # │ 0 │ 1 │ 2 │ 3 │ 0 │ 1 │ 2 │ 3 │ + # └────┴────┴────┴────┴────┴────┴────┴────┘ + # └───────────────┘ └───────────────┘ + # Replica 0 Replica 1 + # (4 GPUs) (4 GPUs) + + left_bundle_count = self.world_size + + pgs = [] + bundle_indices = [] + + for pg in self.resource_pool.pgs[start_pg_index:]: + if left_bundle_count == 0: + break + + left_bundle_count_in_pg = min(left_bundle_count, pg.bundle_count - local_bundle_index) + pg_bundle_indices = [local_bundle_index + idx for idx in range(left_bundle_count_in_pg)] + pgs.append(pg) + bundle_indices.append(pg_bundle_indices) + left_bundle_count -= left_bundle_count_in_pg + local_bundle_index = 0 + + assert left_bundle_count == 0, "all bundle indices should be assigned" + + return pgs, bundle_indices + + async def launch_servers(self): + assert self.resource_pool.pgs is not None, "placement groups are not initialized" + + pgs, bundle_indices = self.get_pgs_and_bundle_indices() + + # Check server process should be launched on the same node as first bundle of first pg. + first_pg_data = placement_group_table(pgs[0]) + node_id = first_pg_data["bundles_to_node_id"][bundle_indices[0][0]] + print(f"TRTLLMReplica: {self.replica_rank}") + print(f"pg node_id: {node_id}") + print(f"pgs: {pgs}") + print(f"bundle_indices: {bundle_indices}") + + # TRTLLMReplica is a 1:1 map from replica to TRTLLMHttpServer. + name = ( + f"trtllm_server_{self.replica_rank}{self.name_suffix}" + if not self.is_reward_model + else f"trtllm_server_reward_{self.replica_rank}{self.name_suffix}" + ) + _server_env_vars = {"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1"} + # Propagate profiling env vars to the Ray actor so that RayExecutor + # (instantiated inside TRTLLMHttpServer) picks them up for inner workers. + for _prof_var in ( + "TLLM_ENABLE_NSYS", + "TLLM_NSYS_OUTPUT_DIR", + "TLLM_USE_TORCHSAMPLER", + ): + if _val := os.environ.get(_prof_var): + _server_env_vars[_prof_var] = _val + server = TRTLLMHttpServer.options( + scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy( + node_id=node_id, + soft=False, + ), + runtime_env={"env_vars": _server_env_vars}, + name=name, + max_concurrency=self.max_concurrency, + ).remote( + config=self.config, + model_config=self.model_config, + is_reward_model=self.is_reward_model, + rollout_mode=self.rollout_mode, + workers=self.workers, + replica_rank=self.replica_rank, + max_colocate_count=self.resource_pool.max_colocate_count, + pgs=pgs, + bundle_indices=bundle_indices, + ) + self.servers.append(server) + + # launch http server in each node + await asyncio.gather(*[server.launch_server.remote() for server in self.servers]) + + # get http server address from first server + server_address, server_port = await self.servers[0].get_server_address.remote() + self._server_handle = self.servers[0] + self._server_address = ( + f"[{server_address}]:{server_port}" + if is_valid_ipv6_address(server_address) + else f"{server_address}:{server_port}" + ) diff --git a/verl/verl/workers/rollout/trtllm_rollout/trtllm_rollout.py b/verl/verl/workers/rollout/trtllm_rollout/trtllm_rollout.py new file mode 100644 index 0000000000000000000000000000000000000000..4d35d020cdd957965568a1cd3a889867ce8667f0 --- /dev/null +++ b/verl/verl/workers/rollout/trtllm_rollout/trtllm_rollout.py @@ -0,0 +1,538 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import asyncio +import base64 +import contextlib +import gc +import logging +import os +import pickle +import threading +from contextlib import asynccontextmanager +from typing import Any, AsyncGenerator, Optional + +import aiohttp +import pynvml +import ray +import torch +import torch.distributed as dist + +try: + from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType +except (ImportError, RuntimeError): + # RuntimeError: FlashInfer's check_cuda_arch() crashes on CPU-only actors + ExecutorMemoryType = None +from torch.distributed.device_mesh import DeviceMesh, init_device_mesh +from torch.multiprocessing.reductions import reduce_tensor + +from verl.utils.device import get_torch_device +from verl.utils.net_utils import is_valid_ipv6_address +from verl.workers.config import HFModelConfig, RolloutConfig +from verl.workers.rollout.base import BaseRollout +from verl.workers.rollout.utils import ensure_async_iterator + +logger = logging.getLogger(__file__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + +# Default configuration constants +DEFAULT_TIMEOUT = 60.0 +DEFAULT_MAX_ATTEMPTS = 3 +DEFAULT_RETRY_DELAY = 2.0 +DEFAULT_MAX_CONNECTIONS = 2000 +DEFAULT_MAX_WAIT_TIME = 300.0 + + +@contextlib.contextmanager +def nvml_context(): + """Context manager for NVML initialization and shutdown. + + Raises: + RuntimeError: If NVML initialization fails + """ + try: + pynvml.nvmlInit() + yield + except pynvml.NVMLError as e: + raise RuntimeError(f"Failed to initialize NVML: {e}") from e + finally: + try: + pynvml.nvmlShutdown() + except pynvml.NVMLError: + pass + + +_NVML_INITIALIZED = False +_NVML_LOCK = threading.Lock() + + +def get_device_uuid(id: str | int) -> str: + """Get the UUID of a CUDA device using NVML.""" + id = int(id) # pynvml expects int; ray.get_gpu_ids() may return str + global _NVML_INITIALIZED + with _NVML_LOCK: + if not _NVML_INITIALIZED: + try: + pynvml.nvmlInit() + _NVML_INITIALIZED = True + except pynvml.NVMLError as e: + raise RuntimeError(f"Failed to initialize NVML: {e}") from e + + # Get the device handle and UUID + try: + handle = pynvml.nvmlDeviceGetHandleByIndex(id) + uuid = pynvml.nvmlDeviceGetUUID(handle) + # Ensure the UUID is returned as a string, not bytes + if isinstance(uuid, bytes): + return uuid.decode("utf-8") + elif isinstance(uuid, str): + return uuid + else: + raise RuntimeError(f"Unexpected UUID type: {type(uuid)} for device {id} (global index: {id})") + except pynvml.NVMLError as e: + raise RuntimeError(f"Failed to get device UUID for device {id} (global index: {id}): {e}") from e + + +async def _read_async_response(resp: aiohttp.ClientResponse) -> dict[str, Any]: + if resp.status == 204 or (resp.content_length == 0): + return {} + + try: + return await resp.json(content_type=None) + except Exception: + try: + text = await resp.text() + except Exception: + return {} + return { + "content_type": (resp.headers.get("Content-Type") or ""), + "text": text, + } + + +class AsyncTRTLLMHttpAdapter: + def __init__( + self, + host: str, + port: int, + timeout: float = DEFAULT_TIMEOUT, + max_attempts: int = DEFAULT_MAX_ATTEMPTS, + retry_delay: float = DEFAULT_RETRY_DELAY, + max_connections: int = DEFAULT_MAX_CONNECTIONS, + ): + self.host = host + self.port = port + self.timeout = timeout + self.max_attempts = max_attempts + self.retry_delay = retry_delay + self.max_connections = max_connections + + @asynccontextmanager + async def _get_session(self) -> aiohttp.ClientSession: + """Context manager for safe session access with proper connection pooling. + + Yields: + aiohttp.ClientSession: Session instance for making HTTP requests + + Note: + This method creates a new session for each request to avoid resource competition + while still maintaining proper connection pooling through the shared connector. + """ + # Create a new session for each request to avoid resource competition + connector = aiohttp.TCPConnector( + limit=self.max_connections, + limit_per_host=self.max_connections // 4, + ttl_dns_cache=300, + use_dns_cache=True, + ) + timeout = aiohttp.ClientTimeout(total=self.timeout) + session = aiohttp.ClientSession(connector=connector, timeout=timeout) + + try: + yield session + finally: + # Always close the session to free up resources + if not session.closed: + await session.close() + + async def _make_async_request( + self, + endpoint: str, + payload: Optional[dict[str, Any]] = None, + timeout: float = DEFAULT_TIMEOUT, + method: str = "POST", + return_status: bool = False, + ) -> dict[str, Any] | int: + """Make an async HTTP request with retry logic and consistent error handling. + + Args: + endpoint (str): The API endpoint to call (without leading slash) + payload (Optional[Dict[str, Any]], optional): The JSON payload to send. + Defaults to empty dict if None. + method (str, optional): HTTP method to use. Defaults to "POST". + + Returns: + Dict[str, Any]: The JSON response from the server + + Raises: + aiohttp.ClientResponseError: If the HTTP request fails with a client/server error + RuntimeError: If all retry attempts are exhausted + + Note: + - Uses exponential backoff for retries + - Logs warnings for timeout and connection errors, errors for HTTP errors + """ + + url = f"http://{self.host}:{self.port}/{endpoint}" + + for attempt in range(self.max_attempts): + try: + async with self._get_session() as session: + if method.upper() == "GET": + async with session.get(url, timeout=timeout) as response: + response.raise_for_status() + return response.status if return_status else await _read_async_response(response) + else: + async with session.post(url, json=payload or {}, timeout=timeout) as response: + response.raise_for_status() + return response.status if return_status else await _read_async_response(response) + + except asyncio.TimeoutError: + logger.warning(f"Async request to {endpoint} timed out (attempt {attempt + 1})") + except aiohttp.ClientConnectorError: + logger.warning(f"Connection error for {endpoint} (attempt {attempt + 1})") + except aiohttp.ClientResponseError as e: + logger.error(f"HTTP error for {endpoint}: {e}") + raise + except Exception as e: + logger.error(f"Unexpected error for {endpoint}: {e}") + if attempt == self.max_attempts - 1: + raise + + if attempt < self.max_attempts - 1: + await asyncio.sleep(self.retry_delay * (2**attempt)) + + raise RuntimeError(f"Failed to complete async request to {endpoint} after {self.max_attempts} attempts") + + async def resume_memory_occupation(self, tags: list[str]): + """Resume GPU memory occupation (async version). + + Similar to AsyncEngine, this method handles first-time weight reloading + by calling release_memory_occupation if needed. + + Args: + tags (Optional[List[str]], optional): List of tags to specify which memory to resume. + If None, resumes all memory. Defaults to None. ["weights", "kv_cache"] + + Returns: + Dict[str, Any]: Server response indicating memory resume status + """ + return await self._make_async_request("resume_memory", {"tags": tags}) + + async def release_memory_occupation(self, tags: list[str]): + """Release GPU memory occupation temporarily (async version). + + Args: + tags (Optional[List[str]], optional): List of tags to specify which memory to release. + If None, releases all memory. Defaults to None. ["weights", "kv_cache"] + + Returns: + Dict[str, Any]: Server response indicating memory release status + """ + return await self._make_async_request("release_memory", {"tags": tags}) + + async def update_weights(self, weights: dict[str, str]): + """Update model weights from tensor data asynchronously. + + Args: + weights: A dictionary that maps the device uuid of the weight handles. + + Returns: + Dict[str, Any]: Server response containing update status + """ + return await self._make_async_request("update_weights", {"weights": weights}) + + +class ServerAdapter(BaseRollout): + # All releasable/resumable weight tags: every ExecutorMemoryType except kv_cache + # (handled separately) and internal tags prefixed with "_". + # Fallback to hard-coded list for trtllm versions that don't export ExecutorMemoryType. + _WEIGHTS_TAGS = ( + [t.value for t in ExecutorMemoryType if t is not ExecutorMemoryType.KV_CACHE and not t.value.startswith("_")] + if ExecutorMemoryType is not None + else [ + "sampler", + "drafter", + "guided_decoder", + "spec_resource_manager", + "model_extra", + "executor_extra", + "model", + "draft_model", + ] + ) + + @staticmethod + def get_full_tags() -> list[str]: + return ServerAdapter._WEIGHTS_TAGS + ["kv_cache"] + + def __init__( + self, config: RolloutConfig, model_config: HFModelConfig, device_mesh: DeviceMesh, replica_rank: int = -1 + ): + if config.get("quantization", None) == "fp8": + FP8_BLOCK_QUANT_KWARGS = { + "activation_scheme": "dynamic", + "fmt": "e4m3", + "quant_method": "fp8", + "weight_block_size": [128, 128], + } + fp8_block_quant_kwargs = dict(FP8_BLOCK_QUANT_KWARGS) + model_config.hf_config.quantization_config = fp8_block_quant_kwargs + super().__init__(config, model_config, device_mesh) + self._adapter = None + self.hybrid_device_mesh = None + self.gpu_id = None + self.is_leader_rank = None + self.replica_rank = None + self.is_dp_rank = None + self._supports_partial_loading = None + + # hybrid mode + if self.device_mesh is not None: + assert device_mesh.mesh_dim_names.index("dp") == 0, "DP dim should always be the first dimension" + + # Clone a new device mesh for CPU backend only (used for internal ranks communication) + device_mesh_kwargs = dict( + mesh_shape=device_mesh.mesh.shape, + mesh_dim_names=device_mesh.mesh_dim_names, + ) + self.hybrid_device_mesh = init_device_mesh("cpu", **device_mesh_kwargs) + + self.hybrid_device_mesh[self.hybrid_device_mesh.mesh_dim_names[1:]]._flatten(mesh_dim_name="exclude_dp") + self.is_leader_rank = self.hybrid_device_mesh["exclude_dp"].get_local_rank() == 0 + logger.info(f"is_dp_leader: {self.is_leader_rank}") + logger.info(f"exclude_dp_rank = {self.hybrid_device_mesh['exclude_dp'].get_local_rank()}") + logger.info(f"exclude_dp_size = {self.hybrid_device_mesh['exclude_dp'].size()}") + self.gpu_id = ray.get_gpu_ids()[0] + self.replica_rank = self.hybrid_device_mesh["dp"].get_local_rank() + assert len(ray.get_gpu_ids()) == 1, "ServerAdapter should run on a single GPU node" + else: + rank = int(os.environ["RANK"]) + self.replica_rank = replica_rank + self.is_leader_rank = rank == 0 + # Required for CUDA IPC handle creation during weight sync for Async RL. + # Reward/ref models skip weight sync so this can be None. + self.gpu_id = ray.get_gpu_ids()[0] + + # Below is required for all modes. + assert self.replica_rank >= 0, "replica_rank is not set" + assert self.is_leader_rank is not None, "is_leader_rank is not set" + + self.node_ip = ray.util.get_node_ip_address().strip("[]") + + async def get_supports_partial_loading(self) -> bool: + """Query and cache whether the model supports partial weight loading.""" + if self._supports_partial_loading is not None: + return self._supports_partial_loading + + await self._init_server_adapter() + try: + self._supports_partial_loading = await self.server_actor.supports_partial_loading.remote() + except Exception as e: + logger.warning(f"Failed to query partial loading support: {e}, defaulting to False") + self._supports_partial_loading = False + + logger.info(f"Model supports partial loading: {self._supports_partial_loading}") + return self._supports_partial_loading + + async def _init_server_adapter(self): + if self._adapter is not None: + return + + # Standalone mode: lazily build the CPU device mesh from the gloo process group + # (initialized by initialize_global_process_group_ray before ServerAdapter construction). + # Reward/ref models that never call resume(), release(), or update_weights() will never build the mesh. + if self.hybrid_device_mesh is None and self.device_mesh is None: + assert dist.is_initialized(), "gloo process group must be initialized before building device mesh" + infer_tp = self.config.tensor_model_parallel_size + infer_pp = getattr(self.config, "pipeline_model_parallel_size", 1) + world_size = dist.get_world_size() + dp = world_size // (infer_tp * infer_pp) + self.hybrid_device_mesh = init_device_mesh( + "cpu", mesh_shape=(dp, infer_tp, infer_pp), mesh_dim_names=["dp", "infer_tp", "infer_pp"] + ) + self.hybrid_device_mesh[self.hybrid_device_mesh.mesh_dim_names[1:]]._flatten(mesh_dim_name="exclude_dp") + self.is_leader_rank = self.hybrid_device_mesh["exclude_dp"].get_local_rank() == 0 + + # Lazy init http server adapter because http server is launched after hybrid engine. + self.server_actor = ray.get_actor(f"trtllm_server_{self.replica_rank}") + server_address, server_port = await self.server_actor.get_server_address.remote() + assert server_address == self.node_ip, f"server address: {server_address} != node_ip: {self.node_ip}" + + logger.debug(f"replica_rank={self.replica_rank}, server address: {server_address}, port: {server_port}") + host = f"[{server_address}]" if is_valid_ipv6_address(server_address) else server_address + self._adapter = AsyncTRTLLMHttpAdapter( + host=host, + port=server_port, + timeout=self.config.server.timeout, + max_attempts=self.config.server.max_attempts, + retry_delay=self.config.server.retry_delay, + max_connections=self.config.server.max_connections, + ) + + async def resume(self, tags: list[str]): + """Resume rollout weights or kv cache in GPU memory. + + Args: + tag: weights or kv_cache. + """ + # Synchronize all ranks before resuming KV cache to ensure non-leader ranks + # have completed actor offloading to CPU, preventing OOM issue. + if "kv_cache" in tags and self.config.free_cache_engine: + group = self.hybrid_device_mesh["exclude_dp"].get_group() if self.hybrid_device_mesh is not None else None + await asyncio.to_thread(dist.barrier, group=group) + if self.is_leader_rank and self.config.free_cache_engine: + if "weights" in tags: + tags = self._WEIGHTS_TAGS + elif "kv_cache" in tags: + tags = ["kv_cache"] + else: + raise ValueError(f"Invalid tag: {tags}") + await self._init_server_adapter() + await self._adapter.resume_memory_occupation(tags=tags) + + async def release(self): + """Release weights and kv cache in GPU memory.""" + if self.is_leader_rank and self.config.free_cache_engine: + await self._init_server_adapter() + tags = self._WEIGHTS_TAGS + ["kv_cache"] + await self._adapter.release_memory_occupation(tags=tags) + + async def update_weights_from_ipc_handles(self, device_handles): + """Update weights from IPC handles.""" + if self.hybrid_device_mesh is not None: + world_size = self.hybrid_device_mesh["exclude_dp"].size() + group = self.hybrid_device_mesh["exclude_dp"].get_group() + else: + world_size = dist.get_world_size() + group = None + + if self.is_leader_rank: + gathered_handles = [None for _ in range(world_size)] + else: + gathered_handles = None + + await asyncio.to_thread( + dist.gather_object, + obj=device_handles, + object_gather_list=gathered_handles, + group_dst=0, + group=group, + ) + + if self.is_leader_rank: + all_handles = {k: v for d in gathered_handles for k, v in d.items()} + await self._adapter.update_weights(all_handles) + + await asyncio.to_thread(dist.barrier, group=group) + + async def update_weights( + self, weights: AsyncGenerator[tuple[str, torch.Tensor], None], global_steps: int = None, **kwargs + ): + """Update the weights of the rollout model. + + Args: + weights: A generator that yields the name of the weight tensor and the tensor itself. + """ + if self.is_leader_rank: + await self._init_server_adapter() + + total_available_bytes = int(self.config.checkpoint_engine.update_weights_bucket_megabytes) * 1024 * 1024 + + if self.config.get("quantization", None) == "fp8": + from verl.utils.trtllm.trtllm_fp8_utils import TRTLLMFP8QuantizerHelper + + fp8_quantizer_helper = TRTLLMFP8QuantizerHelper(self.model_config.hf_config.quantization_config) + weights = fp8_quantizer_helper.quant_weights_by_name( + weights, + dtype=self.model_config.hf_config.dtype, + ) + + try: + device_uuid = get_device_uuid(int(self.gpu_id)) + except Exception as e: + logger.error(f"Failed to get device UUID in update_weights(): {e}") + device_uuid = None + raise e + + cur_available_bytes = total_available_bytes + cur_handles = [] + + async def flush(): + nonlocal cur_available_bytes, cur_handles + if not cur_handles: + return + serialized_device_handles = {device_uuid: base64.b64encode(pickle.dumps(cur_handles)).decode("utf-8")} + await self.update_weights_from_ipc_handles(serialized_device_handles) + cur_available_bytes = total_available_bytes + cur_handles = [] + + # For non-VLM, always use partial loading. For VLM, leader queries and broadcasts to all + # ranks in the DP replica; use get_global_rank(group, 0) since each replica has a different leader. + is_vlm = self.model_config.hf_config is not None and hasattr(self.model_config.hf_config, "vision_config") + if not is_vlm: + supports_partial_loading = True + else: + exclude_dp_group = self.hybrid_device_mesh["exclude_dp"].get_group() + spl_tensor = torch.zeros(1, dtype=torch.int32) + if self.is_leader_rank: + supports_partial_loading = await self.get_supports_partial_loading() + spl_tensor[0] = int(supports_partial_loading) + leader_global_rank = dist.get_global_rank(exclude_dp_group, 0) + await asyncio.to_thread(dist.broadcast, spl_tensor, src=leader_global_rank, group=exclude_dp_group) + supports_partial_loading = bool(spl_tensor.item()) + + async for name, param in ensure_async_iterator(weights): + if supports_partial_loading: + size_in_bytes = param.element_size() * param.numel() + if size_in_bytes > cur_available_bytes: + await flush() + + assert cur_available_bytes >= size_in_bytes, ( + f"cur_available_bytes: {cur_available_bytes:,} size_in_bytes: {size_in_bytes:,} name: {name}" + ) + cur_available_bytes -= size_in_bytes + + # Clone so the IPC handle owns the data to avoid buffer reuse + # before TRT-LLM reads it, which will silently corrupt weights. + handle = reduce_tensor(param.detach().clone()) + cur_handles.append((name, handle)) + + await flush() + + if self.is_leader_rank: + # Finalize update weights + await self._adapter.update_weights(None) + if global_steps is not None: + await self.server_actor.set_global_steps.remote(global_steps) + group = self.hybrid_device_mesh["exclude_dp"].get_group() if self.hybrid_device_mesh is not None else None + await asyncio.to_thread(dist.barrier, group=group) + + del weights + gc.collect() + get_torch_device().empty_cache() + + def _get_attribute(self, name: str): + return getattr(self, name) diff --git a/verl/verl/workers/rollout/trtllm_rollout/trtllm_worker_extension.py b/verl/verl/workers/rollout/trtllm_rollout/trtllm_worker_extension.py new file mode 100644 index 0000000000000000000000000000000000000000..5bcada65746781ff6460236d2f8f8ec0fbae9dcf --- /dev/null +++ b/verl/verl/workers/rollout/trtllm_rollout/trtllm_worker_extension.py @@ -0,0 +1,198 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import base64 +import inspect +from typing import Optional + +# Defer tensorrt_llm imports to avoid FlashInfer's check_cuda_arch() crash +# when this module is loaded on CPU-only Ray actors. The module is normally +# loaded only on GPU workers via string path in trtllm_async_server.py, but +# guard defensively in case of transitive imports. +try: + from tensorrt_llm import serialization + from tensorrt_llm._ray_utils import control_action_decorator + from tensorrt_llm._torch.modules.fused_moe.moe_load_balancer import MoeLoadBalancer + from tensorrt_llm._torch.utils import get_device_uuid + from tensorrt_llm.llmapi.rlhf_utils import WorkerExtension as TrtllmWorkerExtension + from tensorrt_llm.logger import logger +except (ImportError, RuntimeError): + # On CPU actors without CUDA, these imports may fail. + # The class below won't be usable, but the module can be imported safely. + serialization = None + control_action_decorator = lambda f: f # noqa: E731 — identity fallback + MoeLoadBalancer = None + get_device_uuid = None + TrtllmWorkerExtension = object + logger = None + + +class WorkerExtension(TrtllmWorkerExtension): + def __init__(self): + pass + + @control_action_decorator + def supports_partial_loading(self) -> bool: + """Check if the model supports partial weight loading.""" + try: + model = self.engine.model_engine.model + load_weights_args = inspect.getfullargspec(model.load_weights).args + return "allow_partial_loading" in load_weights_args + except Exception as e: + logger.warning(f"Failed to check partial loading support: {e}") + return False + + @control_action_decorator + def update_weights(self, ipc_handles: Optional[dict] = None): + try: + if not hasattr(self.engine.model_engine.model, "first_pre_reload_weights"): + for module in self.engine.model_engine.model.modules(): + if hasattr(module, "pre_reload_weights") and not getattr(module, "_weights_removed", False): + module.pre_reload_weights() + self.engine.model_engine.model.first_pre_reload_weights = True + + if ipc_handles is not None: + logger.info("Update weights from IPC handles") + device_uuid = get_device_uuid(self.device_id) + + if device_uuid not in ipc_handles: + raise ValueError(f"Device UUID {device_uuid} not found in ipc_handles") + + weights = {} + + serialized_handles = ipc_handles[device_uuid] + if isinstance(serialized_handles, str): + # Data is base64-encoded pickled bytes - deserialize it + # using restricted unpickler from tensorrt_llm.serialization + logger.info("Deserializing base64-encoded weight handles") + decoded_data = base64.b64decode(serialized_handles) + # Allow basic builtins and torch tensor reconstruction classes + approved_imports = { + "builtins": [ + "list", + "tuple", + "str", + "int", + "float", + "bool", + "bytes", + "dict", + "NoneType", + "type", + ], + "torch": [ + "Tensor", + "FloatTensor", + "DoubleTensor", + "HalfTensor", + "BFloat16Tensor", + "IntTensor", + "LongTensor", + "ShortTensor", + "CharTensor", + "ByteTensor", + "BoolTensor", + "Size", + "dtype", + "device", + "float32", + "float16", + "bfloat16", + "int32", + "int64", + "int16", + "int8", + "uint8", + "bool", + ], + "torch.multiprocessing.reductions": [ + "rebuild_cuda_tensor", + "rebuild_tensor", + ], + "torch._utils": [ + "_rebuild_tensor_v2", + ], + "torch.storage": [ + "_load_from_bytes", + "_TypedStorage", + "UntypedStorage", + "TypedStorage", + ], + } + all_handles = serialization.loads( + decoded_data, + approved_imports=approved_imports, + ) + + # Verify the result is a list as expected + if not isinstance(all_handles, list): + raise ValueError(f"Deserialized data must be a list, got {type(all_handles).__name__} instead") + else: + # Data is already in the correct format (backward compatibility) + all_handles = serialized_handles + + for param_name, tensor_handle in all_handles: + func, args = tensor_handle + list_args = list(args) + list_args[6] = self.device_id + tensor = func(*list_args) + weights[param_name] = tensor + + logger.info(f"weights key size: {len(weights.keys())}") + + # Check if model supports partial loading and use appropriate strategy + model = self.engine.model_engine.model + load_weights_args = inspect.getfullargspec(model.load_weights).args + supports_partial_loading = "allow_partial_loading" in load_weights_args + + if supports_partial_loading: + self.engine.model_engine.model_loader.reload(model, weights, allow_partial_loading=True) + else: + self.engine.model_engine.model_loader.reload(model, weights, allow_partial_loading=False) + else: + logger.info("Finalize update weights") + for module in self.engine.model_engine.model.modules(): + if hasattr(module, "process_weights_after_loading") and not getattr( + module, "_weights_removed", False + ): + module.process_weights_after_loading() + if hasattr(module, "post_load_weights") and not getattr(module, "_weights_removed", False): + module.post_load_weights() + moe_load_balancer = getattr(self.engine.model_engine, "moe_load_balancer", None) + if isinstance(moe_load_balancer, MoeLoadBalancer): + moe_load_balancer.register_weight_slots_after_to_cuda() + logger.info("moe_load_balancer finalizing model...") + moe_load_balancer.finalize_model() + logger.info("moe_load_balancer finalize model done") + self.engine.reset_prefix_cache() + delattr(self.engine.model_engine.model, "first_pre_reload_weights") + + except Exception as e: + logger.error("Encountered an error in update_weights") + raise e + + def reset_prefix_cache(self) -> None: + """Invalidate the KV cache prefix reuse state after weight updates.""" + self.engine.reset_prefix_cache() + + +# TODO: remove this class and revert the non-VLM path in trtllm_async_server.py +# to use "tensorrt_llm.llmapi.rlhf_utils.WorkerExtension" once verl's TRT-LLM version +# is bumped to include https://github.com/NVIDIA/TensorRT-LLM/pull/13784. +class RlhfWorkerExtension(TrtllmWorkerExtension): + """Minimal extension of TRT-LLM's WorkerExtension for non-VLM RLHF models.""" + + @control_action_decorator + def wait_for_engine_idle(self) -> None: + """Block until the engine has no active or queued requests.""" + pass diff --git a/verl/verl/workers/rollout/utils.py b/verl/verl/workers/rollout/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..28d426dda165ac8597d0f3d452496fffd63bba14 --- /dev/null +++ b/verl/verl/workers/rollout/utils.py @@ -0,0 +1,198 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import asyncio +import logging +import os + +import numpy as np +import ray +import uvicorn +import yaml +from fastapi import FastAPI + +from verl.workers.config.rollout import PrometheusConfig + +logger = logging.getLogger(__file__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + + +def get_max_position_embeddings(hf_config) -> int: + max_len = getattr(hf_config, "max_position_embeddings", None) + if max_len is None: + text_config = getattr(hf_config, "text_config", None) + if text_config is not None: + max_len = getattr(text_config, "max_position_embeddings", None) + + if max_len is None: + raise ValueError("max_position_embeddings not found in HFModelConfig!") + return int(max_len) + + +class _UvicornServerAutoPort(uvicorn.Server): + """Uvicorn Server that reports the system-assigned port when port=0.""" + + def __init__(self, config: uvicorn.Config) -> None: + super().__init__(config) + self.actual_port: int | None = None + self._startup_done: asyncio.Event = asyncio.Event() + + async def startup(self, sockets=None) -> None: + try: + await super().startup(sockets=sockets) + if self.servers and self.config.port == 0: + sock = self.servers[0].sockets[0] + self.actual_port = sock.getsockname()[1] + else: + self.actual_port = self.config.port + finally: + self._startup_done.set() + + async def get_port(self) -> int | None: + await self._startup_done.wait() + return self.actual_port + + +async def run_uvicorn(app: FastAPI, server_args, server_address) -> tuple[int, asyncio.Task]: + app.server_args = server_args + config = uvicorn.Config(app, host=server_address, port=0, log_level="warning") + server = _UvicornServerAutoPort(config) + server_task = asyncio.create_task(server.serve()) + server_port = await server.get_port() + if server_port is None: + # server.startup() failed. await the task to re-raise exception from server.serve() + await server_task + + # Fails on unexpected situation. + raise RuntimeError("Unexpected: HTTP server started without reporting listened port") + logger.info(f"HTTP server started on port {server_port}") + return server_port, server_task + + +async def ensure_async_iterator(iterable): + """Convert an iterable to an async iterator.""" + if hasattr(iterable, "__aiter__"): + async for item in iterable: + yield item + else: + for item in iterable: + yield item + + +def qwen2_5_vl_dedup_image_tokens(prompt_ids: list[int], processor): + """Deduplicate consecutive image tokens in prompt_ids for Qwen2.5-VL, since vLLM will replicate the + <|image_pad|> and <|video_pad|> token by image_data. + For example, + ``` + <|vision_start|><|image_pad|><|image_pad|>...<|image_pad|><|vision_end|> + => + <|vision_start|><|image_pad|><|vision_end|> + ``` + """ + if ( + processor is not None + and hasattr(processor, "image_processor") + and "Qwen2VLImageProcessor" in processor.image_processor.__class__.__name__ + ): + prompt_ids = np.array(prompt_ids) + mask = np.ones(len(prompt_ids), dtype=bool) + is_value = (prompt_ids == processor.image_token_id) | (prompt_ids == processor.video_token_id) + mask[1:] &= ~(is_value[1:] & is_value[:-1]) + return prompt_ids[mask].tolist() + else: + return prompt_ids + + +def update_prometheus_config(config: PrometheusConfig, server_addresses: list[str], rollout_name: str | None = None): + """ + Update Prometheus configuration file with server addresses and reload on first node. + + server_addresses: vllm or sglang server addresses + + rollout_name: name of the rollout backend (e.g., "vllm", "sglang") + """ + + if not server_addresses: + logger.warning("No server addresses available to update Prometheus config") + return + + try: + # Get Prometheus config file path from environment or use default + prometheus_config_json = { + "global": {"scrape_interval": "10s", "evaluation_interval": "10s"}, + "scrape_configs": [ + { + "job_name": "ray", + "file_sd_configs": [{"files": ["/tmp/ray/prom_metrics_service_discovery.json"]}], + }, + {"job_name": "rollout", "static_configs": [{"targets": server_addresses}]}, + ], + } + + # Write configuration file to all nodes + @ray.remote(num_cpus=0) + def write_config_file(config_data, config_path): + os.makedirs(os.path.dirname(config_path), exist_ok=True) + with open(config_path, "w") as f: + yaml.dump(config_data, f, default_flow_style=False, indent=2) + return True + + # Reload Prometheus on all nodes. Only master node should succeed, skip errors on other nodes. + @ray.remote(num_cpus=0) + def reload_prometheus(port): + import socket + import subprocess + + hostname = socket.gethostname() + ip_address = socket.gethostbyname(hostname) + + reload_url = f"http://{ip_address}:{port}/-/reload" + + try: + subprocess.run(["curl", "-X", "POST", reload_url], capture_output=True, text=True, timeout=10) + print(f"Reloading Prometheus on node: {reload_url}") + except Exception: + # Skip errors on non-master nodes + pass + + # Get all available nodes and schedule tasks on each node + nodes = ray.nodes() + alive_nodes = [node for node in nodes if node["Alive"]] + + # Write config files on all nodes + write_tasks = [] + for node in alive_nodes: + node_ip = node["NodeManagerAddress"] + task = write_config_file.options( + resources={"node:" + node_ip: 0.001} # Schedule to specific node + ).remote(prometheus_config_json, config.file) + write_tasks.append(task) + + ray.get(write_tasks) + + server_type = rollout_name.upper() if rollout_name else "rollout" + print(f"Updated Prometheus configuration at {config.file} with {len(server_addresses)} {server_type} servers") + + # Reload Prometheus on all nodes + reload_tasks = [] + for node in alive_nodes: + node_ip = node["NodeManagerAddress"] + task = reload_prometheus.options( + resources={"node:" + node_ip: 0.001} # Schedule to specific node + ).remote(config.port) + reload_tasks.append(task) + + ray.get(reload_tasks) + + except Exception as e: + logger.error(f"Failed to update Prometheus configuration: {e}") diff --git a/verl/verl/workers/rollout/vllm_rollout/__init__.py b/verl/verl/workers/rollout/vllm_rollout/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2ecf113c8394c651655e80fb3780837cd85288c3 --- /dev/null +++ b/verl/verl/workers/rollout/vllm_rollout/__init__.py @@ -0,0 +1,42 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os +from importlib.metadata import PackageNotFoundError, version + +from .vllm_rollout import ServerAdapter # noqa: F401 + + +def get_version(pkg): + try: + return version(pkg) + except PackageNotFoundError: + return None + + +vllm_package_name = "vllm" +vllm_package_version = get_version(vllm_package_name) +if vllm_package_version is None: + raise PackageNotFoundError( + "To use vllm rollout, please ensure the 'vllm' package is properly installed. See " + "https://verl.readthedocs.io/en/latest/start/install.html for more details" + ) + +if "ROCM_PATH" in os.environ: + import re + + match = re.match(r"(\d+\.\d+\.?\d*)", vllm_package_version) + if match: + vllm_package_version = match.group(1) + else: + raise ValueError(f"Warning: Could not parse version format: {vllm_package_version}") diff --git a/verl/verl/workers/rollout/vllm_rollout/bucketed_weight_transfer.py b/verl/verl/workers/rollout/vllm_rollout/bucketed_weight_transfer.py new file mode 100644 index 0000000000000000000000000000000000000000..82256b0beaa6bfac31811192bac3cee90b4f5cc0 --- /dev/null +++ b/verl/verl/workers/rollout/vllm_rollout/bucketed_weight_transfer.py @@ -0,0 +1,333 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Bucketed weight transfer via ZMQ + IPC (or shared memory fallback). + +Not recommended depending on vllm for this file. +""" + +import gc +import logging +import os +from multiprocessing import shared_memory +from typing import Callable, TypedDict + +import torch +import zmq +from torch.multiprocessing.reductions import reduce_tensor + +from verl.utils.device import get_device_id, get_device_name, get_torch_device + +logger = logging.getLogger(__file__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "INFO")) + + +class TensorMetadata(TypedDict): + name: str + shape: torch.Size + dtype: torch.dtype + offset: int + handle: tuple + + +# copy from https://github.com/vllm-project/vllm/blob/main/examples/offline_inference/rlhf_utils.py +def rebuild_ipc(handle: tuple[Callable, tuple], device_id: int | None = None) -> torch.Tensor: + func, args = handle + list_args = list(args) + if device_id is not None: + # the key is to change device id to the current device id + # in case two processes have different CUDA_VISIBLE_DEVICES + list_args[6] = device_id + buffer = func(*list_args) + return buffer + + +def create_shared_memory(size: int, name: str): + """Create shared memory for weight transfer. If already exists, attach to it.""" + try: + shm = shared_memory.SharedMemory(name=name, create=True, size=size) + except FileExistsError: + shm = shared_memory.SharedMemory(name=name) + assert shm.size >= size, f"Stale shm segment '{name}': expected {size} bytes, got {shm.size}" + return shm + + +def rebuild_shared_memory(name: str, size: int, dtype=torch.uint8): + """Rebuild tensor from shared memory.""" + shm = shared_memory.SharedMemory(name=name) + tensor = torch.frombuffer(shm.buf[:size], dtype=dtype) + + return tensor, shm + + +class BucketedWeightSender: + """ + Send model weights via bucketed IPC transfer over ZMQ. + + Packs weight tensors into a fixed-size communication buffer and sends them + in buckets to the receiver. Supports CUDA IPC and shared memory fallback. + + Args: + zmq_handle: ZMQ IPC socket path (e.g., "ipc:///tmp/rl-colocate-zmq-.sock") + bucket_size_mb: Communication buffer size in MB + use_shm: Use shared memory instead of CUDA IPC (for NPU compatibility) + """ + + def __init__( + self, + zmq_handle: str, + bucket_size_mb: int = 512, + use_shm: bool = False, + ): + self.zmq_handle = zmq_handle + self.bucket_size_mb = bucket_size_mb + self.bucket_size = int(bucket_size_mb) << 20 + self.use_shm = use_shm + + self.zmq_context = zmq.Context.instance() + self.socket = None + self.buffer = None + self.shm = None + + async def async_send_weights(self, weights): + """ + Send weights to the receiver. Accepts a sync generator or async iterator. + + Args: + weights: Generator or async iterator yielding (name, tensor) pairs + """ + from verl.workers.rollout.utils import ensure_async_iterator + + try: + self._init_socket() + self._init_buffer() + + # send bucket weights + offset = 0 + bucket_meta: dict[str, TensorMetadata] = {} + # dtype = PrecisionType.to_dtype(self.config.dtype) + async for name, weight in ensure_async_iterator(weights): + # model parameters are in fp32 full precision + # (vermouth1992) we should not force cast weight here because some parameters + # (such as moe gate) have to keep fp32 precision. If a weight is bf16 in the rollout side, + # the rollout should automatically cast on demand. However, this would incur a higher weight + # transfer volume. + # weight = weight.to(dtype, non_blocking=True) + + # fill the tensor bucket + if offset + weight.nbytes > self.bucket_size and len(bucket_meta) > 0: + get_torch_device().synchronize() + self.socket.send_pyobj({"bucket_meta": bucket_meta, "is_last": False}) + self.socket.recv() + bucket_meta = {} + offset = 0 + + if offset + weight.nbytes > self.bucket_size: + assert not self.use_shm, ( + f"Weight {name}({weight.shape}, {weight.dtype}) is too large to fit in the bucket." + f"Please increase rollout.update_weights_bucket_megabytes({self.bucket_size_mb} MB)." + ) + self._direct_send_large_weight(name, weight) + continue + + bucket_meta[name] = { + "name": name, + "shape": weight.shape, + "dtype": weight.dtype, + "offset": offset, + "handle": None, + } + self.buffer[offset : offset + weight.nbytes].copy_(weight.view(-1).view(torch.uint8), non_blocking=True) + offset += weight.nbytes + + # send the last bucket + get_torch_device().synchronize() + self.socket.send_pyobj({"bucket_meta": bucket_meta, "is_last": True}) + self.socket.recv() + finally: + self._cleanup() + + def _init_socket(self): + """Initialize ZMQ REQ socket and bind.""" + if self.zmq_handle.startswith("ipc://"): + ipc_path = self.zmq_handle[len("ipc://") :] + try: + os.remove(ipc_path) + except OSError: + pass + self.socket = self.zmq_context.socket(zmq.REQ) + self.socket.bind(self.zmq_handle) + + def _init_buffer(self): + """build communication buffer""" + buffer, shm = None, None + if not self.use_shm: + buffer = torch.empty(self.bucket_size, dtype=torch.uint8, device=f"{get_device_name()}:{get_device_id()}") + handle = reduce_tensor(buffer) + self.socket.send_pyobj(handle) + else: + import uuid + + # Create unique name for shared memory + shm_name = f"verl_weights_{uuid.uuid4().hex}" + shm = create_shared_memory(self.bucket_size, shm_name) + buffer = torch.frombuffer(shm.buf, dtype=torch.uint8) + + comm_metadata = {"name": shm_name, "size": self.bucket_size} + self.socket.send_pyobj(comm_metadata) + + self.socket.recv() + self.buffer = buffer + self.shm = shm + + def _cleanup(self): + """clean up""" + if self.socket is not None: + self.socket.close() + self.socket = None + if self.zmq_handle.startswith("ipc://"): + ipc_path = self.zmq_handle[len("ipc://") :] + try: + os.remove(ipc_path) + except OSError: + pass + del self.buffer + self.buffer = None + if self.shm is not None: + self.shm.close() + self.shm.unlink() + del self.shm + self.shm = None + gc.collect() + get_torch_device().ipc_collect() + get_torch_device().empty_cache() + + def _direct_send_large_weight(self, name: str, weight: torch.Tensor): + """Send a weight larger than the bucket size via cuda ipc or share memory.""" + logger.debug(f"Direct sending large weight {name}({weight.shape}, {weight.dtype})") + # TODO: support fallback to shared memory + handle = reduce_tensor(weight) + bucket_meta: dict[str, TensorMetadata] = {} + bucket_meta[name] = { + "name": name, + "shape": weight.shape, + "dtype": weight.dtype, + "offset": 0, + "handle": handle, + } + self.socket.send_pyobj({"bucket_meta": bucket_meta, "is_last": False}) + self.socket.recv() + + +class BucketedWeightReceiver: + """ + Receive model weights via bucketed IPC transfer over ZMQ. + + Receives weight tensors from BucketedWeightSender and passes each + bucket to a callback for processing (e.g., loading into the model). + + Args: + zmq_handle: ZMQ IPC socket path (must match sender) + device: Target device for received tensors + use_shm: Use shared memory instead of CUDA IPC + """ + + def __init__( + self, + zmq_handle: str, + device: torch.device, + use_shm: bool = False, + ): + self.zmq_handle = zmq_handle + self.device = device + self.use_shm = use_shm + + self.zmq_context = zmq.Context.instance() + self.socket = None + self.buffer = None + self.shm = None + + def receive_weights(self, on_bucket_received: callable): + """ + Receive weights from sender and process each bucket via callback. + + Args: + on_bucket_received: Callback function(weights: list[(name, tensor)]) called per bucket. + """ + try: + self._init_socket() + self._init_buffer() + + # receive bucket and update weights + while True: + metadata = self.socket.recv_pyobj() + weights, tensor = [], None + for name, meta in metadata["bucket_meta"].items(): + shape, dtype, offset, handle = meta["shape"], meta["dtype"], meta["offset"], meta["handle"] + if handle is not None: + tensor = rebuild_ipc(handle, self.device.index) + weights.append((name, tensor)) + continue + size = dtype.itemsize * shape.numel() + tensor = self.buffer[offset : offset + size].view(dtype=dtype).view(shape) + if self.use_shm: + tensor = tensor.to(self.device) + weights.append((name, tensor)) + on_bucket_received(weights) + get_torch_device().synchronize() + self.socket.send(b"") + del weights, tensor + if metadata["is_last"]: + break + finally: + self._cleanup() + + def _init_socket(self): + """Initialize ZMQ REP socket and connect.""" + self.socket = self.zmq_context.socket(zmq.REP) + self.socket.connect(self.zmq_handle) + + def _init_buffer(self): + """Receive and rebuild communication buffer from sender.""" + comm_metadata = self.socket.recv_pyobj() + buffer, shm = None, None + if not self.use_shm: + handle = comm_metadata + buffer = rebuild_ipc(handle, self.device.index) + assert buffer.dtype == torch.uint8 + else: + shm_name = comm_metadata["name"] + shm_size = comm_metadata["size"] + buffer, shm = rebuild_shared_memory(shm_name, shm_size, dtype=torch.uint8) + self.socket.send(b"") + self.buffer = buffer + self.shm = shm + + def _cleanup(self): + """clean up""" + if self.socket is not None: + self.socket.close() + self.socket = None + # Synchronize before releasing the buffer to ensure all async ops + # referencing it (e.g. clone, .to()) have completed. + get_torch_device().synchronize() + del self.buffer + self.buffer = None + if self.shm is not None: + self.shm.close() + del self.shm + self.shm = None + gc.collect() + get_torch_device().ipc_collect() + get_torch_device().empty_cache() diff --git a/verl/verl/workers/rollout/vllm_rollout/utils.py b/verl/verl/workers/rollout/vllm_rollout/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d9da3a6a43d0bcc4d9b87ccb5a5a83d8ddf53d69 --- /dev/null +++ b/verl/verl/workers/rollout/vllm_rollout/utils.py @@ -0,0 +1,355 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import ctypes +import json +import logging +import os +import platform +import signal +import threading +from types import MethodType +from typing import Any, Literal, Optional, get_args + +import torch +from vllm.outputs import RequestOutput + +from verl.utils.device import is_npu_available +from verl.utils.vllm import TensorLoRARequest, VLLMHijack +from verl.utils.vllm.patch import patch_vllm_moe_model_weight_loader +from verl.utils.vllm.vllm_fp8_utils import apply_vllm_fp8_patches, is_fp8_model, load_quanted_weights + +logger = logging.getLogger(__file__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + +# magic numbers that ensure we are using the same LoRA adapter during the rollout and training process +VLLM_LORA_INT_ID = 123 +VLLM_LORA_NAME = "123" +VLLM_LORA_PATH = "simon_lora_path" + +VLLM_ASCEND_REQUIRED_ENV_VARS = {"VLLM_ALL2ALL_BACKEND": "flashinfer_all2allv", "VLLM_ASCEND_ENABLE_NZ": "0"} + + +def set_death_signal(): + """Kill the current process when the parent process exits.""" + if platform.system() != "Linux": + return + libc = ctypes.CDLL("libc.so.6") + libc.prctl(1, signal.SIGKILL) + if os.getppid() == 1: + os.kill(os.getpid(), signal.SIGKILL) + + +def get_device_uuid(device_id: int) -> str: + from vllm.platforms import current_platform + + # Convert torch.npu.current_device to its corresponding ASCEND_RT_VISIBLE_DEVICES. + if is_npu_available: + if os.getenv("ASCEND_RT_VISIBLE_DEVICES") is not None: + npu_visible_devices = os.environ["ASCEND_RT_VISIBLE_DEVICES"].split(",") + assert device_id < len(npu_visible_devices), f"device_id {device_id} must less than {npu_visible_devices}" + return "NPU-" + npu_visible_devices[device_id] + else: + return f"NPU-{device_id}" + else: + return current_platform.get_device_uuid(device_id) + + +def get_vllm_max_lora_rank(lora_rank: int): + """ + For vLLM, automatically adjusts the `max_lora_rank` to the nearest allowed value. + The allowed values are retrieved from vLLM's MaxLoRARanks type definition. + """ + assert lora_rank > 0, f"lora_rank must be greater than 0, get {lora_rank}" + + try: + from vllm.config.lora import MaxLoRARanks + except Exception: + # FIXME: migrate vllm version https://github.com/vllm-project/vllm/blob/main/vllm/config/lora.py#L25 + MaxLoRARanks = Literal[1, 8, 16, 32, 64, 128, 256, 320, 512] + + vllm_max_lora_ranks = sorted(get_args(MaxLoRARanks)) + if lora_rank > vllm_max_lora_ranks[-1]: + raise ValueError(f"lora_rank must be less than or equal to {vllm_max_lora_ranks[-1]}, but got {lora_rank}") + + for rank in vllm_max_lora_ranks: + if lora_rank <= rank: + return rank + + +# https://github.com/vllm-project/vllm/issues/13175 +def monkey_patch_compute_logits(model, vocab_size: int): + original_compute_logits = model.compute_logits + + def compute_logits( + self, + *args, + **kwargs, + ) -> torch.Tensor: + logits = original_compute_logits(*args, **kwargs) + logits[..., vocab_size:] = float("-inf") + return logits + + model.compute_logits = MethodType(compute_logits, model) + + +class vLLMColocateWorkerExtension: + """ + The class for vLLM's worker to inherit from, in the colocate setting. + By defining an extension class, the code can work no matter what is + the underlying worker class. This way, the code can be compatible + with both vLLM V0 and V1. + NOTE: we define this class in a separate module, and the main module + should pass the full qualified name as `worker_extension_cls` argument. + + Feature support: + 1. LoRA + 2. Online FP8 quantization + """ + + def __new__(cls, **kwargs): + set_death_signal() + + # 1. patch for Lora + VLLMHijack.hijack() + # 2. patch online fp8 quant + if os.environ.get("VERL_VLLM_FP8_QUANT_ENABLED", "0") == "1": + apply_vllm_fp8_patches() + # 3. patch QAT (compressed-tensors NVFP4) for dynamic weight loading + vllm_config = kwargs.get("vllm_config") + quant_config = getattr(vllm_config, "quant_config", None) if vllm_config else None + _is_qat_model = getattr(quant_config, "quant_format", None) == "nvfp4-pack-quantized" + _is_modelopt_qat = type(quant_config).__name__ == "ModelOptNvFp4Config" + if _is_qat_model: + from verl.utils.qat import apply_qat_patches + + apply_qat_patches() + logger.info("Applied QAT (compressed-tensors) patches in vLLM worker subprocess") + elif _is_modelopt_qat: + from verl.utils.modelopt import apply_modelopt_nvfp4_patches + + apply_modelopt_nvfp4_patches() + logger.info("Applied ModelOpt NVFP4 patches in vLLM worker subprocess") + + # TODO: For ascend NPU, when the corresponding vllm-ascend version is upgraded to v0.13.0, + # please remove the VLLM_ASCEND_REQUIRED_ENV_VARS variable replacement action. + # This is only a fix for vllm version < v0.13.0. + if is_npu_available: + for k in VLLM_ASCEND_REQUIRED_ENV_VARS: + if k not in os.environ: + os.environ[k] = VLLM_ASCEND_REQUIRED_ENV_VARS[k] + + instance = super().__new__(cls) + instance._is_qat_model = _is_qat_model + instance._is_modelopt_qat = _is_modelopt_qat + return instance + + def monkey_patch_model(self, vocab_size: int): + # patch compute_logits to avoid sampling OOV token + monkey_patch_compute_logits(self.model_runner.model, vocab_size) + # patch weight loader to support MoE model + patch_vllm_moe_model_weight_loader(self.model_runner.model) + + def update_weights_from_ipc(self, peft_config: dict = None, base_sync_done=False, use_shm: bool = False): + """Update the weights of the rollout model.""" + from vllm.platforms import current_platform + + from verl.workers.rollout.vllm_rollout.bucketed_weight_transfer import BucketedWeightReceiver + + if current_platform.device_type == "npu" and self.device is None: + self.device = torch.device(f"npu:{self.local_rank}") + + # In async mode, make sure the old lora is removed before adding the new one + if peft_config and base_sync_done: + self.remove_lora(VLLM_LORA_INT_ID) + + use_standard_weight_load = not (peft_config and base_sync_done) and not is_fp8_model( + self.model_runner.vllm_config + ) + + if self._is_qat_model: + # QAT (compressed-tensors): Prepare for weight loading BEFORE receiving any buckets + from verl.utils.qat import prepare_qat_for_load_weights + + prepare_qat_for_load_weights(self.model_runner.model, device=self.device) + logger.info("QAT: prepare_qat_for_load_weights completed") + elif self._is_modelopt_qat: + from verl.utils.modelopt.vllm_modelopt_patch import prepare_modelopt_for_weight_reload + + prepare_modelopt_for_weight_reload(self.model_runner.model, device=self.device) + logger.info("ModelOpt: prepare_modelopt_for_weight_reload completed") + elif use_standard_weight_load: + # Re-apply here because async IPC weight sync can happen long after init and lose MoE weight_loader attrs. + patch_vllm_moe_model_weight_loader(self.model_runner.model) + + assert self.device is not None + receiver = BucketedWeightReceiver( + zmq_handle=self._get_zmq_handle(), + device=self.device, + use_shm=use_shm, + ) + receiver.receive_weights( + on_bucket_received=lambda weights: self._update_weights( + weights, peft_config=peft_config, base_sync_done=base_sync_done + ) + ) + + if self._is_qat_model: + # QAT (compressed-tensors): call process_weights_after_loading AFTER all buckets are received + from verl.utils.qat import manual_process_weights_after_loading + + manual_process_weights_after_loading(self.model_runner.model) + logger.info("QAT: process_weights_after_loading completed") + elif self._is_modelopt_qat: + from verl.utils.modelopt.vllm_modelopt_patch import modelopt_process_weights_after_loading + + modelopt_process_weights_after_loading(self.model_runner.model) + logger.info("ModelOpt QAT: process_weights_after_loading completed") + elif use_standard_weight_load: + # Some post-load transforms are non-idempotent; run once after all buckets. + from vllm.model_executor.model_loader.utils import process_weights_after_loading + + model = self.model_runner.model + model_config = self.model_runner.vllm_config.model_config + process_weights_after_loading(model, model_config, self.device) + + def _update_weights(self, weights: list[tuple[str, torch.Tensor]], peft_config: dict, base_sync_done: bool): + if peft_config and base_sync_done: + weights = dict(weights) + lora_request = TensorLoRARequest( + lora_name=VLLM_LORA_NAME, + lora_int_id=VLLM_LORA_INT_ID, + lora_path=VLLM_LORA_PATH, + peft_config=peft_config, + lora_tensors=weights, + ) + self.add_lora(lora_request) + logger.info(f"vLLM load weights, loaded_params: {len(weights)}") + else: + # Add the FP8 related logic here as sharding manager has been deprecated. + # Check if FP8 quantization is enabled and apply appropriate weight loading + if is_fp8_model(self.model_runner.vllm_config): + logger.info(f"FP8 model detected (async): {self.model_runner.vllm_config.quant_config}") + # Convert bf16 weights to fp8 format before loading + loaded_params = load_quanted_weights(weights, self.model_runner) + logger.info(f"FP8 weights loaded (async), loaded_params: {len(loaded_params)}") + else: + logger.info("Loading standard weights (non-FP8, async)") + self.model_runner.model.load_weights(weights) + + def _get_zmq_handle(self) -> str: + """Get ZMQ handle for communication. + Uses Ray job id + replica_rank + local_rank to form the handle so it + matches the sender side regardless of CUDA_VISIBLE_DEVICES differences, + avoids collisions when multiple replicas share the same node, and is + unique per Ray job to avoid cross-job collisions on shared hosts. The + job id is forwarded by the vLLMHttpServer actor as VERL_RAY_JOB_ID and + inherited by this vLLM worker subprocess. + """ + replica_rank = os.environ.get("VERL_REPLICA_RANK", "0") + job_id = os.environ.get("VERL_RAY_JOB_ID", "0") + return f"ipc:///tmp/rl-colocate-zmq-{job_id}-replica-{replica_rank}-rank-{self.local_rank}.sock" + + +class SuppressSignalInThread: + def __enter__(self): + self.original_signal = signal.signal + + def no_op_signal(sig, action): + if threading.current_thread() is not threading.main_thread(): + print(f"Ignored signal {sig} in thread {threading.current_thread().name}") + return + return self.original_signal(sig, action) + + signal.signal = no_op_signal + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + signal.signal = self.original_signal + + +def build_cli_args_from_config(config: dict[str, Any]) -> list[str]: + """ + Convert a config dictionary to CLI arguments for vLLM server. + + Handles different value types appropriately: + - None: skipped + - bool True: adds '--key' + - bool False: skipped + - list: expands to '--key item1 item2 ...' + - empty list: skipped (vLLM uses nargs="+" which requires at least one value) + - dict: JSON serialized + - other: string converted + + Args: + config: Dictionary of configuration key-value pairs + + Returns: + List of CLI argument strings + """ + cli_args = [] + for k, v in config.items(): + if v is None: + continue + if isinstance(v, bool): + if v: + cli_args.append(f"--{k}") + elif isinstance(v, list): + if not v: + # Skip empty lists - vLLM uses nargs="+" which requires at least one value + continue + # Lists need to be expanded as multiple separate arguments + # e.g., --cuda-graph-sizes 1 2 4 8 becomes ['--cuda-graph-sizes', '1', '2', '4', '8'] + cli_args.append(f"--{k}") + cli_args.extend([str(item) for item in v]) + else: + cli_args.append(f"--{k}") + # Use json.dumps for dict to ensure valid JSON format + cli_args.append(json.dumps(v) if isinstance(v, dict) else str(v)) + return cli_args + + +def extract_prompt_logprobs(output: RequestOutput, num_prompt_logprobs: Optional[int], result_dict: dict[str, list]): + """Extract prompt log probabilities from generation output.""" + if num_prompt_logprobs is None: + return + + prompt_logprobs_ls, prompt_ids_ls = [], [] + # NOTE: logprob of first prompt token is None. + for logprobs_dict in output.prompt_logprobs[1:]: + if num_prompt_logprobs == 0: + token_id_str = list(logprobs_dict.keys())[0] + logprob = logprobs_dict[token_id_str].logprob + prompt_logprobs_ls.append([logprob]) + prompt_ids_ls.append([int(token_id_str)]) + else: + prompt_ids = [None] * num_prompt_logprobs + prompt_logprobs = [None] * num_prompt_logprobs + # We get either top-k logprobs or top-k plus the sampled logprob (if sampled token is not in top-k) + assert len(logprobs_dict) in [num_prompt_logprobs, num_prompt_logprobs + 1], len(logprobs_dict) + for token_id_str, token_logprob in logprobs_dict.items(): + rank = token_logprob.rank + if rank > num_prompt_logprobs: + continue # the sampled token is not in the top-k + logprob = token_logprob.logprob + prompt_ids[rank - 1] = int(token_id_str) + prompt_logprobs[rank - 1] = logprob + prompt_logprobs_ls.append(prompt_logprobs) + prompt_ids_ls.append(prompt_ids) + + # NOTE: pad a dummy prompt logprob for last prompt token. + prompt_logprobs_ls.append([0.0] * max(num_prompt_logprobs, 1)) + prompt_ids_ls.append([0] * max(num_prompt_logprobs, 1)) + + result_dict["prompt_ids"] = prompt_ids_ls + result_dict["prompt_logprobs"] = prompt_logprobs_ls diff --git a/verl/verl/workers/rollout/vllm_rollout/vllm_async_server.py b/verl/verl/workers/rollout/vllm_rollout/vllm_async_server.py new file mode 100644 index 0000000000000000000000000000000000000000..4e1b0cff5ed369fb5b20486120eaec4b44b8ac5f --- /dev/null +++ b/verl/verl/workers/rollout/vllm_rollout/vllm_async_server.py @@ -0,0 +1,1082 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import argparse +import asyncio +import inspect +import json +import logging +import os +from pprint import pprint +from typing import Any, Callable, Optional + +import ray +import vllm.entrypoints.cli.serve +from packaging import version +from ray.actor import ActorHandle +from vllm import SamplingParams +from vllm.engine.arg_utils import AsyncEngineArgs +from vllm.entrypoints.cli.serve import run_headless +from vllm.entrypoints.openai.api_server import build_app, init_app_state +from vllm.inputs import TokensPrompt +from vllm.lora.request import LoRARequest +from vllm.outputs import RequestOutput +from vllm.usage.usage_lib import UsageContext +from vllm.v1.engine.async_llm import AsyncLLM + +from verl.utils.config import omega_conf_to_dataclass +from verl.utils.device import get_resource_name, get_visible_devices_keyword, is_torch_npu_available +from verl.utils.net_utils import get_free_port, is_valid_ipv6_address +from verl.utils.profiler import DistProfiler, build_vllm_profiler_args +from verl.utils.tokenizer import normalize_token_ids +from verl.utils.vllm.vllm_fp8_utils import apply_vllm_fp8_patches +from verl.workers.config import HFModelConfig, RolloutConfig +from verl.workers.rollout.replica import RolloutMode, RolloutReplica, TokenOutput +from verl.workers.rollout.utils import get_max_position_embeddings, qwen2_5_vl_dedup_image_tokens, run_uvicorn +from verl.workers.rollout.vllm_rollout.utils import ( + VLLM_LORA_INT_ID, + VLLM_LORA_NAME, + VLLM_LORA_PATH, + SuppressSignalInThread, + build_cli_args_from_config, + extract_prompt_logprobs, + get_vllm_max_lora_rank, +) + +_VLLM_VERSION = version.parse(vllm.__version__) + + +if _VLLM_VERSION > version.parse("0.11.0"): + from vllm.utils.argparse_utils import FlexibleArgumentParser + + if _VLLM_VERSION == version.parse("0.12.0"): + from vllm.entrypoints.harmony_utils import get_encoding + + elif _VLLM_VERSION >= version.parse("0.13.0"): + from vllm.entrypoints.openai.parser.harmony_utils import get_encoding + + else: + get_encoding = None + + if get_encoding is not None and os.getenv("VERL_USE_GPT_OSS", "0") == "1": + get_encoding() +else: + from vllm.utils import FlexibleArgumentParser + + +logger = logging.getLogger(__file__) +logger.setLevel(logging.INFO) + + +class vLLMHttpServer: + """vLLM http server in single node, this is equivalent to launch server with command line: + ``` + vllm serve --tensor-parallel-size=8 ... + ``` + """ + + def __init__( + self, + config, + model_config, + rollout_mode: RolloutMode, + workers: list[ActorHandle], + replica_rank: int, + node_rank: int, + gpus_per_node: int, + nnodes: int, + cuda_visible_devices: str, + ): + """ + Args: + config (RolloutConfig): full config. + model_config (HFModelConfig): model config. + rollout_mode (RolloutMode): rollout mode. + replica_rank (int): replica rank, a replica may contain multiple nodes. + node_rank (int): node rank. + gpus_per_node (int): number of gpus per node. + nnodes (int): number of nodes. + cuda_visible_devices (str): cuda visible devices. + """ + os.environ[get_visible_devices_keyword()] = cuda_visible_devices + os.environ["VERL_REPLICA_RANK"] = str(replica_rank) + # Forward the Ray job id into the vLLM worker subprocess so the + # colocated weight-transfer IPC socket path is unique per Ray job. + # Without this, two concurrent verl jobs on the same node both bind + # the same /tmp/rl-colocate-zmq-replica-0-rank-0.sock and one fails + # with EADDRINUSE; a stale socket from a crashed run trips the same + # error on restart. + os.environ["VERL_RAY_JOB_ID"] = ray.get_runtime_context().get_job_id() + + self.config = self._init_config(config) + self.model_config = self._init_model_config(model_config) + self._validate_configs() + + self.rollout_mode = rollout_mode + self.workers = workers + + self.replica_rank = replica_rank + self.node_rank = node_rank + self.gpus_per_node = gpus_per_node + self.nnodes = nnodes + # model weights version, set by ServerAdapter when update weights. + self.global_steps = None + + if self.rollout_mode != RolloutMode.HYBRID and self.config.load_format == "dummy": + logger.warning(f"rollout mode is {self.rollout_mode}, load_format is dummy, set to auto") + self.config.load_format = "auto" + + # used for http server + self._server_address = ray.util.get_node_ip_address().strip("[]") + self._server_port = None + + # used for controlling vllm server profiler + profiler_config = self.config.profiler + tool_config = None + if profiler_config is not None: + if profiler_config.tool in ["torch", "npu"]: + tool_config = omega_conf_to_dataclass((profiler_config.tool_config or {}).get(profiler_config.tool)) + else: + logger.warning(f"agent loop only support torch and npu profiler, got {profiler_config.tool}") + profiler_config = None + self.profiler_controller = DistProfiler(self.replica_rank, config=profiler_config, tool_config=tool_config) + + # used for data parallel: --data-parallel-address, --data-parallel-rpc-port + if self.node_rank == 0: + self._master_address = self._server_address + # used for torch.distributed.init_process_group + self._master_port, self._master_sock = get_free_port(self._server_address, with_alive_sock=True) + # used for data parallel: --data-parallel-address, --data-parallel-rpc-port + self._dp_rpc_port, self._dp_rpc_sock = get_free_port(self._server_address, with_alive_sock=True) + self._dp_master_port, self._dp_master_sock = get_free_port(self._server_address, with_alive_sock=True) + else: + self._master_address = None + self._master_port = None + self._dp_rpc_port = None + self._dp_master_port = None + + self._post_init(cuda_visible_devices) + + def get_master_address(self): + """Get master address and port for data parallel. + Returns: + tuple: (master_address, master_port, dp_rpc_port) + """ + return self._master_address, self._master_port, self._dp_rpc_port + + def get_server_address(self): + """Get http server address and port.""" + assert self._server_port is not None, "http server is not launched, port is None" + return self._server_address, self._server_port + + @property + def lora_as_adapter(self) -> bool: + return ( + self.model_config.lora_rank > 0 or self.model_config.lora.get("rank", 0) > 0 + ) and not self.model_config.lora.get("merge", False) + + async def collective_rpc( + self, + method: str | Callable, + timeout: float | None = None, + args: tuple = (), + kwargs: dict[str, Any] | None = None, + ): + await self.engine.collective_rpc( + method=method, + timeout=timeout, + args=args, + kwargs=kwargs, + ) + + async def launch_server(self, master_address: str = None, master_port: int = None, dp_rpc_port: int = None): + if self.node_rank != 0: + assert master_address and master_port and dp_rpc_port, ( + "non-master node should provide master_address, master_port and dp_rpc_port" + ) + self._master_address = master_address + self._master_port = master_port + self._dp_rpc_port = dp_rpc_port + + # 1. setup vllm serve cli args + engine_kwargs = self.config.get("engine_kwargs", {}).get(self._get_engine_kwargs_key(), {}) or {} + engine_kwargs = {key: val for key, val in engine_kwargs.items() if val is not None} + if self.config.get("limit_images", None): # support for multi-image data + engine_kwargs["limit_mm_per_prompt"] = {"image": self.config.get("limit_images")} + if self.config.cudagraph_capture_sizes: + engine_kwargs["cuda_graph_sizes"] = self.config.cudagraph_capture_sizes + + self._preprocess_engine_kwargs(engine_kwargs) + + # Override default generation config from hugging face model config, + # user can still override them by passing kwargs in each request. + override_generation_config = self._get_override_generation_config() + logger.info(f"override_generation_config: {override_generation_config}") + + logger.info(f"enable_sleep_mode: {self.config.enable_sleep_mode}") + if not self.config.enable_sleep_mode: + from verl.utils.device import set_expandable_segments + + set_expandable_segments(True) + + quantization, hf_overrides = self._apply_quantization() + + compilation_config = engine_kwargs.pop("compilation_config", None) or {} + if isinstance(compilation_config, str): + compilation_config = json.loads(compilation_config) + compilation_config.setdefault("cudagraph_mode", "FULL_AND_PIECEWISE") + + # FULL cuda graph is not yet supported with DCP, downgrade to PIECEWISE + dcp_size = engine_kwargs.get("decode_context_parallel_size", 1) or 1 + if dcp_size > 1 and compilation_config["cudagraph_mode"] == "FULL_AND_PIECEWISE": + logger.warning( + "FULL cuda graph is not supported with DCP (decode_context_parallel_size=%d), " + "downgrading cudagraph_mode to PIECEWISE.", + dcp_size, + ) + compilation_config["cudagraph_mode"] = "PIECEWISE" + + compilation_config = json.dumps(compilation_config) + args = { + "dtype": self.config.dtype, + "load_format": self.config.load_format, + "skip_tokenizer_init": False, + "distributed_executor_backend": "mp", + "worker_extension_cls": self._get_worker_extension_cls(), + "trust_remote_code": self.model_config.trust_remote_code, + "max_model_len": self.config.max_model_len, + "max_num_seqs": self.config.max_num_seqs, + "enable_chunked_prefill": self.config.enable_chunked_prefill, + "max_num_batched_tokens": self.config.max_num_batched_tokens, + "enable_prefix_caching": self.config.enable_prefix_caching, + "enable_sleep_mode": self.config.enable_sleep_mode, + "logprobs_mode": self.config.logprobs_mode, + "enforce_eager": self.config.enforce_eager, + "gpu_memory_utilization": self.config.gpu_memory_utilization, + "disable_log_stats": self.config.disable_log_stats, + "tensor_parallel_size": self.config.tensor_model_parallel_size, + "seed": self.replica_rank + self.config.get("seed", 0), + "override_generation_config": json.dumps(override_generation_config), + "quantization": quantization, + "hf_overrides": hf_overrides, + "scheduling_policy": self.config.scheduling_policy, + "compilation_config": compilation_config, + **engine_kwargs, + } + + # update profiler args + profiler_args = build_vllm_profiler_args( + self.profiler_controller.config, self.profiler_controller.tool_config, self.replica_rank + ) + if _VLLM_VERSION >= version.parse("0.13.0"): + # vLLM >= 0.13.0 supports profiler config via CLI args; env vars still work but will be deprecated + args.update(profiler_args) + + if self.config.prometheus.enable: + if self.config.prometheus.served_model_name: + # Extract model name from path if it's a full path + served_model_name = self.config.prometheus.served_model_name + if "/" in served_model_name: + # If it's a full path, extract the last part as model name + served_model_name = served_model_name.split("/")[-1] + args["served_model_name"] = served_model_name + + if self.config.mtp is not None and self.config.mtp.enable and self.config.mtp.enable_rollout: + speculative_config = { + "method": self.config.mtp.method, + "num_speculative_tokens": self.config.mtp.num_speculative_tokens, + } + args["speculative_config"] = speculative_config + + if self.config.data_parallel_size > 1: + assert self.gpus_per_node % self.config.tensor_model_parallel_size == 0, ( + "gpus_per_node should be divisible by tensor_model_parallel_size" + ) + data_parallel_size_local = self.gpus_per_node // self.config.tensor_model_parallel_size + assert len(self.workers) == data_parallel_size_local * self.config.tensor_model_parallel_size, ( + f"num workers ({len(self.workers)}) should be equal to " + f"dp_size_local ({data_parallel_size_local}) * tp_size ({self.config.tensor_model_parallel_size})" + ) + dp_args = { + "data_parallel_size": self.config.data_parallel_size, + "data_parallel_size_local": data_parallel_size_local, + "data_parallel_start_rank": self.node_rank * data_parallel_size_local, + "data_parallel_address": self._master_address, + "data_parallel_rpc_port": self._dp_rpc_port, + } + args.update(dp_args) + + args.update({"enable_expert_parallel": self.config.expert_parallel_size > 1}) + + # used for torch.distributed.init_process_group + if self.nnodes > 1: + args.update( + { + "master_addr": self._master_address, + "master_port": self._master_port, + "node_rank": self.node_rank, + "nnodes": self.nnodes, + "data_parallel_address": self._master_address, + "data_parallel_rpc_port": self._dp_rpc_port, + } + ) + + # update lora-related args + lora_rank = self.model_config.lora.get("rank", 0) + if lora_rank <= 0: + lora_rank = ( + self.model_config.lora_rank + ) # FIXME: fallback to lora_rank for now, we should unify lora settings. + + if self.model_config.lora.get("merge", False): + lora_rank = 0 + + if lora_rank > 0: + lora_args = { + "enable_lora": True, + "max_loras": 1, + "max_lora_rank": get_vllm_max_lora_rank(lora_rank), + } + if self.model_config.lora.get("fully_sharded_loras", False): + lora_args["fully_sharded_loras"] = True + args.update(lora_args) + + if self.config.enable_rollout_routing_replay: + args.update({"enable_return_routed_experts": True}) + + server_args = ["serve", self.model_config.local_path] + build_cli_args_from_config(args) + + if self.replica_rank == 0: + pprint(server_args) + + CMD_MODULES = self._get_cli_modules() + parser = FlexibleArgumentParser(description=self._get_cli_description()) + subparsers = parser.add_subparsers(required=False, dest="subparser") + cmds = {} + for cmd_module in CMD_MODULES: + new_cmds = cmd_module.cmd_init() + for cmd in new_cmds: + cmd.subparser_init(subparsers).set_defaults(dispatch_function=cmd.cmd) + cmds[cmd.name] = cmd + server_args = parser.parse_args(args=server_args) + server_args.model = server_args.model_tag + if server_args.subparser in cmds: + cmds[server_args.subparser].validate(server_args) + + # 3. launch server + if self.node_rank == 0: + await self.run_server(server_args) + else: + await self.run_headless(server_args) + + async def run_server(self, args: argparse.Namespace): + engine_args = AsyncEngineArgs.from_cli_args(args) + usage_context = UsageContext.OPENAI_API_SERVER + vllm_config = engine_args.create_engine_config(usage_context=usage_context) + vllm_config.parallel_config.data_parallel_master_port = self._dp_master_port + + fn_args = set(dict(inspect.signature(AsyncLLM.from_vllm_config).parameters).keys()) + kwargs = {} + if "enable_log_requests" in fn_args: + kwargs["enable_log_requests"] = engine_args.enable_log_requests + if "disable_log_stats" in fn_args: + kwargs["disable_log_stats"] = engine_args.disable_log_stats + + engine_client = AsyncLLM.from_vllm_config(vllm_config=vllm_config, usage_context=usage_context, **kwargs) + + # Don't keep the dummy data in memory + await engine_client.reset_mm_cache() + await engine_client.collective_rpc( + method="monkey_patch_model", kwargs={"vocab_size": len(self.model_config.tokenizer)} + ) + + build_app_sig = inspect.signature(build_app) + supported_tasks: tuple[Any, ...] = () + if "supported_tasks" in build_app_sig.parameters: + supported_tasks = await engine_client.get_supported_tasks() + app = build_app(args, supported_tasks) + else: + app = build_app(args) + + init_app_sig = inspect.signature(init_app_state) + if "vllm_config" in init_app_sig.parameters: + await init_app_state(engine_client, vllm_config, app.state, args) + elif "supported_tasks" in init_app_sig.parameters: + await init_app_state(engine_client, app.state, args, supported_tasks) + else: + await init_app_state(engine_client, app.state, args) + if self.replica_rank == 0 and self.node_rank == 0: + logger.info(f"Initializing a V1 LLM engine with config: {vllm_config}") + + self.engine = engine_client + self._server_port, self._server_task = await run_uvicorn(app, args, self._server_address) + + async def run_headless(self, args: argparse.Namespace): + """Run headless server in a separate thread.""" + args.api_server_count = 0 + + def run_headless_wrapper(): + with SuppressSignalInThread(): + run_headless(args) + + def on_run_headless_done(future: asyncio.Future): + try: + exc = future.exception() + if exc: + logger.exception(f"run_headless failed with exception: {exc}") + else: + logger.warning("run_headless completed successfully, but it's not expected.") + except Exception as e: + logger.exception(f"get result from run_headless failed: {e}") + finally: + os._exit(1) + + self.task = asyncio.create_task(asyncio.to_thread(run_headless_wrapper)) + self.task.add_done_callback(on_run_headless_done) + + async def generate( + self, + prompt_ids: list[int], + sampling_params: dict[str, Any], + request_id: str, + image_data: Optional[list[Any]] = None, + video_data: Optional[list[Any]] = None, + audio_data: Optional[list[Any]] = None, + mm_processor_kwargs: Optional[dict[str, Any]] = None, + priority: int = 0, + ) -> TokenOutput: + """Generate sequence with token-in-token-out.""" + prompt_ids = normalize_token_ids(prompt_ids) + + # Calculate the maximum possible new tokens based on available context space + # This serves as a safety upper bound + max_possible_tokens = self.config.max_model_len - len(prompt_ids) + if max_possible_tokens < 0: + raise ValueError( + f"Prompt length ({len(prompt_ids)}) exceeds the model's maximum context length " + f"({self.config.max_model_len})." + ) + + # Determine max_tokens from sampling_params or use configured response_length as default + if "max_tokens" in sampling_params: + max_tokens = sampling_params.pop("max_tokens") + elif "max_new_tokens" in sampling_params: + # support sglang-style 'max_new_tokens' param + max_tokens = sampling_params.pop("max_new_tokens") + else: + # Default to a calculation that considers configured lengths + # Cap max_tokens by response_length to ensure tensor alignment, + # and by remaining budget to prevent OOM in multi-turn rollouts. + max_tokens = min( + self.config.response_length, self.config.prompt_length + self.config.response_length - len(prompt_ids) + ) + + # Clamp max_tokens to the valid range [0, max_possible_tokens] + max_tokens = max(0, min(max_tokens, max_possible_tokens)) + + assert max_tokens <= max_possible_tokens, ( + f"max_tokens {max_tokens} exceeds available context space {max_possible_tokens}" + ) + sampling_params["logprobs"] = 0 if sampling_params.pop("logprobs", False) else None + sampling_params.setdefault("repetition_penalty", self.config.get("repetition_penalty", 1.0)) + sampling_params = SamplingParams(max_tokens=max_tokens, **sampling_params) + prompt_ids = qwen2_5_vl_dedup_image_tokens(prompt_ids, self.model_config.processor) + multi_modal_data = {} + if image_data is not None: + multi_modal_data["image"] = image_data + if video_data is not None: + multi_modal_data["video"] = video_data + if audio_data is not None: + multi_modal_data["audio"] = audio_data + + prompt_kwargs = {"prompt_token_ids": prompt_ids, "multi_modal_data": multi_modal_data} + if mm_processor_kwargs: + prompt_kwargs["mm_processor_kwargs"] = mm_processor_kwargs + try: + prompt = TokensPrompt(**prompt_kwargs) + except TypeError: + prompt = prompt_kwargs + + # Add lora request + lora_request = None + if self.lora_as_adapter: + # Make sure we also check that the lora is already loaded in the engine + lora_loaded = VLLM_LORA_INT_ID in await self.engine.list_loras() + if lora_loaded: + lora_request = LoRARequest( + lora_name=VLLM_LORA_NAME, lora_int_id=VLLM_LORA_INT_ID, lora_path=VLLM_LORA_PATH + ) + + generator = self.engine.generate( + prompt=prompt, + sampling_params=sampling_params, + request_id=request_id, + lora_request=lora_request, + priority=priority, + ) + + # Get final response + final_res: Optional[RequestOutput] = None + async for output in generator: + final_res = output + assert final_res is not None + + # Handle abort case: when the request is aborted by pause_generation(abort), + # outputs may be empty. Return empty results with stop_reason="aborted" + # instead of crashing with "IndexError: list index out of range". + if not final_res.outputs: + return TokenOutput( + token_ids=[], + log_probs=None, + routed_experts=None, + stop_reason="aborted", + ) + + extra_fields = {"global_steps": self.global_steps} + extract_prompt_logprobs( + output=final_res, + num_prompt_logprobs=sampling_params.prompt_logprobs, + result_dict=extra_fields, + ) + token_ids = final_res.outputs[0].token_ids + log_probs = None + if sampling_params.logprobs is not None: + log_probs = [logprobs[token_ids[i]].logprob for i, logprobs in enumerate(final_res.outputs[0].logprobs)] + + routed_experts = None + if self.config.enable_rollout_routing_replay: + routed_experts = final_res.outputs[0].routed_experts + + # Determine stop reason from finish_reason + finish_reason = final_res.outputs[0].finish_reason + if finish_reason == "abort": + stop_reason = "aborted" + elif finish_reason in ("stop", "length"): + stop_reason = "completed" + else: + stop_reason = finish_reason # for more stop reason in the future + + num_preempted = None + + if hasattr(final_res.outputs[0], "num_preempted"): + num_preempted = final_res.outputs[0].num_preempted + + return TokenOutput( + token_ids=token_ids, + log_probs=log_probs, + routed_experts=routed_experts, + stop_reason=stop_reason, + num_preempted=num_preempted, + extra_fields=extra_fields, + ) + + async def wake_up(self): + if self.node_rank != 0: + return + + if self.rollout_mode == RolloutMode.HYBRID: + # In hybrid mode, rollout is wake up in `update_weights` + raise ValueError(f"wake_up not support rollout_mode {self.rollout_mode}") + elif self.rollout_mode == RolloutMode.COLOCATED: + # Directly call engine to wake up without sync weights. + await self.engine.wake_up(tags=self._get_wake_up_tags()) + await self.engine.reset_prefix_cache() + elif self.rollout_mode == RolloutMode.STANDALONE: + logger.info("skip wake_up in standalone mode") + + async def sleep(self): + if self.node_rank != 0 or not self.config.free_cache_engine: + return + + if self.rollout_mode == RolloutMode.HYBRID: + await self._sleep_hybrid() + elif self.rollout_mode == RolloutMode.COLOCATED: + await self.engine.sleep(level=1) + elif self.rollout_mode == RolloutMode.STANDALONE: + logger.info("skip sleep in standalone mode") + + async def clear_kv_cache(self): + if self.node_rank == 0: + await self.engine.reset_prefix_cache() + + async def release_kv_cache(self): + """Release only kv_cache GPU memory, keeping model weights intact. + # TODO: support true release of kv_cache + """ + if self.node_rank != 0 or not self.config.free_cache_engine: + return + + async def resume_kv_cache(self): + """Restore kv_cache GPU memory after a weight sync. Counterpart to release_kv_cache().""" + if self.node_rank != 0: + return + + async def start_profile(self, **kwargs): + if ( + self.profiler_controller.check_enable() + and self.profiler_controller.check_this_rank() + and self.profiler_controller.is_discrete_mode() + ): + await self.engine.start_profile(**kwargs) + + async def stop_profile(self): + if ( + self.profiler_controller.check_enable() + and self.profiler_controller.check_this_rank() + and self.profiler_controller.is_discrete_mode() + ): + await self.engine.stop_profile() + + async def set_global_steps(self, global_steps: int): + """Set the global steps of the model weights.""" + self.global_steps = global_steps + + async def wait_for_requests_to_drain(self): + await self.engine.wait_for_requests_to_drain() + + async def abort_all_requests(self, reset_prefix_cache: bool = True) -> dict[str, Any]: + """Abort all ongoing generation requests. + + On vLLM >= 0.12.0, uses AsyncLLM.pause_generation() to abort in-flight + requests, drain, and clear caches. The engine remains paused after this + call — use resume_generation() to accept new requests (e.g. before + validation). + + On vLLM < 0.12.0, manually aborts each request and resets prefix cache. + + Returns: + dict[str, Any]: Dictionary containing: + - aborted_count: Number of requests aborted + - request_ids: List of aborted request IDs + """ + try: + if _VLLM_VERSION >= version.parse("0.12.0"): + # Snapshot request IDs before pausing for reporting + request_ids = list(self.engine.output_processor.request_states.keys()) + + # pause_generation with wait_for_inflight_requests=False will: + # 1. Set engine to paused state (blocks new generate calls) + # 2. Abort all in-flight requests + # 3. Wait for requests to drain + # 4. Clear prefix and mm caches if clear_cache=True + await self.engine.pause_generation( + wait_for_inflight_requests=False, + clear_cache=reset_prefix_cache, + ) + else: + # Take an atomic snapshot to avoid race conditions with the vLLM engine thread + request_states_snapshot = list(self.engine.output_processor.request_states.items()) + request_ids = [req_id for req_id, _ in request_states_snapshot] + + if not request_ids: + return {"aborted_count": 0, "request_ids": []} + + # For each request, create an abort output and put it to its queue + # This allows the generator to receive the aborted result + from vllm.v1.engine import FinishReason + + for _, req_state in request_states_snapshot: + request_output = req_state.make_request_output( + [], pooling_output=None, finish_reason=FinishReason.ABORT, stop_reason=None + ) + req_state.queue.put(request_output) + + # Abort requests in the output processor and engine core + self.engine.output_processor.abort_requests(request_ids) + await self.engine.engine_core.abort_requests_async(request_ids) + + # Try to reset prefix cache to ensure clean state + if reset_prefix_cache: + await self.clear_kv_cache() + logger.info("Prefix cache reset after abort") + + logger.info(f"Aborted {len(request_ids)} requests: {request_ids}") + return {"aborted_count": len(request_ids), "request_ids": request_ids} + + except Exception as e: + logger.error(f"Error aborting requests: {e}") + return {"aborted_count": 0, "request_ids": [], "error": str(e)} + + async def resume_generation(self): + """Resume generation after abort_all_requests (pause_generation). + + Only effective on vLLM >= 0.12.0 where pause_generation is used. + No-op on older versions. + """ + if self.node_rank != 0: + return + if _VLLM_VERSION >= version.parse("0.12.0"): + await self.engine.resume_generation() + + async def abort_request(self, request_id: str, reset_prefix_cache: bool = True) -> dict[str, Any]: + """Abort a specific generation request. + + Args: + request_id: The ID of the request to abort. + + Returns: + dict[str, Any]: Dictionary containing abort result. + """ + try: + request_states = self.engine.output_processor.request_states + req_state = request_states.get(request_id) + + if req_state is None: + return {"aborted": False, "error": f"Request {request_id} not found"} + + # Create abort output and put it to the queue + from vllm.v1.engine import FinishReason + + request_output = req_state.make_request_output( + [], pooling_output=None, finish_reason=FinishReason.ABORT, stop_reason=None + ) + req_state.queue.put(request_output) + + # Abort in output processor and engine core + self.engine.output_processor.abort_requests([request_id]) + await self.engine.engine_core.abort_requests_async([request_id]) + + # Try to reset prefix cache to ensure clean state + if reset_prefix_cache: + await self.clear_kv_cache() + logger.info(f"Prefix cache reset after abort request {request_id}") + + logger.info(f"Aborted request: {request_id}") + return {"aborted": True, "request_id": request_id} + + except Exception as e: + logger.error(f"Error aborting request {request_id}: {e}") + return {"aborted": False, "request_id": request_id, "error": str(e)} + + # ----------------------------------------------------------------------- + # Hook methods for subclass overrides + # ----------------------------------------------------------------------- + + def _init_config(self, config): + """Initialise config. Override when a specific dataclass_type is needed.""" + return omega_conf_to_dataclass(config) + + def _init_model_config(self, model_config): + """Initialise model_config. Override when a specific dataclass_type is needed.""" + return omega_conf_to_dataclass(model_config, dataclass_type=HFModelConfig) + + def _validate_configs(self) -> None: + """Validate config/model_config after initialisation.""" + max_position_embeddings = get_max_position_embeddings(self.model_config.hf_config) + if self.config.max_model_len is None: + self.config.max_model_len = max_position_embeddings + else: + if self.config.max_model_len > max_position_embeddings: + raise ValueError( + f"max_model_len ({self.config.max_model_len}) should be less than or equal to " + f"max_position_embeddings ({max_position_embeddings})" + ) + + def _post_init(self, cuda_visible_devices: str) -> None: + """Called at the end of __init__. Default logs server metadata.""" + logger.info( + f"{self.__class__.__name__}, replica_rank: {self.replica_rank}, node_rank: {self.node_rank}, " + f"{get_visible_devices_keyword()}: {cuda_visible_devices}, " + f"master_address: {self._master_address}, master_port: {self._master_port}, " + f"data_parallel_rpc_port: {self._dp_rpc_port}, data_parallel_master_port: {self._dp_master_port}" + ) + + def _get_engine_kwargs_key(self) -> str: + """Return the key under config.engine_kwargs for this engine (e.g. 'vllm').""" + return "vllm" + + def _preprocess_engine_kwargs(self, engine_kwargs: dict) -> None: + """Mutate engine_kwargs in-place before the CLI args dict is built. No-op by default.""" + pass + + def _get_override_generation_config(self) -> dict: + """Return the override_generation_config dict.""" + # Override default generation config from hugging face model config, + # user can still override them by passing kwargs in each request. + return dict( + temperature=self.config.temperature, + top_k=self.config.top_k, + top_p=self.config.top_p, + repetition_penalty=1.0, + max_new_tokens=self.config.response_length, + ) + + def _apply_quantization(self) -> tuple[Optional[str], dict]: + """Process quantization config. Returns (quantization_str, hf_overrides).""" + quantization = self.config.quantization + hf_overrides = {} + + if is_torch_npu_available(check_device=False): + from verl.utils.vllm.npu_vllm_patch import check_vllm_ascend_before_server_launch + + check_vllm_ascend_before_server_launch() + + # Handle QAT (Quantization-Aware Training) configuration + qat_config_dict = getattr(self.config, "qat", {}) or {} + if qat_config_dict.get("enable", False): + from verl.utils.qat import QATConfig, load_quantization_config + + qat_config = QATConfig(**qat_config_dict) + quantization_config_dict = load_quantization_config(qat_config) + quant_method = quantization_config_dict.get("quant_method", None) + + if quant_method == "modelopt": + from verl.utils.modelopt import apply_modelopt_nvfp4_patches + + apply_modelopt_nvfp4_patches() + quantization = "modelopt" + elif quant_method == "compressed-tensors": + from verl.utils.qat import apply_qat_patches + + apply_qat_patches() + quantization = "compressed-tensors" + else: + raise ValueError(f"Unsupported quant_method: {quant_method}") + + logger.info(f"QAT quantization config injected (quant_method={quant_method})") + hf_overrides["quantization_config"] = quantization_config_dict + elif quantization is not None: + # Handle other quantization methods (fp8, torchao) + _SUPPORTED_QUANTIZATION = ["fp8", "torchao", "ascend"] + if quantization not in _SUPPORTED_QUANTIZATION: + raise ValueError(f"Currently only support {_SUPPORTED_QUANTIZATION} quantization, got: {quantization}") + + if quantization == "fp8": + # Ignore MoE router layers for FP8 quantization + all_mlp_gate_layers = [] + for layer in range(self.model_config.hf_config.num_hidden_layers): + all_mlp_gate_layers.append(f"model.layers.{layer}.mlp.gate") + + FP8_BLOCK_QUANT_KWARGS = { + "activation_scheme": "dynamic", + "fmt": "e4m3", + "quant_method": "fp8", + "weight_block_size": [128, 128], + "ignored_layers": all_mlp_gate_layers, + } + hf_overrides["quantization_config"] = dict(FP8_BLOCK_QUANT_KWARGS) + # Apply vllm fp8 patches + # Will remove the patch after vllm support on-the-fly quant for rollout natively. + apply_vllm_fp8_patches() + # for subprocesses patching + os.environ["VERL_VLLM_FP8_QUANT_ENABLED"] = "1" + + if quantization is not None and self.config.quantization_config_file is not None: + hf_overrides["quantization_config_file"] = self.config.quantization_config_file + + return quantization, hf_overrides + + def _get_worker_extension_cls(self) -> str: + """Return the fully-qualified colocate worker extension class name.""" + return "verl.workers.rollout.vllm_rollout.utils.vLLMColocateWorkerExtension" + + def _get_cli_modules(self) -> list: + """Return the list of CLI command modules used for argument parsing.""" + return [vllm.entrypoints.cli.serve] + + def _get_cli_description(self) -> str: + """Return the description string for the CLI argument parser.""" + return "vLLM CLI" + + def _get_wake_up_tags(self) -> list[str]: + """Return the tags passed to engine.wake_up(). Default includes kv_cache.""" + return ["kv_cache", "weights"] + + async def _sleep_hybrid(self): + """HYBRID sleep: lora adapters only need level=1; full weights need level=2.""" + # Don't use engine.sleep(level=2) here + # lora only update adapter weights, so set sleep level to 1 + # vllm_ascend not support sleep_level now. Enabling EP during training may lead to accuracy issues. + if self.lora_as_adapter or is_torch_npu_available(check_device=False): + sleep_level = 1 + else: + sleep_level = 2 + await self.engine.collective_rpc("sleep", kwargs={"level": sleep_level}) + if _VLLM_VERSION >= version.parse("0.17.0"): + await self.engine.reset_encoder_cache() + + +class vLLMReplica(RolloutReplica): + def __init__( + self, + replica_rank: int, + config: RolloutConfig, + model_config: HFModelConfig, + gpus_per_node: int = 8, + is_reward_model: bool = False, + is_teacher_model: bool = False, + name_suffix: str = "", + ): + super().__init__( + replica_rank, config, model_config, gpus_per_node, is_reward_model, is_teacher_model, name_suffix + ) + self.server_class = ray.remote(vLLMHttpServer) + + async def launch_servers(self): + """Launch http server in each node.""" + assert len(self.workers) == self.world_size, ( + f"worker number {len(self.workers)} not equal to world size {self.world_size}" + ) + + self._validate_launch_requirements() + + # get (node_id, CUDA_VISIBLE_DEVICES) of all workers + worker_infos = await asyncio.gather( + *[ + worker.__ray_call__.remote( + lambda self: ( + ray.get_runtime_context().get_node_id(), + ray.get_runtime_context().get_accelerator_ids()[get_resource_name()][0], + ) + ) + for worker in self.workers + ] + ) + worker_cuda_visible_devices = [worker_info[1] for worker_info in worker_infos] + worker_node_ids = [worker_info[0] for worker_info in worker_infos] + + # create server actor in each node with node affinity and cuda visible devices + nnodes, gpus_per_replica_node = self.nnodes, self.gpus_per_replica_node + for node_rank in range(nnodes): + workers = self.workers[node_rank * gpus_per_replica_node : (node_rank + 1) * gpus_per_replica_node] + node_cuda_visible_devices = ",".join( + worker_cuda_visible_devices[node_rank * gpus_per_replica_node : (node_rank + 1) * gpus_per_replica_node] + ) + node_id = worker_node_ids[node_rank * gpus_per_replica_node] + prefix = self._get_server_name_prefix() + if self.is_reward_model: + name = f"{prefix}server_reward_{self.replica_rank}_{node_rank}{self.name_suffix}" + elif self.is_teacher_model: + name = f"{prefix}server_teacher_{self.replica_rank}_{node_rank}{self.name_suffix}" + else: + name = f"{prefix}server_{self.replica_rank}_{node_rank}{self.name_suffix}" + server = self.server_class.options( + scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy( + node_id=node_id, + soft=False, + ), + runtime_env={ + "env_vars": { + "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1", + "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES": "1", + # To prevent hanging or crash during synchronization of weights between actor and rollout + # in disaggregated mode. See: + # https://docs.vllm.ai/en/latest/usage/troubleshooting.html?h=nccl_cumem_enable#known-issues + # https://github.com/vllm-project/vllm/blob/c6b0a7d3ba03ca414be1174e9bd86a97191b7090/vllm/worker/worker_base.py#L445 + "NCCL_CUMEM_ENABLE": "0", + } + }, + name=name, + max_concurrency=self.max_concurrency, + ).remote( + config=self.config, + model_config=self.model_config, + rollout_mode=self.rollout_mode, + workers=workers, + replica_rank=self.replica_rank, + node_rank=node_rank, + gpus_per_node=gpus_per_replica_node, + nnodes=nnodes, + cuda_visible_devices=node_cuda_visible_devices, + ) + self.servers.append(server) + + # launch http server in each node + master_address, master_port, dp_rpc_port = await self.servers[0].get_master_address.remote() + await asyncio.gather( + *[ + server.launch_server.remote( + master_address=master_address, master_port=master_port, dp_rpc_port=dp_rpc_port + ) + for server in self.servers + ] + ) + + # get http server address from first server + server_address, server_port = await self.servers[0].get_server_address.remote() + self._server_handle = self.servers[0] + self._server_address = ( + f"[{server_address}]:{server_port}" + if is_valid_ipv6_address(server_address) + else f"{server_address}:{server_port}" + ) + + async def sleep(self): + """Sleep each rollout server.""" + # Drain DP engines for safe sleep. + await self.servers[0].wait_for_requests_to_drain.remote() + await asyncio.gather(*[server.sleep.remote() for server in self.servers]) + + async def abort_all_requests(self) -> dict[str, Any]: + """Abort all ongoing generation requests across all servers. + + Returns: + dict[str, Any]: Combined abort results from all servers. + """ + results = await asyncio.gather(*[server.abort_all_requests.remote() for server in self.servers]) + + total_aborted = sum(r.get("aborted_count", 0) for r in results) + all_request_ids = [] + for r in results: + all_request_ids.extend(r.get("request_ids", [])) + + return { + "aborted_count": total_aborted, + "request_ids": all_request_ids, + "server_results": results, + } + + async def resume_generation(self): + """Resume generation on all servers after abort_all_requests.""" + await asyncio.gather(*[server.resume_generation.remote() for server in self.servers]) + + async def abort_request(self, request_id: str) -> dict[str, Any]: + """Abort a specific request. Tries all servers since we don't know which one has it. + + Args: + request_id: The ID of the request to abort. + + Returns: + dict[str, Any]: Abort result. + """ + # TODO(petersh6): we should only abort on the server that has the request. + results = await asyncio.gather(*[server.abort_request.remote(request_id) for server in self.servers]) + + for r in results: + if r.get("aborted", False): + return r + + return {"aborted": False, "request_id": request_id, "error": "Request not found on any server"} + + async def release_kv_cache(self): + # Drain all in-flight requests so that vLLM worker threads go idle + # before we touch engine.release_kv_cache() + await self.servers[0].wait_for_requests_to_drain.remote() + await asyncio.gather(*[server.release_kv_cache.remote() for server in self.servers]) + + # ----------------------------------------------------------------------- + # Hook methods for subclass overrides + # ----------------------------------------------------------------------- + + def _validate_launch_requirements(self) -> None: + """Validate requirements before launching. Override in subclasses.""" + # NOTE: We always use MP Executor backend whether it's single-node or multi-node. + # For multi-node without DP (e.g TP=16), need vllm>=0.11.1, https://github.com/vllm-project/vllm/pull/23691 + if self.config.data_parallel_size == 1 and self.nnodes > 1: + assert _VLLM_VERSION >= version.parse("0.11.1"), ( + "For multi-node MP Executor, either (1) set data_parallel_size > 1 or (2) upgrade vLLM to >= 0.11.1" + ) + + def _get_server_name_prefix(self) -> str: + """Return the Ray actor name prefix (e.g. 'vllm_').""" + return "vllm_" diff --git a/verl/verl/workers/rollout/vllm_rollout/vllm_rollout.py b/verl/verl/workers/rollout/vllm_rollout/vllm_rollout.py new file mode 100644 index 0000000000000000000000000000000000000000..c9f55cb1e86c429a847082c26c59cd3972685c32 --- /dev/null +++ b/verl/verl/workers/rollout/vllm_rollout/vllm_rollout.py @@ -0,0 +1,217 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +The vllm_rollout that can be applied in different backend +When working with FSDP: +- Use DTensor weight loader (recommended) or HF weight loader +- Utilize state_dict from the FSDP to synchronize the weights among tp ranks in vLLM +When working with Megatron: +- Use Megatron weight loader +- During training, only the current pp stage holds the parameters +- Before inference, broadcast the parameters of the current pp rank + to all other pp ranks (all pp ranks holds all the parameters) +- Bind the parameters to the inference engine +- Do inference in tp. pp is treated as additional dp +- After inference, all the parameters that doesn't belong to this pp rank is freed. +""" + +import logging +import os +import time +from typing import Any, Generator, Optional + +import ray +import torch +from packaging import version as vs +from torch.distributed.device_mesh import DeviceMesh + +from verl import DataProto +from verl.third_party.vllm import VLLM_SLEEP_LEVEL, get_version +from verl.utils.device import get_device_id, is_support_ipc +from verl.workers.config import HFModelConfig, RolloutConfig +from verl.workers.rollout.base import BaseRollout +from verl.workers.rollout.vllm_rollout.bucketed_weight_transfer import BucketedWeightSender +from verl.workers.rollout.vllm_rollout.utils import get_device_uuid + +logger = logging.getLogger(__file__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "INFO")) + + +def _check_vllm_version_for_sleep_level(): + # https://github.com/vllm-project/vllm/issues/25171 + minver = "0.11.0" + current_version = get_version("vllm") + if not current_version: + logger.warning("Could not determine vLLM version, assuming an older version for sleep_level configuration.") + return False + return vs.parse(current_version) >= vs.parse(minver) + + +class ServerAdapter(BaseRollout): + """ + vLLM server adapter used in native async mode, serve as a client to request vLLM server + to resume/release/update weights and kv_cache. + """ + + def __init__( + self, + config: RolloutConfig, + model_config: HFModelConfig, + device_mesh: DeviceMesh, + replica_rank: int = -1, + ): + super().__init__(config, model_config, device_mesh) + self.server_handle: ray.actor.ActorHandle = None + + rank = int(os.environ["RANK"]) + local_world_size = int(os.environ["RAY_LOCAL_WORLD_SIZE"]) + rollout_world_size = ( + self.config.tensor_model_parallel_size + * self.config.data_parallel_size + * self.config.pipeline_model_parallel_size + ) + if replica_rank == -1: + self.replica_rank = rank // rollout_world_size + else: + self.replica_rank = replica_rank + self.rollout_rank = rank % rollout_world_size + self.node_rank = self.rollout_rank // local_world_size + + if config.layered_summon or (config.expert_parallel_size > 1 and not _check_vllm_version_for_sleep_level()): + logger.warning("Setting the sleep level to 1 may cause a memory overflow.") + self.sleep_level = 1 + else: + self.sleep_level = VLLM_SLEEP_LEVEL + + self.device_uuid = get_device_uuid(get_device_id()) + # Use replica_rank + node-local rank to form ZMQ handle instead of GPU UUID, + # because CheckpointEngineWorker and vLLM worker may see different GPU UUIDs + # when CUDA_VISIBLE_DEVICES differs between processes (common on ROCm/AMD). + # Must use node-local rank (not rollout_rank) so it matches vLLM worker's + # local_rank on every node. Include replica_rank to avoid collisions when + # multiple replicas share a node, and the Ray job id so two independent + # verl jobs on the same host (or a new run after a crashed one with a + # stale socket file) cannot collide on the shared /tmp namespace. + local_rank = self.rollout_rank % local_world_size + job_id = ray.get_runtime_context().get_job_id() + self.zmq_handle = f"ipc:///tmp/rl-colocate-zmq-{job_id}-replica-{self.replica_rank}-rank-{local_rank}.sock" + + self.use_shm = not is_support_ipc() + if self.use_shm: + logger.warning( + "IPC is not supported on your devices. Falling back to shared memory for weight transfer, " + "which may cause performance degradation. If you are using Ascend NPUs, please ensure that " + "your software and CANN toolkit versions meet the requirements for IPC support. (Ascend HDK version " + ">= 25.3.rc1 and CANN toolkit version >= 8.3.RC1)" + ) + + async def _execute_method( + self, + method: str, + non_block: bool = False, + timeout: Optional[float] = None, + args: tuple = (), + kwargs: Optional[dict] = None, + ) -> Any: + """Execute method on inference engine via ray. + + Args: + method: The method name to execute on the server. + non_block: If True, execute the method asynchronously and return immediately. + timeout: Timeout for the collective_rpc call. + args: Positional arguments for the method. + kwargs: Keyword arguments for the method. + + Returns: + The result of the method execution, or None if non_block=True. + """ + if self.rollout_rank != 0: + return None + + # Lazy init http server adapter because http server is launched after hybrid engine. + if self.server_handle is None: + prefix = self._get_server_name_prefix() + self.server_handle = ray.get_actor(f"{prefix}server_{self.replica_rank}_{self.node_rank}") + + future = self.server_handle.collective_rpc.remote(method, timeout=timeout, args=args, kwargs=kwargs) + return future if non_block else await future + + async def resume(self, tags: list[str]): + """Resume rollout weights or kv cache in GPU memory. + + Args: + tags: weights or kv_cache. + """ + if self.config.free_cache_engine: + await self._execute_method("wake_up", kwargs={"tags": tags}) + + async def release(self): + """Release weights and kv cache in GPU memory.""" + if self.config.free_cache_engine: + await self._execute_method("sleep", kwargs={"level": self.sleep_level}) + + @torch.no_grad() + async def update_weights( + self, weights: Generator[tuple[str, torch.Tensor], None, None], global_steps: int = None, **kwargs + ): + """Update model weights via CUDA IPC (fallback to shared memory if IPC not supported) to inference workers.""" + start_time = time.time() + + future = await self._execute_method( + "update_weights_from_ipc", + non_block=True, + kwargs={**kwargs, "use_shm": self.use_shm}, + ) + + bucket_size_mb = self.config.checkpoint_engine.update_weights_bucket_megabytes + sender = BucketedWeightSender( + zmq_handle=self.zmq_handle, + bucket_size_mb=bucket_size_mb, + use_shm=self.use_shm, + ) + await sender.async_send_weights(weights) + + if future is not None: + await future + + # reset prefix cache after updating weights + if self.rollout_rank == 0: + await self.server_handle.clear_kv_cache.remote() + if global_steps is not None: + await self.server_handle.set_global_steps.remote(global_steps) + + if self.replica_rank == 0 and self.rollout_rank == 0: + logger.info(f"update_weights done, time cost: {time.time() - start_time:.2f}s") + + def _get_server_name_prefix(self) -> str: + """Return the Ray actor name prefix matching the rollout type (e.g. 'vllm_').""" + return f"{self.config.get('name', 'vllm')}_" + + def generate_sequences(self, prompts: DataProto) -> DataProto: + """Batch generate sequences in sync mode. + + Note: ServerAdapter uses async server mode and does not support synchronous + generation. Since SPMD mode was retired (PR #4411), the generation workflow + should use the async server interface instead. + + Raises: + NotImplementedError: Always raised as sync generation is not supported. + """ + raise NotImplementedError( + "ServerAdapter does not support synchronous generate_sequences(). " + "The vLLM SPMD mode was retired in PR #4411. For batch generation, " + "please use the async server interface via vLLMReplica and LLMServerClient, " + "or use HFRollout for synchronous generation. " + "See https://github.com/verl-project/verl/issues/4682 for more details." + ) diff --git a/verl/verl/workers/utils/__init__.py b/verl/verl/workers/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1cd1e8433dffa0b3ba420be3e346f4f5cd062014 --- /dev/null +++ b/verl/verl/workers/utils/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/verl/verl/workers/utils/__pycache__/__init__.cpython-312.pyc b/verl/verl/workers/utils/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..920e5374531a3a7bfd23ee531cdaceb7465ec15c Binary files /dev/null and b/verl/verl/workers/utils/__pycache__/__init__.cpython-312.pyc differ diff --git a/verl/verl/workers/utils/__pycache__/losses.cpython-312.pyc b/verl/verl/workers/utils/__pycache__/losses.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6d613ed7536f01d75721a81b78d87b6fce05bb94 Binary files /dev/null and b/verl/verl/workers/utils/__pycache__/losses.cpython-312.pyc differ diff --git a/verl/verl/workers/utils/__pycache__/padding.cpython-312.pyc b/verl/verl/workers/utils/__pycache__/padding.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7dff81b2001fae4ecf2667e4a0fb9fc5489948aa Binary files /dev/null and b/verl/verl/workers/utils/__pycache__/padding.cpython-312.pyc differ diff --git a/verl/verl/workers/utils/losses.py b/verl/verl/workers/utils/losses.py new file mode 100644 index 0000000000000000000000000000000000000000..c865b9dd17052835e7281d8b98c6e9e98d54f5b7 --- /dev/null +++ b/verl/verl/workers/utils/losses.py @@ -0,0 +1,186 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch +from tensordict import TensorDict + +from verl.trainer.ppo.core_algos import agg_loss, compute_value_loss, get_policy_loss_fn, kl_penalty +from verl.utils import tensordict_utils as tu +from verl.utils.dataset.dataset_utils import DatasetPadMode +from verl.utils.metric import AggregationType, Metric +from verl.utils.torch_functional import masked_mean, masked_sum +from verl.workers.config import ActorConfig, CriticConfig +from verl.workers.utils.padding import no_padding_2_padding + + +def sft_loss(config: ActorConfig, model_output, data: TensorDict, dp_group=None): + pad_mode = tu.get_non_tensor_data(data=data, key="pad_mode", default=DatasetPadMode.NO_PADDING) + dp_size = data["dp_size"] + batch_num_tokens = data["batch_num_tokens"] + + log_prob = model_output["log_probs"] + + if pad_mode == DatasetPadMode.NO_PADDING: + # log_prob and loss mask are nested tensors of shape [bsz, j1] + # for each sample, loss mask shape is [1, prompt_length + response_length] + loss_mask = data["loss_mask"] + + log_prob_flatten = log_prob.values() + loss_mask_flatten = loss_mask.values() + + # left-shift the loss mask by one token to align with log_prob + loss_mask_flatten = torch.roll(loss_mask_flatten, shifts=-1, dims=0) + + # NOTE: loss is averaged over all tokens in the batch across all data parallel groups, + # For FSDP backend, the loss is directly used for backward; while for Megatron backend, + # the loss should be scaled by `num_microbatches` for pp schedule. + loss = -masked_sum(log_prob_flatten, loss_mask_flatten) / batch_num_tokens * dp_size + else: + response_mask = data["response_mask"].to(bool) + loss = -masked_sum(log_prob, response_mask) / batch_num_tokens * dp_size + + return loss, {} + + +def ppo_loss(config: ActorConfig, model_output, data: TensorDict, dp_group=None): + """Computes ppo loss from model output (log_prob, entropy, values, etc. ) and old_log_probs from data.""" + log_prob = no_padding_2_padding(model_output["log_probs"], data) + entropy = model_output.get("entropy", None) + if entropy is not None: + entropy = no_padding_2_padding(entropy, data) + + # global batch info for loss aggregation + config.global_batch_info["dp_size"] = data["dp_size"] + config.global_batch_info["batch_num_tokens"] = data["batch_num_tokens"] + config.global_batch_info["global_batch_size"] = data["global_batch_size"] + config.global_batch_info["loss_scale_factor"] = config.loss_scale_factor + + # assumes that if any of the global batch info is set, the policy_loss_fn will + # normalize using dp_size/global_bsz/global_token; in this case, metric aggregation should be SUM + # to reflect the mean loss over the global batch + if ( + data["dp_size"] > 1 + or data["batch_num_tokens"] is not None + or data["global_batch_size"] is not None + or config.loss_scale_factor is not None + ): + metric_aggregation = AggregationType.SUM + else: + metric_aggregation = AggregationType.MEAN + + metrics = {} + + # select fields and convert to padded tensor + fields = ["response_mask", "old_log_probs", "advantages"] + if "rollout_is_weights" in data: + fields.append("rollout_is_weights") + if "ref_log_prob" in data: + fields.append("ref_log_prob") + data = data.select(*fields).to_padded_tensor() + + response_mask = data["response_mask"].to(bool) + # compute policy loss + old_log_prob = data["old_log_probs"] + advantages = data["advantages"] + rollout_is_weights = data.get("rollout_is_weights", None) + + loss_agg_mode = config.loss_agg_mode + + loss_mode = config.policy_loss.get("loss_mode", "vanilla") + + policy_loss_fn = get_policy_loss_fn(loss_mode) + pg_loss, pg_metrics = policy_loss_fn( + old_log_prob=old_log_prob, + log_prob=log_prob, + advantages=advantages, + response_mask=response_mask, + loss_agg_mode=loss_agg_mode, + config=config, + rollout_is_weights=rollout_is_weights, + ) + + # AggregationType.MEAN for pg metrics: assumes policy_loss_fn normalizes by local_bsz/local_tokens + # Ex: in compute_policy_loss_vanilla, pg_metrics are pg_clipfrac, ppo_kl, pg_clipfrac_lower + pg_metrics = Metric.from_dict(pg_metrics, aggregation=AggregationType.MEAN) + + metrics.update(pg_metrics) + metrics["actor/pg_loss"] = Metric(value=pg_loss, aggregation=metric_aggregation) + policy_loss = pg_loss + + # add entropy loss + if entropy is not None: + entropy_loss = agg_loss( + loss_mat=entropy, loss_mask=response_mask, loss_agg_mode=loss_agg_mode, **config.global_batch_info + ) + entropy_coeff = config.entropy_coeff + policy_loss -= entropy_coeff * entropy_loss + metrics["actor/entropy_loss"] = Metric(value=entropy_loss, aggregation=metric_aggregation) + + # add kl loss + if config.use_kl_loss: + ref_log_prob = data["ref_log_prob"] + # compute kl loss + kld = kl_penalty(logprob=log_prob, ref_logprob=ref_log_prob, kl_penalty=config.kl_loss_type) + kl_loss = agg_loss( + loss_mat=kld, loss_mask=response_mask, loss_agg_mode=config.loss_agg_mode, **config.global_batch_info + ) + + policy_loss += kl_loss * config.kl_loss_coef + metrics["kl_loss"] = Metric(value=kl_loss, aggregation=metric_aggregation) + metrics["kl_coef"] = config.kl_loss_coef + + return policy_loss, metrics + + +def value_loss(config: CriticConfig, model_output, data: TensorDict, dp_group=None): + """value loss + + Args: + config: CriticConfig + model_output: model output from the model + data: the input to the model + dp_group: data paralle group + + Returns: + value loss + """ + vpreds = no_padding_2_padding(model_output["values"], data) # (bsz, response_length) + + # select fields and convert to padded tensor + data = data.select("values", "returns", "response_mask").to_padded_tensor() + values = data["values"] + returns = data["returns"] + response_mask = data["response_mask"].to(bool) + + vf_loss, vf_clipfrac = compute_value_loss( + vpreds=vpreds, + values=values, + returns=returns, + response_mask=response_mask, + cliprange_value=config.cliprange_value, + loss_agg_mode=config.loss_agg_mode, + ) + + metrics = {} + + metrics.update( + { + "critic/vf_loss": vf_loss.detach().item(), + "critic/vf_clipfrac": vf_clipfrac.detach().item(), + "critic/vpred_mean": masked_mean(vpreds, response_mask).detach().item(), + } + ) + + return vf_loss, metrics diff --git a/verl/verl/workers/utils/padding.py b/verl/verl/workers/utils/padding.py new file mode 100644 index 0000000000000000000000000000000000000000..ed8d145520571f60452579a85ffba981dd8d64ec --- /dev/null +++ b/verl/verl/workers/utils/padding.py @@ -0,0 +1,231 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +import torch.nn.functional as F +from tensordict import TensorDict + +from verl.utils import tensordict_utils as tu +from verl.utils.attention_utils import index_first_axis, unpad_input + + +def left_right_2_no_padding(data: TensorDict) -> TensorDict: + """ + Convert TensorDict from left-right padding to no-padding format. + + Args: + data: TensorDict with "input_ids", "attention_mask", "response_mask", "position_ids" + + Returns: + data: TensorDict with + - Tensor includes NestedTensors like "input_ids", "loss_mask", "position_ids" + - NonTensorData includes "max_seq_len", "max_response_len", "indices" + + Note: + 1. the return input_ids/position_ids/loss_mask are nested tensor. + 2. we will remove "attention_mask", "response" in the return data, but "response_mask" is kept. + """ + assert "input_ids" in data, "input_ids is required in left-right padding data" + assert "attention_mask" in data, "attention_mask is required in left-right padding data" + assert "response_mask" in data, "response_mask is required in left-right padding data" + assert "position_ids" in data, "position_ids is required in left-right padding data" + + input_ids = data.pop("input_ids") + attention_mask = data["attention_mask"] + response_mask = data["response_mask"] + position_ids = data["position_ids"] # (bs, seq_len) or # (bs, 4, seq_len) + + max_seq_len, max_response_len = input_ids.shape[1], response_mask.shape[1] + tu.assign_non_tensor_data(data, "max_seq_len", max_seq_len) + tu.assign_non_tensor_data(data, "max_response_len", max_response_len) + + input_ids_rmpad, indices, cu_seqlens, *_ = unpad_input(input_ids.unsqueeze(-1), attention_mask) + tu.assign_non_tensor_data(data, "indices", indices) + + input_ids_nested = torch.nested.nested_tensor_from_jagged(input_ids_rmpad.squeeze(-1), offsets=cu_seqlens) + + position_ids_list = [] + for i in range(attention_mask.shape[0]): + curr_mask = attention_mask[i].bool() + curr_pos_ids = position_ids[i] + if curr_pos_ids.dim() == 1: # (seq_len,) + valid_ids = curr_pos_ids[curr_mask] + else: # (4, seq_len) + valid_ids = curr_pos_ids[:, curr_mask] + position_ids_list.append(valid_ids) + position_ids_nested = torch.nested.as_nested_tensor(position_ids_list, layout=torch.jagged) + + data["input_ids"] = input_ids_nested + data["position_ids"] = position_ids_nested + data["loss_mask"] = data["response_mask"] + + routed_experts = data.get("routed_experts", None) + if routed_experts is not None and not routed_experts.is_nested: + if routed_experts.max() <= 255: + routed_experts = routed_experts.to(torch.uint8) + routed_experts_rmpad = index_first_axis(routed_experts.unsqueeze(-1).flatten(0, 1), indices) + routed_experts_nested = torch.nested.nested_tensor_from_jagged( + routed_experts_rmpad.squeeze(-1), offsets=cu_seqlens + ) + data["routed_experts"] = routed_experts_nested + + # (bsz, seqlen, topk) + teacher_logprobs = data.get("teacher_logprobs", None) + teacher_ids = data.get("teacher_ids", None) + if teacher_logprobs is not None and teacher_ids is not None: + teacher_logprobs_rmpad = index_first_axis(teacher_logprobs.unsqueeze(-1).flatten(0, 1), indices) + teacher_ids_rmpad = index_first_axis(teacher_ids.unsqueeze(-1).flatten(0, 1), indices) + teacher_logprobs_nested = torch.nested.nested_tensor_from_jagged( + teacher_logprobs_rmpad.squeeze(-1), offsets=cu_seqlens + ) + teacher_ids_nested = torch.nested.nested_tensor_from_jagged(teacher_ids_rmpad.squeeze(-1), offsets=cu_seqlens) + data["teacher_logprobs"] = teacher_logprobs_nested + data["teacher_ids"] = teacher_ids_nested + + return data + + +def no_padding_2_padding(tensor: torch.Tensor, data: TensorDict) -> torch.Tensor: + """Slice response from unpad model output. + + Args: + tensor: a nested tensor or a tensor of shape (total_nnz,*), + total_nnz is the total number of tokens across all sequences in the batch + + data: TensorDict with "prompts", "responses", "attention_mask" + + Returns: + tensor: sliced response tensor of shape [bsz, max_response_len, *] + """ + values = tensor.values() if tensor.is_nested else tensor + prompt_ids = data["prompts"] + response_ids = data["responses"] + + max_response_len = tu.get_non_tensor_data(data=data, key="max_response_len", default=-1) + + if prompt_ids.is_nested: + prompt_lens = prompt_ids.offsets().diff() + response_lens = response_ids.offsets().diff() + if max_response_len < 0: + max_response_len = response_lens.max().item() + else: + attention_mask = data["attention_mask"] + assert not attention_mask.is_nested + prompt_lens = attention_mask[:, : prompt_ids.shape[1]].sum(dim=1) + response_lens = attention_mask[:, prompt_ids.shape[1] :].sum(dim=1) + max_response_len = response_ids.shape[1] + + sequence_lens = prompt_lens + response_lens + sequence_offsets = sequence_lens.cumsum(dim=0) + assert sequence_offsets[-1].item() == values.shape[0] + assert not prompt_lens.eq(0).any(), f"seq_offset - resp_len - 1 assumes prompt_len > 0. Got {prompt_lens}" + + response_list = [] + # Skip padding dimensions after sequence dimensions, if any. + skip_padding = (0, 0) * (values.ndim - 1) + for resp_len, seq_offset in zip(response_lens, sequence_offsets, strict=True): + pad_size = max_response_len - resp_len + # left-shift model output by one token for log_probs/values + response_list.append(F.pad(values[seq_offset - resp_len - 1 : seq_offset - 1], (*skip_padding, 0, pad_size))) + + output = torch.stack(response_list, dim=0) + return output + + +def build_attention_mask_from_nested(input_ids: torch.Tensor, max_seq_len: int | None = None) -> torch.Tensor: + """Build a padded full-sequence attention mask from nested input ids.""" + assert input_ids.is_nested, "input_ids must be a nested tensor" + device = input_ids.values().device + seq_lens = input_ids.offsets().diff().to(device=device) + if max_seq_len is None: + max_seq_len = int(seq_lens.max().item()) + positions = torch.arange(max_seq_len, device=device).unsqueeze(0) + return (positions < seq_lens.unsqueeze(1)).to(torch.int32) + + +def embeds_padding_2_no_padding(data: TensorDict) -> TensorDict: + """ + Convert TensorDict from prompt embeds with padding to no-padding format. + + Currently we expect the prompt embedding mask to be [1111000...] format, + which means the valid tokens are continuous and start from the left. + + Args: + data: TensorDict with "prompt_embeds", "prompt_embeds_mask", + "negative_prompt_embeds", "negative_prompt_embeds_mask" + + Returns: + data: TensorDict with + - Tensor includes NestedTensors "prompt_embeds", "prompt_embeds_mask", + "negative_prompt_embeds", "negative_prompt_embeds_mask" + """ + + def _to_nested(embeds: torch.Tensor, mask: torch.Tensor): + """Strip padding from (bs, seq_len, dim) embeds using the boolean mask and return nested tensors.""" + embeds_list, mask_list = [], [] + for i in range(mask.shape[0]): + curr_mask = mask[i].bool() + embeds_list.append(embeds[i, curr_mask, :]) + mask_list.append(curr_mask[curr_mask]) + return ( + torch.nested.as_nested_tensor(embeds_list, layout=torch.jagged), + torch.nested.as_nested_tensor(mask_list, layout=torch.jagged), + ) + + data["prompt_embeds"], data["prompt_embeds_mask"] = _to_nested(data["prompt_embeds"], data["prompt_embeds_mask"]) + + if isinstance(data.get("negative_prompt_embeds", None), torch.Tensor): + data["negative_prompt_embeds"], data["negative_prompt_embeds_mask"] = _to_nested( + data["negative_prompt_embeds"], data["negative_prompt_embeds_mask"] + ) + + return data + + +def response_from_nested(tensor: torch.Tensor, response_mask: torch.Tensor) -> torch.Tensor: + """Extract response from nested model output. + + Args: + tensor: a nested tensor with shape (bsz, prompt_len + response_len) + response_mask: a nested tensor with shape (bsz, response_len) + + Returns: + tensor: a nested tensor with shape (bsz, response_len) + """ + values, offsets = tensor.values(), tensor.offsets() + response_lens = response_mask.offsets().diff() + response_list = [] + for resp_len, seq_offset in zip(response_lens, offsets[1:], strict=True): + # left-shift model output by one token for log_probs/values + response_list.append(values[seq_offset - resp_len - 1 : seq_offset - 1]) + return torch.nested.as_nested_tensor(response_list, layout=torch.jagged) + + +def response_to_nested(tensor: torch.Tensor, response_mask: torch.Tensor) -> torch.Tensor: + """Convert padded response tensor to nested tensor. + + Args: + tensor: a tensor with shape (bsz, response_len) + response_mask: a nested tensor with shape (bsz, response_len) + + Returns: + tensor: a nested tensor with shape (bsz, response_len) + """ + assert response_mask.is_nested + response_lens = response_mask.offsets().diff() + response_list = [] + for i in range(tensor.shape[0]): + response_list.append(tensor[i, : response_lens[i]]) + + return torch.nested.as_nested_tensor(response_list, layout=torch.jagged)